Support streaming trace viewer (#1128)

Author: trisolaran

tensorboard/components/tf_trace_viewer/BUILD

@@ -11,6 +11,9 @@ tf_web_library(
         "@org_chromium_catapult_vulcanized_trace_viewer//:trace_viewer_full.html",
     ],
     path = "/tf-trace-viewer",
+    deps = [
+        ":tf-trace-viewer-helper",
+    ],
 )
 
 tf_web_library(
@@ -23,3 +26,10 @@ tf_web_library(
     ],
 )
 
+tf_web_library(
+    name = "tf-trace-viewer-helper",
+    srcs = [
+        "tf-trace-viewer-helper.ts",
+    ],
+    path = "/tf-trace-viewer",
+)

tensorboard/components/tf_trace_viewer/tf-trace-viewer-helper.ts

@@ -0,0 +1,67 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the 'License');
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an 'AS IS' BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+/**
+ * @fileoverview Helper utilities for the trace viewer within TensorBoard's profile plugin.
+ */
+
+namespace tf_component_traceviewer {
+  /** Amount of zooming allowed before re-fetching. */
+  export const ZOOM_RATIO = 8;
+
+  /** Minimum safety buffer relative to viewport size. */
+  export const PRESERVE_RATIO = 2;
+
+  /** Amount to fetch relative to viewport size. */
+  export const FETCH_RATIO = 3;
+
+  export interface Range {
+    min: number;
+    max: number;
+  }
+
+  /**
+   * Expand the input range by scale, keep the center invariant.
+   */
+  export function expand(range: Range, scale: number) : Range {
+    var width = range.max - range.min;
+    var mid = range.min + width / 2;
+    return {
+      min: mid - scale * width / 2,
+      max: mid + scale * width / 2,
+    };
+  }
+  /**
+   * Check if range is within (totally included) in bounds.
+   */
+  export function within(range: Range, bounds: Range): boolean {
+    return bounds.min <= range.min && range.max <= bounds.max;
+  }
+  /**
+   * Return length of the range.
+   */
+  export function length(range: Range): number {
+    return range.max - range.min;
+  }
+  /**
+   * Return the intersection of two ranges.
+   */
+  export function intersect(range: Range, bounds: Range): Range {
+    return {
+      min: Math.max(range.min, bounds.min),
+      max: Math.min(range.max, bounds.max),
+    };
+  }
+}  // namespace tf_component_traceviewer

tensorboard/components/tf_trace_viewer/tf-trace-viewer.html

@@ -50,9 +50,15 @@
     transition-delay: 0.3s;
   }
 </style>
+<script src="tf-trace-viewer-helper.js"></script>
 <script>
   "use strict";
 
+  /* tf-trace-viewer will work in two modes: static mode and streaming mode.
+   * in static mode, data are load at 'ready' time,
+   * in streaming mode, data are load on demand when resolution and view port is changed.
+   * static mode limit the amount of trace that we can collect and show to the users.
+   */
   Polymer({
     is: "tf-trace-viewer",
     properties: {
@@ -61,6 +67,7 @@
         type: String,
         value: null,
       },
+      // _traceData is used for static mode.
       _traceData: {
         type: Object,
         observer: "_traceDataChanged"
@@ -69,7 +76,16 @@
       _traceContainer: Object,
       _traceModel: Object,
       _throbber: Object,
+      _isStreaming: { type: Boolean, value: false },
+      _loadedRange: Object,
+      _loadedTraceEents: Object,
+      _fullBounds: Object,
+      _isLoading: { type: Boolean, value: false },
+      _dirty: { type: Boolean, value: false },
+      _model: Object,
+      _resolution: { type: Number, value: 1000 },
     },
+
     ready: function() {
       // Initiate the trace viewer app.
       this._traceContainer = document.createElement("track-view-container");
@@ -97,44 +113,224 @@
           var components = parts[i].split('=');
           if (components[0] == "trace_data_url") {
             this.traceDataUrl = decodeURIComponent(components[1]);
-            break;
+          } else if (components[0] == "is_streaming") {
+            this._isStreaming = components[1] === 'true';
           }
         }
       }
 
+      if (!this.traceDataUrl) {
+        this._displayOverlay("Trace data URL is not provided.", "Trace Viewer");
+        return null;
+      }
       this._throbber.className = "active";
+
       this._loadTrace();
     },
+
     _loadTrace : function() {
-      if (!this.traceDataUrl) {
-        this._displayOverlay("Trace data URL is not provided.", "Trace Viewer");
-        return null;
+      if (!this._isStreaming) {
+        // Send HTTP request to get the trace data.
+        var req = new XMLHttpRequest();
+        req.open('GET', this.traceDataUrl, true);
+
+	req.onreadystatechange = event => {
+          if (req.readyState !== 4) {
+            return;
+          }
+	  window.setTimeout(() => {
+            if (req.status === 200) {
+              this._throbber.className = "inactive";
+              this.set("_traceData", req.responseText);
+            } else {
+              this._displayOverlay(req.status, "Failed to fetch data");
+            }
+          }, 0);
+        };
+        req.send(null);
+      } else {
+        this._loadStreamingTrace();
       }
-      // Send HTTP request to get the trace data.
-      var req = new XMLHttpRequest();
-      var is_binary = / [.] gz$ /.test(this.traceDataUrl) ||
-                      / [.] zip$ /.test(this.traceDataUrl);
-      req.overrideMimeType('text/plain; charset=x-user-defined');
-      req.open('GET', this.traceDataUrl, true);
-      if (is_binary) {
-        req.responseType = 'arraybuffer';
+    },
+
+    // Something has changed, so consider reloading the data:
+    //   - if we have zoomed in enough to need more detail
+    //   - if we have scrolled too close to missing data regions
+    // We ensure there's only ever one request in flight.
+    _maybeLoad : function() {
+      if (this._isLoading || this._resolution == 0) return;
+      // We have several ranges of interest:
+      //             [viewport]           - what's on-screen
+      //         [----preserve----]       - issue loads to keep this full of data
+      //     [---------fetch----------]   - fetch this much data with each load
+      // [-----------full bounds--------] - the whole profile
+      var viewport = this._trackViewRange(this._traceViewer.trackView);
+      var PRESERVE_RATIO = tf_component_traceviewer.PRESERVE_RATIO;
+      var preserve = tf_component_traceviewer.intersect(
+          tf_component_traceviewer.expand(viewport, PRESERVE_RATIO), this._fullBounds);
+      var FETCH_RATIO = tf_component_traceviewer.FETCH_RATIO;
+      var fetch = tf_component_traceviewer.expand(viewport, FETCH_RATIO);
+      var zoomFactor = tf_component_traceviewer.length(this._loadedRange) /
+	               tf_component_traceviewer.length(fetch);
+      if (!tf_component_traceviewer.within(preserve, this._loadedRange) ||
+          zoomFactor > tf_component_traceviewer.ZOOM_RATIO) {
+        console.log("loading more data: ", {
+          zoomFactor: zoomFactor,
+          loadedRange: this._loadedRange,
+          viewport: viewport,
+          preserve: preserve,
+          fetch: fetch,
+        });
+        this._loadTrace(fetch, /*replaceModel=*/false);
       }
+    },
 
-      req.onreadystatechange = function(event) {
-        if (req.readyState !== 4) {
-          return;
-        }
-        window.setTimeout(function() {
-          if (req.status === 200) {
+    _loadStreamingTrace : function(requestedRange, replaceModel) {
+      var success = true;
+      this._isLoading = true;
+
+      this._loadJSON(requestedRange).
+	  then((data) => { this._updateModel(data, replaceModel); }).
+	  then(() => { this._updateView(requestedRange); }).
+	  catch((err) => { this._displayOverlay("Trace Viewer", err);})
+	  .then(() => {
+            this._isLoading = false;
             this._throbber.className = "inactive";
-            this.set("_traceData", is_binary ? req.response : req.responseText);
-          } else {
-            this._displayOverlay(req.status, "Failed to fetch data");
+            // Don't immediately load new data after the very first load. When
+            // we first load the trace viewer, the actual view is not properly
+            // initialized and we get an incorrect viewport leading to a spurious
+            // load of data.
+            if (success && requestedRange) this._maybeLoad();
+          });
+    },
+
+    // Loads a time window (the whole trace if requestedRange is null).
+    // Returns a promise for the JSON event data.
+    _loadJSON : function(requestedRange) {
+      // Set up an XMLHTTPRequest to the JSON endpoint, populating range and
+      // resolution if appropriate.
+      var requestURL = this._buildBaseURL();
+      var ZOOM_RATIO = tf_component_traceviewer.ZOOM_RATIO;
+      requestURL.searchParams.set("resolution", this._resolution * ZOOM_RATIO);
+      if (requestedRange != null) {
+        requestURL.searchParams.set("start_time_ms", requestedRange.min);
+        requestURL.searchParams.set("end_time_ms", requestedRange.max);
+      }
+
+      return new Promise(function(resolve, reject) {
+        var xhr = new XMLHttpRequest();
+        xhr.open('GET', requestURL);
+        xhr.onload = function() {
+          var contentType = this.getResponseHeader('Content-Type');
+          if (this.status !== 200 ||
+              !contentType.startsWith('application/json')) {
+            var msg = requestURL + ' could not be loaded';
+            if (contentType.startsWith('text/plain')) {
+              msg = msg + ': ' + xhr.statusText;
+            }
+            reject(msg);
           }
-        }.bind(this), 0);
-      }.bind(this);
-      req.send(null);
+          resolve(xhr.response);
+        };
+        xhr.onerror = function () {
+          reject(requestURL + 'could not be loaded: ' + xhr.statusText);
+        };
+        xhr.send();
+      });
     },
+    // Decodes the JSON trace events, removes all events that were loaded before
+    // and serializes to JSON again.
+    _filterKnownTraceEvents: function(data) {
+      var traceEvents = data.traceEvents;
+      data.traceEvents = [];
+      for (var i = 0; i < traceEvents.length; i++) {
+        // This is inefficient as we are serializing the events we just
+        // deserialized. If this becomes a problem in practice, we should assign
+        // IDs on the server.
+        var asString = JSON.stringify(traceEvents[i]);
+        if (!this._loadedTraceEvents.has(asString)) {
+          this._loadedTraceEvents.add(asString);
+          data.traceEvents.push(traceEvents[i]);
+        }
+      }
+      return data;
+    },
+
+    // Updates the model with data returned by the JSON endpoint.
+    // If replaceModel is true, the data set is completely replaced; otherwise,
+    // the new data is merged with the old data.
+    // Returns a void promise.
+    _updateModel: function(data, replaceModel) {
+      data = JSON.parse(data);
+      if (!this._model /* first load */ || replaceModel) {
+        this._dirty = true;
+        this._model = new tr.Model();
+        this._loadedTraceEvents = new Set();
+      } else {
+        // Delete metadata and displayTimeUnits as otherwise traceviewer
+        // accumulates them.
+        delete data['metadata'];
+        delete data['displayTimeUnit'];
+      }
+
+      data = this._filterKnownTraceEvents(data);
+      if (data.traceEvents.length > 0) {
+        var opt = new tr.importer.ImportOptions();
+        opt.shiftWorldToZero = false;
+        new tr.importer.Import(this._model, opt).importTraces([data]);
+        this._dirty = true;
+      }
+      return Promise.resolve();
+    },
+
+    // Updates the view based on the current model.
+    _updateView: function(requestedRange) {
+      if (requestedRange == null) {
+        this._fullBounds = {min: this._model.bounds.min, max: this._model.bounds.max};
+        this._loadedRange = tf_component_traceviewer.expand(
+	    this._fullBounds, tf_component_traceviewer.FETCH_RATIO);
+      } else {
+        this._loadedRange = requestedRange;
+      }
+      if (!this._dirty){
+        return;
+      }
+      this._dirty = false;
+      // We can't assign the model until the viewer is attached. This may be
+      // delayed indefinitely if the tab is backgrounded. This version of polymer
+      // doesn't provide a direct way to observe the viewer being attached.
+      // This is a hack: the browser won't paint until the viewer is attached.
+      window.requestAnimationFrame(function() {
+        this._traceViewer.model = this._model;
+        if (this._traceViewer.trackView != null) {  // Only initialized if data in nonempty!
+          // Wait 200ms to let an animated zoom/pan operation complete. Ideally,
+          // we could just explicitly wait for its end.
+
+          this._traceViewer.trackView.viewport.addEventListener(
+              "change", () => setTimeout(this._maybeLoad.bind(this), 200));
+        }
+        this._traceViewer.viewTitle = "";
+      }.bind(this));
+    },
+
+    // Access the {min, max} range of a trackView.
+    _trackViewRange: function(trackView) {
+      var xfm = trackView.viewport.currentDisplayTransform;
+      const pixelRatio = window.devicePixelRatio || 1;
+      const devicePixelWidth = pixelRatio * trackView.viewWidth_;
+      return {
+        min: xfm.xViewToWorld(0),
+        max: xfm.xViewToWorld(devicePixelWidth),
+      };
+    },
+
+    // Builds a base URL for fetching json data. The URL will be assembled with
+    // all filtering URL parameters, except resolution and range.
+    _buildBaseURL: function() {
+      var requestURL = new URL(this.traceDataUrl, window.location.href);
+      return requestURL;
+    },
+
     _traceDataChanged: function(data) {
       if (!data) {
         this._displayOverlay("Trace Viewer", "No trace to display...");
@@ -152,6 +348,7 @@
             'Import error', tr.b.normalizeException(err).message);
       });
     },
+
     _displayOverlay: function(title, content) {
       var overlay = new tr.ui.b.Overlay();
       overlay.textContent = content;