From 6eb8c2bab657ea5a2c596651959bb5bb74de9a07 Mon Sep 17 00:00:00 2001
From: William Chargin <wchargin@gmail.com>
Date: Wed, 23 Sep 2020 16:53:48 -0700
Subject: [PATCH] =?UTF-8?q?scalars:=20don=E2=80=99t=20multiplex=20fetches?=
 =?UTF-8?q?=20in=20Colab?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Summary:
In #4050, we added multiplexing to scalar chart fetches for performance.
This works nicely in local TensorBoard and public Colab, but the POST
request machinery isn’t yet supported in the Google-internal version of
Colab, so #4050 caused a regression there. This patch conditionally
rolls back the change for Colab only. This includes rolling it back for
public Colab, where it worked fine. We hope that this isn’t too big of a
problem, since we expect that most Colab users are inspecting datasets
with few runs (generated from within Colab) rather than massive
hyperparameter sweeps. We also hope that we can simply revert this patch
once the Google-internal behavior is fixed.

Jupyter environments are unaffected (and still work).

We used to have a global `window.TENSORBOARD_ENV` that listed whether we
were in Colab, but we removed that in #2798 because we thought that it
was no longer necessary. Since in that PR we switched to create an
iframe rather than manually linking and loading the HTML (much nicer, to
be sure), we now plumb this information via a query parameter. This also
has the advantage that it’s easy to test the Colab codepaths by simply
adding that query parameter to a normal TensorBoard instance.

Test Plan:
First, test that normal TensorBoard (`bazel run //tensorboard`) still
works with multiplexing, and that adding `?tensorboardColab=true` causes
it to send multiple GET requests instead. Then, build the Pip package
(`bazel run //tensorboard/pip_package:extract_pip_packages`), upload it
to public Colab, and install it into the runtime. Verify that the scalar
charts still work there, albeit without multiplexing. Finally,
cherry-pick these changes into google3 via a test sync and follow the
test plan at <http://cl/333398676> to verify the fix.

wchargin-branch: scalars-nomux-colab
wchargin-source: 453503f1dea196985e889be483c2a7675cc87aa1
---
 tensorboard/notebook.py                       |  3 +-
 .../tf_scalar_dashboard/tf-scalar-card.ts     | 34 +++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/tensorboard/notebook.py b/tensorboard/notebook.py
index 8e6e024901..b4c09d4558 100644
--- a/tensorboard/notebook.py
+++ b/tensorboard/notebook.py
@@ -345,7 +345,8 @@ def _display_colab(port, height, display_handle):
 
     shell = """
         (async () => {
-            const url = await google.colab.kernel.proxyPort(%PORT%, {"cache": true});
+            const url = new URL(await google.colab.kernel.proxyPort(%PORT%, {'cache': true}));
+            url.searchParams.set('tensorboardColab', 'true');
             const iframe = document.createElement('iframe');
             iframe.src = url;
             iframe.setAttribute('width', '100%');
diff --git a/tensorboard/plugins/scalar/tf_scalar_dashboard/tf-scalar-card.ts b/tensorboard/plugins/scalar/tf_scalar_dashboard/tf-scalar-card.ts
index 4e88141549..dab3471429 100644
--- a/tensorboard/plugins/scalar/tf_scalar_dashboard/tf-scalar-card.ts
+++ b/tensorboard/plugins/scalar/tf_scalar_dashboard/tf-scalar-card.ts
@@ -264,6 +264,40 @@ export class TfScalarCard extends PolymerElement {
     items,
     onLoad,
     onFinish
+  ) => {
+    const inColab =
+      new URLSearchParams(window.location.search).get('tensorboardColab') ===
+      'true';
+    if (inColab) {
+      // Google-internal Colab doesn't support HTTP POST requests, so we fall
+      // back to HTTP GET (even though public Colab supports POST).
+      return this._requestDataGet(items, onLoad, onFinish);
+    } else {
+      return this._requestDataPost(items, onLoad, onFinish);
+    }
+  };
+
+  _requestDataGet: RequestDataCallback<RunTagItem, ScalarDatum[] | null> = (
+    items,
+    onLoad,
+    onFinish
+  ) => {
+    const router = getRouter();
+    const baseUrl = router.pluginRoute('scalars', '/scalars');
+    Promise.all(
+      items.map((item) => {
+        const url = addParams(baseUrl, {tag: item.tag, run: item.run});
+        return this.requestManager
+          .request(url)
+          .then((data) => void onLoad({item, data}));
+      })
+    ).finally(() => void onFinish());
+  };
+
+  _requestDataPost: RequestDataCallback<RunTagItem, ScalarDatum[] | null> = (
+    items,
+    onLoad,
+    onFinish
   ) => {
     const router = getRouter();
     const url = router.pluginRoute('scalars', '/scalars_multirun');