colab: use proxyPort for dynamic plugin

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';