colab: use proxyPort for dynamic plugin (#2798)

proxyPort is a relatively new API that lets colab proxy a server from a
port to url given by the method. It installs a ServiceWorker when
invoked so a request is proxied to the correct kernel.

This API an iframe to be nestable and, hence, enable dynamic plugin

With this change, we now removed the colab not supported page.

Author: stephanwlee

tensorboard/components/tf_tensorboard/tf-tensorboard.html

@@ -428,21 +428,6 @@ <h3>
         }
       }
     </style>
-    <template id="colabNotSupported" preserve-content>
-      <div style="max-width: 540px; margin: 80px auto 0 auto;">
-        <h3>Dynamic plugin isn’t supported in Colab yet.</h3>
-        <p>
-          Please see
-          <a
-            href="https://github.com/tensorflow/tensorboard/issues/1913"
-            rel="noopener"
-            target="_blank"
-            >GitHub issue #1913</a
-          >
-          for more information.
-        </p>
-      </div>
-    </template>
   </template>
   <script src="autoReloadBehavior.js"></script>
   <script>
@@ -659,11 +644,6 @@ <h3>Dynamic plugin isn’t supported in Colab yet.</h3>
           type: Boolean,
           value: false,
         },
-        _inColab: {
-          type: Boolean,
-          value: () => !!(window.TENSORBOARD_ENV || {}).IN_COLAB,
-          readOnly: true,
-        },
       },
       observers: [
         '_updateSelectedDashboardFromActive(' +
@@ -877,14 +857,6 @@ <h3>Dynamic plugin isn’t supported in Colab yet.</h3>
       },
 
       _renderPluginIframe(container, selectedDashboard, loadingMechanism) {
-        if (this._inColab) {
-          const templateContent = this.$.colabNotSupported.content;
-          const node = document.importNode(templateContent, true);
-          node.id = 'dashboard'; // used in `_selectedDashboardComponent`
-          this.scopeSubtree(node);
-          container.appendChild(node);
-          return;
-        }
         const iframe = document.createElement('iframe');
         iframe.id = 'dashboard'; // used in `_selectedDashboardComponent`
         const srcUrl = new URL('data/plugin_entry.html', window.location.href);

tensorboard/notebook.py

@@ -339,61 +339,28 @@ def _display_colab(port, height, display_handle):
   """
   import IPython.display
   shell = """
-    <div id="root"></div>
-    <script>
-      (function() {
-        window.TENSORBOARD_ENV = window.TENSORBOARD_ENV || {};
-        window.TENSORBOARD_ENV["IN_COLAB"] = true;
-        document.querySelector("base").href = "https://localhost:%PORT%";
-        function fixUpTensorboard(root) {
-          const tftb = root.querySelector("tf-tensorboard");
-          // Disable the fragment manipulation behavior in Colab. Not
-          // only is the behavior not useful (as the iframe's location
-          // is not visible to the user), it causes TensorBoard's usage
-          // of `window.replace` to navigate away from the page and to
-          // the `localhost:<port>` URL specified by the base URI, which
-          // in turn causes the frame to (likely) crash.
-          tftb.removeAttribute("use-hash");
-        }
-        function executeAllScripts(root) {
-          // When `script` elements are inserted into the DOM by
-          // assigning to an element's `innerHTML`, the scripts are not
-          // executed. Thus, we manually re-insert these scripts so that
-          // TensorBoard can initialize itself.
-          for (const script of root.querySelectorAll("script")) {
-            const newScript = document.createElement("script");
-            newScript.type = script.type;
-            newScript.textContent = script.textContent;
-            root.appendChild(newScript);
-            script.remove();
-          }
-        }
-        function setHeight(root, height) {
-          // We set the height dynamically after the TensorBoard UI has
-          // been initialized. This avoids an intermediate state in
-          // which the container plus the UI become taller than the
-          // final width and cause the Colab output frame to be
-          // permanently resized, eventually leading to an empty
-          // vertical gap below the TensorBoard UI. It's not clear
-          // exactly what causes this problematic intermediate state,
-          // but setting the height late seems to fix it.
-          root.style.height = `${height}px`;
-        }
-        const root = document.getElementById("root");
-        fetch(".")
-          .then((x) => x.text())
-          .then((html) => void (root.innerHTML = html))
-          .then(() => fixUpTensorboard(root))
-          .then(() => executeAllScripts(root))
-          .then(() => setHeight(root, %HEIGHT%));
-      })();
-    </script>
-  """.replace("%PORT%", "%d" % port).replace("%HEIGHT%", "%d" % height)
-  html = IPython.display.HTML(shell)
+  (async () => {
+      const url = await google.colab.kernel.proxyPort(%PORT%, {"cache": true});
+      const iframe = document.createElement('iframe');
+      iframe.src = url;
+      iframe.setAttribute('width', '100%');
+      iframe.setAttribute('height', '%HEIGHT%');
+      iframe.setAttribute('frameborder', 0);
+      document.body.appendChild(iframe);
+  })();
+  """
+  replacements = [
+      ("%PORT%", "%d" % port),
+      ("%HEIGHT%", "%d" % height),
+  ]
+  for (k, v) in replacements:
+    shell = shell.replace(k, v)
+  script = IPython.display.Javascript(shell)
+
   if display_handle:
-    display_handle.update(html)
+    display_handle.update(script)
   else:
-    IPython.display.display(html)
+    IPython.display.display(script)
 
 
 def _display_ipython(port, height, display_handle):

tensorboard/plugins/profile/tf_profile_dashboard/tf-profile-dashboard.html

@@ -49,21 +49,6 @@
 
 <dom-module id="tf-profile-dashboard">
   <template>
-    <template is="dom-if" if="[[_isState(_topLevelState, 'IN_COLAB')]]">
-      <div style="max-width: 540px; margin: 80px auto 0 auto;">
-        <h3>Profiling isn’t supported in Colab yet.</h3>
-        <p>
-          Please see
-          <a
-            href="https://github.com/tensorflow/tensorboard/issues/1913"
-            rel="noopener"
-            target="_blank"
-            >GitHub issue #1913</a
-          >
-          for more information.
-        </p>
-      </div>
-    </template>
     <paper-toast id="toast" duration="0" text="" always-on-top>
       <paper-button onclick="toast.toggle()" class="yellow-button"
         >Close now!</paper-button
@@ -462,13 +447,6 @@ <h3>No profile data was found.</h3>
      * @enum {string}
      */
     const TopLevelState = {
-      /**
-       * Indicates that we are in a Colab notebook environment. The
-       * profile dashboard does not work well in Colab; the trace viewer
-       * does not work at all. If in Colab, we'll show only an explanatory
-       * message.
-       */
-      IN_COLAB: 'IN_COLAB',
       /**
        * Indicates that there are no runs with profile data.
        */
@@ -514,16 +492,10 @@ <h3>No profile data was found.</h3>
             type: Array,
             observer: '_activeHostsChanged',
           },
-          _inColab: {
-            type: Boolean,
-            value: () => !!(window.TENSORBOARD_ENV || {}).IN_COLAB,
-            readOnly: true,
-          },
           /** @type {TopLevelState} */
           _topLevelState: {
             type: String,
-            computed:
-              '_computeTopLevelState(_inColab, _dataNotFound, progress)',
+            computed: '_computeTopLevelState(_dataNotFound, progress)',
             readOnly: true,
           },
           /**
@@ -950,8 +922,7 @@ <h3>No profile data was found.</h3>
           // Otherwise, remove the seperator, e.g. "host1_" => "host1".
           return host.slice(0, -1);
         },
-        _computeTopLevelState: function(inColab, dataNotFound, progress) {
-          if (inColab) return 'IN_COLAB';
+        _computeTopLevelState: function(dataNotFound, progress) {
           if (dataNotFound) return 'DATA_NOT_FOUND';
           if (!progress || progress.value < 100) return 'LOADING';
           return 'ACTIVE';