A few bug fixes and improvements of TPU profile plugin (#1245)

1) Fix a bug that streaming trace viewer loads wrong host names
2) Fix a bug that streaming trace viewer doesn't fetch dynamic traces
3) Add throbler to streaming trace viewer when waiting for data to load
4) Show overview page always as the first tool
5) Improve color code of overview page: green => good, orange => moderate
   red => bad

Author: qiuminxu

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

@@ -123,12 +123,11 @@
         this._displayOverlay("Trace data URL is not provided.", "Trace Viewer");
         return null;
       }
-      this._throbber.className = "active";
-
       this._loadTrace();
     },
 
-    _loadTrace : function() {
+    _loadTrace : function(requestedRange, replaceModel) {
+      this._throbber.className = "active";
       if (!this._isStreaming) {
         // Send HTTP request to get the trace data.
         var req = new XMLHttpRequest();
@@ -149,7 +148,7 @@
         };
         req.send(null);
       } else {
-        this._loadStreamingTrace();
+        this._loadStreamingTrace(requestedRange, replaceModel);
       }
     },
 
@@ -174,13 +173,6 @@
 	               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);
       }
     },

tensorboard/plugins/profile/overview_page/overview-page.html

@@ -44,7 +44,7 @@
      .steptime-average {
        font-weight: bold;
        font-style: italic;
-       color: red;
+       color: var(--summary-color, green);
      }
      .table-style {
        table-layout: auto;
@@ -71,7 +71,7 @@
        text-transform: capitalize;
      }
      div.bottleneck-statement {
-       color: red;
+       color: var(--summary-color, green);
        font-weight: bolder;
        font-style: italic;
      }
@@ -219,6 +219,7 @@
        var runEnvironmentJson = this._data[2];
        var recommendationJson = this._data[3];
 
+       var statement = recommendationJson.p['statement'];
        this.set('_host_idle_time_percent', generalAnalysisJson.p['host_idle_time_percent']);
        this.set('_device_idle_time_percent', generalAnalysisJson.p['device_idle_time_percent']);
        this.set('_mxu_utilization_percent', generalAnalysisJson.p['mxu_utilization_percent']);
@@ -234,6 +235,13 @@
        this.set('_build_target', runEnvironmentJson.p['build_target']);
        this.set('_statement', recommendationJson.p['statement']);
        this.set('_show_error_message', this._error_message!='');
+
+       if (statement.includes('HIGHLY')) {
+         this.customStyle['--summary-color'] = 'red';
+       } else if (statement.includes('MODERATE')) {
+         this.customStype['--summary-color'] = 'orange';
+       } 
+
        this.updateStyles();
 
        this._showDeviceStepChart(inputAnalysisJson);

tensorboard/plugins/profile/profile_plugin.py

@@ -180,11 +180,18 @@ def index_impl(self):
         removed_tool = 'trace_viewer@' if self.stub is None else 'trace_viewer'
         if removed_tool in run_to_tools[run]:
           run_to_tools[run].remove(removed_tool)
+      run_to_tools[run].sort()
+      op = 'overview_page'
+      if op in run_to_tools[run]:
+        # keep overview page at the top of the list
+        run_to_tools[run].remove(op)
+        run_to_tools[run].insert(0, op)
     return run_to_tools
 
   @wrappers.Request.application
   def tools_route(self, request):
     run_to_tools = self.index_impl()
+
     return http_util.Respond(request, run_to_tools, 'application/json')
 
   def host_impl(self, run, tool):
@@ -260,7 +267,8 @@ def data_impl(self, request):
       grpc_request.repository_root = self.plugin_logdir
       grpc_request.session_id = run[:-1]
       grpc_request.tool_name = 'trace_viewer'
-      grpc_request.host_name = host
+      # Remove the trailing dot if present
+      grpc_request.host_name = host.rstrip('.')
 
       grpc_request.parameters['resolution'] = request.args.get('resolution')
       if request.args.get('start_time_ms') is not None: