From db2251ac7cb60e88c5c96e68778ee1fc2c79e7aa Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Mon, 1 Jul 2024 11:03:44 +0200 Subject: [PATCH] Make it possible to configure how `rerun_notebook` fetches the widget (#6680) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Jan Procházka --- rerun_js/web-viewer/index.ts | 9 ++++ rerun_notebook/README.md | 45 ++++++++++++++++ rerun_notebook/package-lock.json | 2 +- rerun_notebook/src/rerun_notebook/__init__.py | 41 +++++++++++++-- rerun_notebook/src/rerun_notebook/__main__.py | 26 ++++++++++ .../src/rerun_notebook/asset_server.py | 52 +++++++++++++++++++ scripts/ci/publish_wheels.py | 23 ++++++-- 7 files changed, 190 insertions(+), 8 deletions(-) create mode 100644 rerun_notebook/src/rerun_notebook/__main__.py create mode 100644 rerun_notebook/src/rerun_notebook/asset_server.py diff --git a/rerun_js/web-viewer/index.ts b/rerun_js/web-viewer/index.ts index 446c1a6a9dcb..2f3837ff61e1 100644 --- a/rerun_js/web-viewer/index.ts +++ b/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; diff --git a/rerun_notebook/README.md b/rerun_notebook/README.md index 4ff01f0240a4..2e6d4aa2922d 100644 --- a/rerun_notebook/README.md +++ b/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: diff --git a/rerun_notebook/package-lock.json b/rerun_notebook/package-lock.json index 52e3302c605b..3f7ce291c556 100644 --- a/rerun_notebook/package-lock.json +++ b/rerun_notebook/package-lock.json @@ -17,7 +17,7 @@ }, "../rerun_js/web-viewer": { "name": "@rerun-io/web-viewer", - "version": "0.16.0", + "version": "0.17.0-alpha.8", "license": "MIT", "devDependencies": { "dts-buddy": "^0.3.0", diff --git a/rerun_notebook/src/rerun_notebook/__init__.py b/rerun_notebook/src/rerun_notebook/__init__.py index d7d36afe6bfb..d90939de8ca9 100644 --- a/rerun_notebook/src/rerun_notebook/__init__.py +++ b/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) diff --git a/rerun_notebook/src/rerun_notebook/__main__.py b/rerun_notebook/src/rerun_notebook/__main__.py new file mode 100644 index 000000000000..92e695b10b45 --- /dev/null +++ b/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() diff --git a/rerun_notebook/src/rerun_notebook/asset_server.py b/rerun_notebook/src/rerun_notebook/asset_server.py new file mode 100644 index 000000000000..7fe52a4481ff --- /dev/null +++ b/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 diff --git a/scripts/ci/publish_wheels.py b/scripts/ci/publish_wheels.py index 12216ed841ff..b84b19be565e 100755 --- a/scripts/ci/publish_wheels.py +++ b/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()