From db4eff2361189b95ee131bcf0963d7b12e7b2f85 Mon Sep 17 00:00:00 2001
From: Googler <noreply@google.com>
Date: Wed, 18 Dec 2024 23:41:44 -0800
Subject: [PATCH] Add `get_node_execution_count` API to Tflex UI backend

### Why
This is for rendering execution count per node for large async run graphs before fetching actual data for better usability.

![Example usage of the API for async runs](https://screenshot.googleplex.com/7t9pqYsdpRSAbiA.png)

PiperOrigin-RevId: 707801411
---
 src/app/directed_acyclic_graph_raw.ng.html | 7 ++++---
 src/app/directed_acyclic_graph_raw.ts      | 8 +++++---
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/src/app/directed_acyclic_graph_raw.ng.html b/src/app/directed_acyclic_graph_raw.ng.html
index 3e0fd10..3bf4fb3 100644
--- a/src/app/directed_acyclic_graph_raw.ng.html
+++ b/src/app/directed_acyclic_graph_raw.ng.html
@@ -385,15 +385,16 @@
         This causes the DAG to open up, fade-in the label before hiding it,
         which looks very glitchy
       -->
-      <span *ngIf="group.treatAsLoop" class="fade-in" [attr.aria-hidden]="isGroupExpanded(group)">
-        <span class="iteration-counter"> {{ getIterationsFor(group).length }} iterations </span>
-      </span>
       @if (showGroupLabel(group)) {
         <div class="group-label">
           <span class="fade-in">
             <span> {{ group.groupLabel }} </span>
           </span>
         </div>
+      } @else if (group.treatAsLoop) {
+        <span class="fade-in" [attr.aria-hidden]="isGroupExpanded(group)">
+          <span class="iteration-counter"> {{ getIterationsFor(group).length }} iterations </span>
+        </span>
       }
       <button
         class="expand-toggle"
diff --git a/src/app/directed_acyclic_graph_raw.ts b/src/app/directed_acyclic_graph_raw.ts
index fa5dd94..f829644 100644
--- a/src/app/directed_acyclic_graph_raw.ts
+++ b/src/app/directed_acyclic_graph_raw.ts
@@ -460,13 +460,15 @@ export class DagRaw implements DoCheck, OnInit, OnDestroy {
   }
 
   showGroupLabel(group: DagGroup) {
-    if (group.treatAsLoop || !group.groupLabel) return false;
+    if (!group.groupLabel) return false;
     if (this.isGroupExpanded(group) && group.hasControlNode &&
         !group.hideControlNodeOnExpand) {
-      return false
+      return false;
     }
-    return true
+
+    return true;
   }
+
   broadcastIterChange(
       group: DagGroup, iterationNode: GroupIterationRecord['iterationNode']) {
     (group as any)._cachedSelection = iterationNode;