Request health pills (but not re-layout) on each graph dashboard reload.

The most recent health pill per node is then rendered.
Change: 147087895

5 files changed, 88 insertions, 55 deletions

diff --git a/tensorflow/tensorboard/components/tf_graph/tf-graph-scene.html b/tensorflow/tensorboard/components/tf_graph/tf-graph-scene.html
index 8f1321138b..24b7bb7b21 100644
--- a/tensorflow/tensorboard/components/tf_graph/tf-graph-scene.html
+++ b/tensorflow/tensorboard/components/tf_graph/tf-graph-scene.html
@@ -702,16 +702,13 @@ Polymer({
       readOnly: true,
       value: ['#762a83', '#af8dc3', '#f7da0b', '#a2c96f', '#1f8926', '#0909c6'],
     },
-    /**
-     * A mapping from node name to a list of tf.graph.scene.HealthPills. Updated when the client
-     * receives a health pills response.
-     */
-    nodeNamesToHealthPills: {},
+    // A mapping between node name to the tf.graph.scene.HealthPill to render.
+    nodeNamesToHealthPill: Object,
   },
   observers: [
     '_colorByChanged(colorBy)',
     '_buildAndFit(renderHierarchy)',
-    '_updateHealthPills(nodeNamesToHealthPills)',
+    '_updateHealthPills(nodeNamesToHealthPill)',
   ],
   getNode: function(nodeName) {
     return this.renderHierarchy.getRenderNodeByName(nodeName);
@@ -754,7 +751,7 @@ Polymer({
     }.bind(this));
     // Update the minimap again when the graph is done animating.
     setTimeout(function() {
-      this._updateHealthPills(this.nodeNamesToHealthPills);
+      this._updateHealthPills(this.nodeNamesToHealthPill);
       this.minimap.update();
     }.bind(this), tf.graph.layout.PARAMS.animation.duration);
   },
@@ -904,8 +901,8 @@ Polymer({
   getEdgeGroup: function(e) {
     return this._edgeGroupIndex[e];
   },
-  _updateHealthPills: function(nodeNamesToHealthPills) {
-    tf.graph.scene.addHealthPills(this.$.svg, nodeNamesToHealthPills, this.healthPillColors);
+  _updateHealthPills: function(nodeNamesToHealthPill) {
+    tf.graph.scene.addHealthPills(this.$.svg, nodeNamesToHealthPill, this.healthPillColors);
   },
   /**
    * Update node and annotation node of the given name.
diff --git a/tensorflow/tensorboard/components/tf_graph/tf-graph.html b/tensorflow/tensorboard/components/tf_graph/tf-graph.html
index 6c542f1720..8002c6afb9 100644
--- a/tensorflow/tensorboard/components/tf_graph/tf-graph.html
+++ b/tensorflow/tensorboard/components/tf_graph/tf-graph.html
@@ -64,6 +64,7 @@ paper-button {
           selected-node="[[selectedNode]]"
           color-by="[[colorBy]]"
           progress="[[progress]]"
+          node-names-to-health-pill="[[nodeNamesToHealthPill]]"
     ></tf-graph-scene>
   </div>
 </div>
@@ -118,6 +119,8 @@ Polymer({
       type: Boolean,
       value: true
     },
+    // A mapping between node name to the tf.graph.scene.HealthPill to render.
+    nodeNamesToHealthPill: Object,
   },
   observers: [
     '_statsChanged(stats, devicesForStats)',
diff --git a/tensorflow/tensorboard/components/tf_graph_board/tf-graph-board.html b/tensorflow/tensorboard/components/tf_graph_board/tf-graph-board.html
index 2bc16b15c0..a0032cf6c7 100644
--- a/tensorflow/tensorboard/components/tf_graph_board/tf-graph-board.html
+++ b/tensorflow/tensorboard/components/tf_graph_board/tf-graph-board.html
@@ -142,6 +142,7 @@ paper-progress {
               color-by="[[colorBy]]"
               color-by-params="{{colorByParams}}"
               progress="{{progress}}"
+              node-names-to-health-pill="[[nodeNamesToHealthPill]]"
     ></tf-graph>
   </div>
   <div id="info">
@@ -186,6 +187,8 @@ Polymer({
       type: Object,
       notify: true
     },
+    // A mapping between node name to the tf.graph.scene.HealthPill to render.
+    nodeNamesToHealthPill: Object,
     // Private API: Data routing between child components.
     _selectedNode: String,
     // The enum value of the include property of the selected node.
diff --git a/tensorflow/tensorboard/components/tf_graph_common/lib/scene/scene.ts b/tensorflow/tensorboard/components/tf_graph_common/lib/scene/scene.ts
index 518d89e82a..5e9e2504d5 100644
--- a/tensorflow/tensorboard/components/tf_graph_common/lib/scene/scene.ts
+++ b/tensorflow/tensorboard/components/tf_graph_common/lib/scene/scene.ts
@@ -473,18 +473,18 @@ export function positionEllipse(ellipse, cx: number, cy: number,
  * Renders a health pill for an op atop a node.
  */
 function _addHealthPill(
-    nodeGroupElement: SVGElement, healthPills: HealthPill[],
+    nodeGroupElement: SVGElement, healthPill: HealthPill,
     nodeInfo: render.RenderGroupNodeInfo, colors: string[]) {
   // Check if text already exists at location.
   d3.select(nodeGroupElement.parentNode)
       .selectAll('.health-pill-group')
       .remove();
 
-  if (!nodeInfo || !healthPills || !healthPills.length) {
+  if (!nodeInfo || !healthPill) {
     return;
   }
 
-  let lastHealthPillData = healthPills[healthPills.length - 1].value;
+  let lastHealthPillData = healthPill.value;
 
   // For now, we only visualize the 6 values that summarize counts of tensor
   // elements of various categories: -Inf, negative, 0, positive, Inf, and NaN.
@@ -492,10 +492,10 @@ function _addHealthPill(
 
   let healthPillWidth = 60;
   let healthPillHeight = 10;
-  let healthPill = document.createElementNS(svgNamespace, 'svg');
-  healthPill.classList.add('health-pill-group');
-  healthPill.setAttribute('width', String(healthPillWidth));
-  healthPill.setAttribute('height', String(healthPillHeight));
+  let healthPillSvg = document.createElementNS(svgNamespace, 'svg');
+  healthPillSvg.classList.add('health-pill-group');
+  healthPillSvg.setAttribute('width', String(healthPillWidth));
+  healthPillSvg.setAttribute('height', String(healthPillHeight));
 
   let totalCount = lastHealthPillData[1];
   // Create 1 rectangle for each category.
@@ -515,34 +515,38 @@ function _addHealthPill(
         'x', String(totalWidthDividedByTotalCount * totalCountSoFar));
     rect.setAttribute('fill', colors[i]);
     totalCountSoFar += lastHealthPillOverview[i];
-    Polymer.dom(healthPill).appendChild(rect);
+    Polymer.dom(healthPillSvg).appendChild(rect);
   }
 
   // Center this health pill just right above the node for the op.
-  healthPill.setAttribute(
+  healthPillSvg.setAttribute(
       'x', String(nodeInfo.x - healthPillWidth + nodeInfo.width / 2));
-  healthPill.setAttribute(
+  healthPillSvg.setAttribute(
       'y',
       String(nodeInfo.y - healthPillHeight - nodeInfo.coreBox.height / 2 - 2));
-  Polymer.dom(nodeGroupElement.parentNode).appendChild(healthPill);
+  Polymer.dom(nodeGroupElement.parentNode).appendChild(healthPillSvg);
 }
 
 /**
  * Adds health pills (which visualize tensor summaries) to a graph group.
  * @param svgRoot The root SVG element of the graph to add heath pills to.
- * @param nodeNamesToHealthPills An object mapping node names to health pills.
+ * @param nodeNamesToHealthPills An object mapping node name to health pill.
  * @param colors A list of colors to use.
  */
 export function addHealthPills(
-    svgRoot: SVGElement, nodeNamesToHealthPills: {[key: string]: HealthPill[]},
+    svgRoot: SVGElement, nodeNamesToHealthPill: {[key: string]: HealthPill},
     colors: string[]) {
-  let svgRootSelection = d3.select(svgRoot);
+  if (!nodeNamesToHealthPill) {
+    // No health pill information available.
+    return;
+  }
 
+  let svgRootSelection = d3.select(svgRoot);
   svgRootSelection.selectAll('g.nodeshape')
       .each(function(nodeInfo: render.RenderGroupNodeInfo) {
         // The element is the first item of a d3 selection.
         _addHealthPill(
-            this, nodeNamesToHealthPills[nodeInfo.node.name], nodeInfo, colors);
+            this, nodeNamesToHealthPill[nodeInfo.node.name], nodeInfo, colors);
       });
 };
 
diff --git a/tensorflow/tensorboard/components/tf_graph_dashboard/tf-graph-dashboard.html b/tensorflow/tensorboard/components/tf_graph_dashboard/tf-graph-dashboard.html
index 5b3857e6f5..e9ab46d5a5 100644
--- a/tensorflow/tensorboard/components/tf_graph_dashboard/tf-graph-dashboard.html
+++ b/tensorflow/tensorboard/components/tf_graph_dashboard/tf-graph-dashboard.html
@@ -104,46 +104,72 @@ Polymer({
   ],
   properties: {
     _datasets: Object,
-    _renderHierarchy: Object,
-    backend: {type: Object, observer: 'reload'},
+    _renderHierarchy: {type: Object, observer: '_renderHierarchyChanged'},
+    backend: {type: Object, observer: '_backendChanged'},
     debuggerDataEnabled: Boolean,
     runs: Array
   },
+  listeners: {
+    'node-toggle-expand': '_handleNodeToggleExpand',
+  },
   reload: function() {
-    let requestPromises = [
-        this.backend.graphRuns(),
-        this.backend.runMetadataRuns()];
-    if (this.debuggerDataEnabled && this._renderHierarchy) {
-      // Request debugger data on graph reloads.
-      requestPromises.push(this._renderHierarchy.getNamesOfRenderedOps());
+    if (!this.debuggerDataEnabled ||
+        !this._renderHierarchy ||
+        this._datasetsEmpty(this._datasets)) {
+      // Do not load debugger data if the feature is disabled or if the graph itself has not loaded
+      // yet. We need the graph to load so that we know which nodes to request health pills for.
+      return;
     }
-    // TODO(chizeng): Prevent the graph layout itself from refreshing when the reload button is hit.
-    // We should only update the data atop the graph such as debugger data. The graph itself does
-    // not change on successive reloads.
-    Promise.all(requestPromises)
-    .then(function(result) {
-      var runsWithGraph = result[0].sort(VZ.Sorting.compareTagNames);
-      var runToMetadata = result[1];
-      var datasets = _.map(runsWithGraph, function(runName) {
-        return {
-          name: runName,
-          path: this.backend.router.graph(runName, tf.graph.LIMIT_ATTR_SIZE,
-            tf.graph.LARGE_ATTRS_KEY),
-          runMetadata: runToMetadata[runName] ? _.map(
-            runToMetadata[runName].sort(VZ.Sorting.compareTagNames), function(tag) {
-              return {
-                tag: tag,
-                path: this.backend.router.runMetadata(tag, runName)
-              };
-            }, this) : []
-        };
-      }, this);
-      this.set('_datasets', datasets);
+
+    // Request debugger data on graph reloads, but do not re-request the graph itself. The graph
+    // would not change across reloads.
+    this._requestHealthPills();
+  },
+  _backendChanged: function(backend) {
+    Promise.all([backend.graphRuns(), backend.runMetadataRuns()])
+      .then(function(result) {
+        var runsWithGraph = result[0].sort(VZ.Sorting.compareTagNames);
+        var runToMetadata = result[1];
+        var datasets = _.map(runsWithGraph, function(runName) {
+          return {
+            name: runName,
+            path: backend.router.graph(
+                runName, tf.graph.LIMIT_ATTR_SIZE, tf.graph.LARGE_ATTRS_KEY),
+            runMetadata: runToMetadata[runName] ? _.map(
+              runToMetadata[runName].sort(VZ.Sorting.compareTagNames), function(tag) {
+                return {
+                  tag: tag,
+                  path: backend.router.runMetadata(tag, runName)
+                };
+              }, this) : []
+          };
+        }, this);
+        this.set('_datasets', datasets);
+      }.bind(this));
+  },
+  _requestHealthPills: function() {
+    this.backend.healthPills(this._renderHierarchy.getNamesOfRenderedOps()).then(function(result) {
+      // For now, show the latest health pill for each node.
+      // TODO(chizeng): Let the user select which step to show.
+      var mapping = {};
+      for (let nodeName in result) {
+        mapping[nodeName] = result[nodeName].pop();
+      }
+      this.$$('#graphboard').set('nodeNamesToHealthPill', mapping);
     }.bind(this));
   },
   _datasetsEmpty: function(datasets) {
     return !datasets || !datasets.length;
-  }
+  },
+  _renderHierarchyChanged: function(renderHierarchy) {
+    // Reload any data on the graph when the render hierarchy (which determines which nodes are
+    // rendered) changes.
+    this.reload();
+  },
+  _handleNodeToggleExpand: function() {
+    // Nodes were toggled. We may need to request health pills for more nodes.
+    this._requestHealthPills();
+  },
 });
 })();
 </script>