Support proxied URLs for TensorBoard running in IPython (#3674)

Users will frequently want to run TensorBoard in environments where they are not allowed to open up or expose additional ports. For those cases jupyter-server-proxy (and other proxy software) allows forwarding web traffic to other services running on the system without opening new ports. (The Google Colab integration also works around this same issue of not being able to directly expose the port that TensorBoard runs on.)

This PR allows users to optionally pass the TensorBoard traffic through a proxy by setting the TENSORBOARD_PROXY_URL environment variable. When this variable is not set TensorBoard will behave as normal.

This adds support for a new optional environment variable TENSORBOARD_PROXY_URL which if set will be used as the URL's path for the IFRAME displayed in the notebook when running %tensorboard.

tensorboard/notebook.py

@@ -23,6 +23,7 @@
 import datetime
 import errno
 import json
+import os
 import random
 import shlex
 import sys
@@ -378,18 +379,35 @@ def _display_ipython(port, height, display_handle):
       <script>
         (function() {
           const frame = document.getElementById(%JSON_ID%);
-          const url = new URL("/", window.location);
-          url.port = %PORT%;
+          const url = new URL(%URL%, window.location);
+          const port = %PORT%;
+          if (port) {
+            url.port = port;
+          }
           frame.src = url;
         })();
       </script>
-  """
-    replacements = [
-        ("%HTML_ID%", html_escape(frame_id, quote=True)),
-        ("%JSON_ID%", json.dumps(frame_id)),
-        ("%PORT%", "%d" % port),
-        ("%HEIGHT%", "%d" % height),
-    ]
+    """
+    proxy_url = os.environ.get("TENSORBOARD_PROXY_URL")
+    if proxy_url is not None:
+        # Allow %PORT% in $TENSORBOARD_PROXY_URL
+        proxy_url = proxy_url.replace("%PORT%", "%d" % port)
+        replacements = [
+            ("%HTML_ID%", html_escape(frame_id, quote=True)),
+            ("%JSON_ID%", json.dumps(frame_id)),
+            ("%HEIGHT%", "%d" % height),
+            ("%PORT%", "0"),
+            ("%URL%", json.dumps(proxy_url)),
+        ]
+    else:
+        replacements = [
+            ("%HTML_ID%", html_escape(frame_id, quote=True)),
+            ("%JSON_ID%", json.dumps(frame_id)),
+            ("%HEIGHT%", "%d" % height),
+            ("%PORT%", "%d" % port),
+            ("%URL%", json.dumps("/")),
+        ]
+
     for (k, v) in replacements:
         shell = shell.replace(k, v)
     iframe = IPython.display.HTML(shell)