Make it possible to configure how rerun_notebook fetches the widget (…

…#6680)

Co-authored-by: Jan Procházka <honza.spacir@gmail.com>

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()