Let the user view health pills at any step.

This involves adding a toggle to the health pills info box in the graph visualizer. When that toggle is enabled, Tensorboard makes a request for health pills at step X when the user moves the slider.

This feature can be very slow because it requires reading from disk. Viewing health pills at say step 100,000 could take minutes to an hour. We must design ways to make this faster (for instance, have the debugger write events at a much greater frequency only after it encounters a bad value).
Change: 150379929

4 files changed, 224 insertions, 22 deletions

diff --git a/tensorflow/tensorboard/components/tf_backend/backend.ts b/tensorflow/tensorboard/components/tf_backend/backend.ts
index b4be89163e..b87ced2ec2 100644
--- a/tensorflow/tensorboard/components/tf_backend/backend.ts
+++ b/tensorflow/tensorboard/components/tf_backend/backend.ts
@@ -200,8 +200,14 @@ module TF.Backend {
     /**
      * Returns a promise for requesting the health pills for a list of nodes.
      */
-    public healthPills(nodeNames: string[]): Promise<HealthPillsResponse> {
+    public healthPills(nodeNames: string[], step?: number):
+        Promise<HealthPillsResponse> {
       let postData = {'node_names': JSON.stringify(nodeNames)};
+      if (step !== undefined) {
+        // The user requested health pills for a specific step. This request
+        // might be slow since the backend reads events sequentially from disk.
+        postData['step'] = step;
+      }
       return this.requestManager.request(this.router.healthPills(), postData);
     }
 
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 c2ad0d09f7..5909172fbe 100644
--- a/tensorflow/tensorboard/components/tf_graph_board/tf-graph-board.html
+++ b/tensorflow/tensorboard/components/tf_graph_board/tf-graph-board.html
@@ -157,7 +157,11 @@ paper-progress {
               highlighted-node="{{_highlightedNode}}"
               color-by="[[colorBy]]"
               color-by-params="[[colorByParams]]"
+              debugger-data-enabled="[[debuggerDataEnabled]]"
+              are-health-pills-loading="[[areHealthPillsLoading]]"
               node-names-to-health-pills="[[nodeNamesToHealthPills]]"
+              all-steps-mode-enabled="{{allStepsModeEnabled}}"
+              specific-health-pill-step="{{specificHealthPillStep}}"
               health-pill-step-index="{{healthPillStepIndex}}"
     ></tf-graph-info>
   </div>
@@ -190,8 +194,26 @@ Polymer({
       type: Object,
       notify: true
     },
+    // Whether debugger data is enabled for this instance of Tensorboard.
+    debuggerDataEnabled: Boolean,
+    // Whether health pills are currently being loaded.
+    areHealthPillsLoading: Boolean,
     // A mapping between node name to the tf.graph.scene.HealthPill to render.
     nodeNamesToHealthPills: Object,
+    // Whether the user can request health pills for individual steps from the server. This can be
+    // slow compared the default of showing sampled health pills.
+    allStepsModeEnabled: {
+      type: Boolean,
+      notify: true,
+      value: false,
+    },
+    // Relevant if allStepsModeEnabled. The specific step for which to fetch health pills from the
+    // server for.
+    specificHealthPillStep: {
+      type: Number,
+      notify: true,
+      value: 0,
+    },
     // The step of health pills to show throughout the graph.
     healthPillStepIndex: Number,
     // Private API: Data routing between child components.
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 b0e296742a..d62e4ccedc 100644
--- a/tensorflow/tensorboard/components/tf_graph_dashboard/tf-graph-dashboard.html
+++ b/tensorflow/tensorboard/components/tf_graph_dashboard/tf-graph-dashboard.html
@@ -75,7 +75,11 @@ out-hierarchy-params="{{_hierarchyParams}}"
         graph="[[_graph]]"
         hierarchy-params="[[_hierarchyParams]]"
         progress="[[_progress]]"
+        debugger-data-enabled="[[debuggerDataEnabled]]"
+        are-health-pills-loading="[[_areHealthPillsLoading]]"
         node-names-to-health-pills="[[_nodeNamesToHealthPills]]"
+        all-steps-mode-enabled="{{allStepsModeEnabled}}"
+        specific-health-pill-step="{{specificHealthPillStep}}"
         health-pill-step-index="[[_healthPillStepIndex]]"
         render-hierarchy="{{_renderHierarchy}}"
         stats="[[_stats]]"
@@ -111,18 +115,38 @@ Polymer({
     _renderHierarchy: {type: Object, observer: '_renderHierarchyChanged'},
     backend: {type: Object, observer: '_backendChanged'},
     debuggerDataEnabled: Boolean,
+    allStepsModeEnabled: Boolean,
+    specificHealthPillStep: {type: Number, value: 0},
     healthPillsToggledOn: {type: Boolean, value: true, observer: '_healthPillsToggledOnChanged'},
+    // Whether health pills are currently being loaded, in which case we may want to say show a
+    // spinner.
+    _areHealthPillsLoading: Boolean,
     // Maps the names of nodes to an array of health pills (HealthPillDatums).
     _nodeNamesToHealthPills: {
       type: Object,
       value: {},
     },
     _healthPillStepIndex: Number,
-    runs: Array
+    // A strictly increasing ID. Each request for health pills has a unique ID. This helps us
+    // identify stale requests.
+    _healthPillRequestId: {type: Number, value: 1},
+    // The setTimeout ID for the pending request for health pills at a specific step.
+    _healthPillStepRequestTimerId: Number,
+    // The request for health pills at a specific step (as opposed to all sampled health pills) may
+    // involve slow disk reads. Hence, we throttle to 1 of those requests every this many ms.
+    _healthPillStepRequestTimerDelay: {
+      type: Number,
+      value: 500,
+      readOnly: true,
+    },
+    runs: Array,
   },
   listeners: {
     'node-toggle-expand': '_handleNodeToggleExpand',
   },
+  observers: [
+    '_maybeFetchHealthPillsAtSpecificStep(allStepsModeEnabled, specificHealthPillStep)',
+  ],
   reload: function() {
     if (!this.debuggerDataEnabled ||
         !this.healthPillsToggledOn ||
@@ -161,12 +185,51 @@ Polymer({
       }.bind(this));
   },
   _requestHealthPills: function() {
-    this.backend.healthPills(this._renderHierarchy.getNamesOfRenderedOps()).then(function(result) {
+    this.set('_areHealthPillsLoading', true);
+    const requestId = ++this._healthPillRequestId;
+
+    if (this._healthPillStepRequestTimerId !== null) {
+      // A request for health pills is already scheduled to be initiated. Clear it, and schedule a
+      // new request.
+      window.clearTimeout(this._healthPillStepRequestTimerId);
+      this._healthPillStepRequestTimerId = null;
+    }
+
+    if (this.allStepsModeEnabled) {
+      // This path may be slow. Schedule network requests to start some time later. If another
+      // request is scheduled in the mean time, drop this current request.
+      this._healthPillStepRequestTimerId = setTimeout(function() {
+        this._healthPillStepRequestTimerId = null;
+        this._initiateNetworkRequestForHealthPills(requestId);
+      }.bind(this), this._healthPillStepRequestTimerDelay);
+    } else {
+      // The user is fetching sampled steps. This path is fast, so no need to throttle. Directly
+      // fetch the health pills across the network.
+      this._initiateNetworkRequestForHealthPills(requestId);
+    }
+  },
+  // Initiates the network request for health pills. Do not directly call this method - network
+  // requests may be throttled. Instead, call _requestHealthPills, which uses this method.
+  _initiateNetworkRequestForHealthPills: function(requestId) {
+    if (this._healthPillRequestId !== requestId) {
+      // This possibly scheduled request was outdated before it was even sent across the network. Do
+      // not bother initiating it.
+      return;
+    }
+
+    const specificStep = this.allStepsModeEnabled ? this.specificHealthPillStep : undefined;
+    this.backend.healthPills(this._renderHierarchy.getNamesOfRenderedOps(), specificStep).then(
+        function(result) {
       if (!this.healthPillsToggledOn) {
         // The user has opted to hide health pills via the toggle button.
         return;
       }
 
+      if (requestId !== this._healthPillRequestId) {
+        // This response is no longer relevant.
+        return;
+      }
+
       // Set the index for which step to show for the health pills. By default, show the last step.
       // A precondition we assume (that Tensorboard's reservoir sampling guarantees) is that all
       // node names should be mapped to the same number of steps.
@@ -176,6 +239,8 @@ Polymer({
       }
 
       this.set('_nodeNamesToHealthPills', result);
+      this.set('_areHealthPillsLoading', false);
+      this.set('_healthPillStepRequestTimerId', null);
     }.bind(this));
   },
   _datasetsEmpty: function(datasets) {
@@ -199,6 +264,15 @@ Polymer({
       this.set('_nodeNamesToHealthPills', {});
     }
   },
+  // Fetch health pills for a specific step if applicable.
+  _maybeFetchHealthPillsAtSpecificStep: function(allStepsModeEnabled, specificHealthPillStep) {
+    if (!this._renderHierarchy) {
+      // The graph is not ready yet.
+      return;
+    }
+
+    this._requestHealthPills();
+  },
 });
 })();
 </script>
diff --git a/tensorflow/tensorboard/components/tf_graph_info/tf-graph-info.html b/tensorflow/tensorboard/components/tf_graph_info/tf-graph-info.html
index 1fb3c5a8de..45347fb1de 100644
--- a/tensorflow/tensorboard/components/tf_graph_info/tf-graph-info.html
+++ b/tensorflow/tensorboard/components/tf_graph_info/tf-graph-info.html
@@ -17,6 +17,7 @@ limitations under the License.
 
 <link rel="import" href="../polymer/polymer.html">
 <link rel="import" href="../paper-slider/paper-slider.html">
+<link rel="import" href="../paper-spinner/paper-spinner-lite.html">
 <link rel="import" href="../tf-graph-common/tf-graph-common.html">
 <link rel="import" href="tf-node-info.html">
 
@@ -78,6 +79,17 @@ h2 {
    * Apparently, the paper-slider lacks a mixin for those padding values. */
   width: calc(100% + 31px);
 }
+
+#health-pills-loading-spinner {
+  width: 20px;
+  height: 20px;
+  vertical-align: top;
+}
+
+#health-pill-step-number-input {
+  text-align: center;
+  vertical-align: top;
+}
 </style>
 <template is="dom-if" if="{{selectedNode}}">
   <paper-material elevation="1" class="card">
@@ -91,19 +103,49 @@ h2 {
     </tf-node-info>
   </paper-material>
 </template>
-<template is="dom-if" if="[[_healthPillsAvailable(nodeNamesToHealthPills)]]">
+<template is="dom-if" if="[[_healthPillsAvailable(debuggerDataEnabled, nodeNamesToHealthPills)]]">
   <paper-material elevation="1" class="card health-pill-legend">
-    <template is="dom-if" if="[[_maxStepIndex]]">
-      <h2>
-        Step of Health Pills: [[_currentStepDisplayValue]]
-      </h2>
+    <div class="title">
+      Enable all (not just sampled) steps. Requires slow disk read.
+    </div>
+    <paper-toggle-button id="enableAllStepsModeToggle" checked="{{allStepsModeEnabled}}">
+    </paper-toggle-button>
+    <h2>
+      Step of Health Pills:
+      <template is="dom-if" if="[[allStepsModeEnabled]]">
+        <input type="number"
+               id="health-pill-step-number-input"
+               min="0"
+               max="[[_biggestStepEverSeen]]"
+               value="{{specificHealthPillStep::input}}">
+      </template>
+      <template is="dom-if" if="[[!allStepsModeEnabled]]">
+        [[_currentStepDisplayValue]]
+      </template>
+
+      <paper-spinner-lite active
+                          hidden$=[[!areHealthPillsLoading]]
+                          id="health-pills-loading-spinner"></paper-spinner-lite>
+    </h2>
+    <template is="dom-if" if="[[allStepsModeEnabled]]">
       <paper-slider
             id="health-pill-step-slider"
-            immediate-value="{{healthPillStepIndex}}"
-            max="[[_maxStepIndex]]"
+            immediate-value="{{specificHealthPillStep}}"
+            max="[[_biggestStepEverSeen]]"
             snaps
             step="1"
-            value="{{healthPillStepIndex}}"></paper-slider>
+            value="{{specificHealthPillStep}}"></paper-slider>
+    </template>
+    <template is="dom-if" if="[[!allStepsModeEnabled]]">
+      <template is="dom-if" if="[[_maxStepIndex]]">
+        <paper-slider
+              id="health-pill-step-slider"
+              immediate-value="{{healthPillStepIndex}}"
+              max="[[_maxStepIndex]]"
+              snaps
+              step="1"
+              value="{{healthPillStepIndex}}"></paper-slider>
+      </template>
     </template>
     <h2>
       Health Pill
@@ -141,6 +183,13 @@ h2 {
         type: Number,
         notify: true,
       },
+      // Only relevant if we are in all steps mode, in which case the user may want to view health
+      // pills for a specific step.
+      specificHealthPillStep: {
+        type: Number,
+        value: 0,
+        notify: true,
+      },
       colorBy: String,
       // Two-ways
       selectedNode: {
@@ -156,6 +205,11 @@ h2 {
         type: Number,
         notify: true
       },
+      // Whether debugger data is enabled for this instance of Tensorboard.
+      debuggerDataEnabled: Boolean,
+      // Whether health pills are currently being loaded, in which case we show a spinner (and the
+      // current health pills shown might be out of date).
+      areHealthPillsLoading: Boolean,
       healthPillEntries: {
         type: Array,
         value: tf.graph.scene.healthPillEntries,
@@ -163,7 +217,19 @@ h2 {
       },
       healthPillValuesForSelectedNode: {
         type: Array,
-        computed: '_computeHealthPillForNode(nodeNamesToHealthPills, healthPillStepIndex, selectedNode)',
+        computed: '_computeHealthPillForNode(nodeNamesToHealthPills, healthPillStepIndex, selectedNode, allStepsModeEnabled, areHealthPillsLoading)',
+      },
+      // When all-steps mode is enabled, the user can request health pills for any step. In this
+      // mode, Tensorboard makes a request every time the user drags the slider to a different step.
+      allStepsModeEnabled: {
+        type: Boolean,
+        notify: true,
+      },
+      // The biggest step value ever seen. Used to determine what steps of health pills to let the
+      // user fetch in all steps mode.
+      _biggestStepEverSeen: {
+        type: Number,
+        computed: '_computeBiggestStepEverSeen(nodeNamesToHealthPills)',
       },
       _maxStepIndex: {
         type: Number,
@@ -171,7 +237,7 @@ h2 {
       },
       _currentStepDisplayValue: {
         type: String,
-        computed: '_computeCurrentStepDisplayValue(nodeNamesToHealthPills, healthPillStepIndex)',
+        computed: '_computeCurrentStepDisplayValue(nodeNamesToHealthPills, healthPillStepIndex, allStepsModeEnabled, specificHealthPillStep, areHealthPillsLoading)',
       },
     },
     listeners: {
@@ -188,12 +254,11 @@ h2 {
     _nodeListItemMouseout: function() {
       this.highlightedNode = null;
     },
-    _healthPillsAvailable: function(nodeNamesToHealthPills) {
-      let count = 0;
-      for (let nodeName in nodeNamesToHealthPills) {
-        return true;
-      }
-      return false;
+    _healthPillsAvailable: function(debuggerDataEnabled, nodeNamesToHealthPills) {
+      // So long as there is a mapping (even if empty) from node name to health pills, show the
+      // legend and slider. We do that because, even if no health pills exist at the current step,
+      // the user may desire to change steps, and the slider must show for the user to do that.
+      return debuggerDataEnabled && nodeNamesToHealthPills;
     },
     _computeTensorCountString: function(healthPillValuesForSelectedNode, valueIndex) {
       if (!healthPillValuesForSelectedNode) {
@@ -204,7 +269,12 @@ h2 {
       return healthPillValuesForSelectedNode[valueIndex].toFixed(0);
     },
     _computeHealthPillForNode: function(
-        nodeNamesToHealthPills, healthPillStepIndex, selectedNode) {
+        nodeNamesToHealthPills, healthPillStepIndex, selectedNode, allStepsModeEnabled, areHealthPillsLoading) {
+      if (areHealthPillsLoading) {
+        // Health pills are loading. Do not render data that is out of date.
+        return null;
+      }
+
       if (!selectedNode) {
         // No node is selected.
         return null;
@@ -216,7 +286,9 @@ h2 {
         return null;
       }
 
-      const healthPill = healthPills[healthPillStepIndex];
+      // If all steps mode is enabled, we use the first health pill in the list because the JSON
+      // response from the server is a mapping between node name and a list of 1 health pill.
+      const healthPill = healthPills[allStepsModeEnabled ? 0 : healthPillStepIndex];
       if (!healthPill) {
         // This node lacks a health pill at the current step.
         return null;
@@ -225,16 +297,44 @@ h2 {
       // The health pill count values start at 2. Each health pill contains 6 values.
       return healthPill.value.slice(2, 8);
     },
-    _computeCurrentStepDisplayValue: function(nodeNamesToHealthPills, healthPillStepIndex) {
+    _computeCurrentStepDisplayValue: function(
+        nodeNamesToHealthPills,
+        healthPillStepIndex,
+        allStepsModeEnabled,
+        specificHealthPillStep,
+        areHealthPillsLoading) {
+      if (allStepsModeEnabled) {
+        // The user seeks health pills for specific step from the server.
+        return specificHealthPillStep.toFixed(0);
+      }
+
+      if (areHealthPillsLoading) {
+        // The current step is undefined.
+        return 0;
+      }
+
       for (let nodeName in nodeNamesToHealthPills) {
         // All nodes have the same number of steps stored, so only examine 1 node. We cannot
         // directly index into the nodeNamesToHealthPills object because we do not have a key.
+        // If all steps mode is enabled, we only have 1 step to show.
         return nodeNamesToHealthPills[nodeName][healthPillStepIndex].step.toFixed(0);
       }
 
       // The current step could not be computed.
       return 0;
     },
+    _computeBiggestStepEverSeen: function(nodeNamesToHealthPills) {
+      for (let nodeName in nodeNamesToHealthPills) {
+        // All nodes have the same number of steps stored, so only examine 1 node.
+        // The index is 1 less than the count. Tensorboard backend logic guarantees that the length
+        // of the array will be greater than 1.
+        var healthPills = nodeNamesToHealthPills[nodeName];
+        return Math.max(this._biggestStepEverSeen, healthPills[healthPills.length - 1].step);
+      }
+
+      // No steps seen so far. Default to 0.
+      return this._biggestStepEverSeen || 0;
+    },
     _computeMaxStepIndex: function(nodeNamesToHealthPills) {
       for (let nodeName in nodeNamesToHealthPills) {
         // All nodes have the same number of steps stored, so only examine 1 node.