Allow Jupyter+TensorBoard to work w/ hostname other than localhost

Signed-off-by: Cliff Woolley <jwoolley@nvidia.com>

Author: cliffwoolley

tensorboard/notebook.py

@@ -395,11 +395,31 @@ def _display_colab(port, height, display_handle):
 
 def _display_ipython(port, height, display_handle):
   import IPython.display
-  iframe = IPython.display.IFrame(
-      src="http://localhost:%d" % port,
-      height=height,
-      width="100%",
-  )
+  import cgi
+  import json
+  import random
+
+  frame_id = "tensorboard-frame-{:08x}".format(random.randrange(0, 1 << 64))
+  shell = """
+    <iframe id="%HTML_ID%" width="100%" height="%HEIGHT%" frameborder="0"></iframe>
+    <script>
+      (function() {
+        const frame = document.getElementById(%JSON_ID%);
+        const url = new URL("/", window.location);
+        url.port = %PORT%;
+        frame.src = url;
+      })();
+    </script>
+  """
+  replacements = [
+    ("%HTML_ID%", cgi.escape(frame_id, quote=True)),
+    ("%JSON_ID%", json.dumps(frame_id)),
+    ("%PORT%", "%d" % port),
+    ("%HEIGHT%", "%d" % height),
+  ]
+  for (k, v) in replacements:
+    shell = shell.replace(k, v)
+  iframe = IPython.display.HTML(shell)
   if display_handle:
     display_handle.update(iframe)
   else: