Make it possible to configure how rerun_notebook fetches the widget

rerun_js/web-viewer/index.ts

@@ -76,6 +76,10 @@ type EventsWithoutValue = {
 
 type Cancel = () => void;
 
+function delay(ms: number) {
+  return new Promise(resolve => setTimeout(resolve, ms));
+}
+
 export class WebViewer {
   #id = randomId();
   #handle: WebHandle | null = null;
@@ -115,6 +119,11 @@ export class WebViewer {
     this.#canvas.id = this.#id;
     parent.append(this.#canvas);
 
+    // This yield appears to be necessary to ensure that the canvas is attached to the DOM
+    // and visible. Without it we get occasionally get a panic about a failure to find a canvas
+    // element with the given ID.
+    await delay(0);
+
     let WebHandle_class = await load();
     if (this.#state !== "starting") return;
 

rerun_notebook/README.md

@@ -18,6 +18,51 @@ There are several reasons for this package to be separate from the main `rerun-s
 - `rerun-notebook` uses [hatch](https://hatch.pypa.io/) as package backend, and benefits from the [hatch-jupyter-builder](https://github.com/jupyterlab/hatch-jupyter-builder) plug-in. Since `rerun-sdk` must use [Maturin](https://www.maturin.rs), it would the package management more complex.
 - Developer experience: building `rerun-notebook` implies building `rerun_js`, which is best avoided when iterating on `rerun-sdk` outside of notebook environments.
 
+## Ways to access the widget asset
+
+Even though `rerun_notebook` ships with the widget asset bundled in, by default it will try to load the asset
+from `https://app.rerun.io`. This is because the way anywiget transmits the asset at the moment results in
+[a memory leak](https://github.com/manzt/anywidget/issues/613) of the entire module for each cell execution.
+
+If your network does not allow you to access `app.rerun.io`, the behavior can be changed by setting the
+the `RERUN_NOTEBOOK_ASSET` environment variable before you import `rerun_notebook`. This variable must
+be set prior to your import because `AnyWidget` stores the resource on the widget class instance
+once at import time.
+
+### Inlined asset
+Setting:
+```
+RERUN_NOTEBOOK_ASSET=inline
+```
+Will cause `rerun_notebook` to directly transmit the inlined asset to the widget over Jupyter comms.
+This will be the most portable way to use the widget, but is currently known to leak memory and
+has some performance issues in environments such as Google colab.
+
+### Locally served asset
+Setting:
+```
+RERUN_NOTEBOOK_ASSET=serve-local
+```
+Will cause `rerun_notebook` to launch a thread serving the asset from the local machine during
+the lifetime of the kernel. This will be the best way to use the widget in a notebook environment
+when your notebook server is running locally.
+
+### Manually hosted asset
+Setting:
+```
+RERUN_NOTEBOOK_ASSET=https://your-hosted-asset-url.com/widget.js
+```
+Will cause `rerun_notebook` to load the asset from the provided URL. This is the most flexible way to
+use the widget, but requires you to host the asset yourself.
+
+The `rerun_notebook` package has a minimal server that can be used to serve the asset nanually by running:
+```
+python -m rerun_notebook serve
+```
+
+However, any hosting platform can be used to serve the asset, as long as it is accessible to the notebook
+and has appropriate CORS headers set. See: `asset_server.py` for a simple example.
+
 ## Run from source
 
 Use pixi:

rerun_notebook/src/rerun_notebook/__init__.py

@@ -2,8 +2,8 @@
 
 import importlib.metadata
 import logging
+import os
 import pathlib
-import sys
 import time
 from typing import Any, Literal
 
@@ -20,10 +20,39 @@
 Panel = Literal["top", "blueprint", "selection", "time"]
 PanelState = Literal["expanded", "collapsed", "hidden"]
 
+WIDGET_PATH = pathlib.Path(__file__).parent / "static" / "widget.js"
+CSS_PATH = pathlib.Path(__file__).parent / "static" / "widget.css"
+
+# We need to bootstrap the value of ESM_MOD before the Viewer class is instantiated.
+# This is because AnyWidget will process the resource files during `__init_subclass__`
+
+# We allow customization through the `RERUN_NOTEBOOK_ASSET` environment variable.
+# The default value is hosted at `app.rerun.io`.
+# The value can be set to `serve-local` to use the hosted asset server.
+# The value can be set to `inline` to use the local widget.js file.
+# The value can be set to a URL to use a custom asset server.
+#  One way to run a custom asset server is to run `python -m rerun_notebook serve`.
+ASSET_MAGIC_SERVE = "serve-local"
+ASSET_MAGIC_INLINE = "inline"
+
+ASSET_ENV = os.environ.get("RERUN_NOTEBOOK_ASSET", f"https://app.rerun.io/version/{__version__}/widget.js")
+
+if ASSET_ENV == ASSET_MAGIC_SERVE:
+    from .asset_server import serve_assets
+
+    bound_addr = serve_assets(background=True)
+    ESM_MOD = f"http://localhost:{bound_addr[1]}/widget.js"
+elif ASSET_ENV == ASSET_MAGIC_INLINE:
+    ESM_MOD = WIDGET_PATH
+else:
+    ESM_MOD = ASSET_ENV
+    if not (ASSET_ENV.startswith("http://") or ASSET_ENV.startswith("https://")):
+        raise ValueError(f"RERUN_NOTEBOOK_ASSET_URL should be a URL starting with http or https. Found: {ASSET_ENV}")
+
 
 class Viewer(anywidget.AnyWidget):
-    _esm = pathlib.Path(__file__).parent / "static" / "widget.js"
-    _css = pathlib.Path(__file__).parent / "static" / "widget.css"
+    _esm = ESM_MOD
+    _css = CSS_PATH
 
     _width = traitlets.Int(allow_none=True).tag(sync=True)
     _height = traitlets.Int(allow_none=True).tag(sync=True)
@@ -85,7 +114,11 @@ def block_until_ready(self, timeout=5.0) -> None:
         with jupyter_ui_poll.ui_events() as poll:
             while self._ready is False:
                 if time.time() - start > timeout:
-                    logging.warning("Timed out waiting for viewer to become ready.")
+                    logging.warning(
+                        f"""Timed out waiting for viewer to become ready. Make sure: {ESM_MOD} is accessible.
+If not, consider setting `RERUN_NOTEBOOK_ASSET`. Consult https://pypi.org/project/rerun-notebook/{__version__}/ for details.
+"""
+                    )
                     return
                 poll(1)
                 time.sleep(0.1)

rerun_notebook/src/rerun_notebook/__main__.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+import argparse
+
+from .asset_server import serve_assets
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    subparsers = parser.add_subparsers(dest="command", help="Which command to run")
+
+    serve_parser = subparsers.add_parser("serve")
+    serve_parser.add_argument("--bind-address", default="localhost")
+    serve_parser.add_argument("--port", type=int, default=8080)
+
+    args = parser.parse_args()
+
+    if args.command == "serve":
+        serve_assets(args.bind_address, args.port)
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()

rerun_notebook/src/rerun_notebook/asset_server.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import http.server
+import socketserver
+
+from . import WIDGET_PATH
+
+resource_data: bytes | None = None
+
+
+class AssetHandler(http.server.SimpleHTTPRequestHandler):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def do_GET(self):
+        if self.path == "/widget.js":  # remap this path
+            self.send_response(200)
+            self.send_header("Access-Control-Allow-Origin", "*")
+            self.send_header("Content-type", "text/javascript")
+            self.end_headers()
+            if resource_data is not None:
+                self.wfile.write(resource_data)
+        else:
+            # Serve other requests normally
+            self.send_error(404, "File Not Found")
+
+    def log_message(self, format, *args):
+        # Disable logging
+        return
+
+
+def serve_assets(bind_address: str = "localhost", port: int = 0, background=False) -> socketserver._AfInetAddress:
+    global resource_data
+    if resource_data is None:
+        with open(WIDGET_PATH, "rb") as f:
+            resource_data = f.read()
+
+    httpd = socketserver.TCPServer((bind_address, port), AssetHandler)
+
+    bound_addr = httpd.server_address
+    print(f"Serving rerun notebook assets at http://{bound_addr[0]}:{bound_addr[1]}")
+
+    if background:
+        import threading
+
+        thread = threading.Thread(target=httpd.serve_forever)
+        thread.daemon = True
+        thread.start()
+    else:
+        httpd.serve_forever()
+
+    return bound_addr

scripts/ci/publish_wheels.py

@@ -14,6 +14,7 @@
 import shutil
 import subprocess
 import sys
+import zipfile
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 
@@ -35,15 +36,29 @@ def run(
 def check_version(expected_version: str) -> None:
     wheels = list(Path("wheels").glob("*.whl"))
 
-    for wheel in wheels:
-        wheel_version = wheel.stem.split("-")[1]
+    for whl in wheels:
+        wheel_version = whl.stem.split("-")[1]
         if canonicalize_version(wheel_version) != expected_version:
-            print(f"Unexpected version: {wheel_version} (expected: {expected_version}) in {wheel.name}")
+            print(f"Unexpected version: {wheel_version} (expected: {expected_version}) in {whl.name}")
             sys.exit(1)
 
     print(f"All wheel versions match the expected version: {expected_version}")
 
 
+def publish_notebook_asset() -> None:
+    bucket = Gcs("rerun-open").bucket("rerun-web-viewer")
+    wheels = list(Path("wheels").glob("*.whl"))
+    for whl in wheels:
+        if whl.name.startswith("rerun_notebook-"):
+            wheel_version = whl.stem.split("-")[1]
+            with zipfile.ZipFile(whl, "r") as archive:
+                # Extract the specified file to the target directory
+                archive.extract("rerun_notebook/static/widget.js", "extracted")
+                bucket.blob(f"version/{wheel_version}/widget.js").upload_from_filename(
+                    "extracted/rerun_notebook/static/widget.js"
+                )
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(description="Publish wheels to PyPI")
     parser.add_argument("--version", required=True, help="Version to expect")
@@ -76,6 +91,8 @@ def main() -> None:
         },
     )
 
+    publish_notebook_asset()
+
 
 if __name__ == "__main__":
     main()