From 4edf07874f8db2f3bd2db058afc9178e8b4e0c5c Mon Sep 17 00:00:00 2001 From: Anders Hafreager Date: Thu, 9 Jul 2026 07:45:15 +0200 Subject: [PATCH] Make the notebook site cross-origin isolated on GitHub Pages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The KOKKOS multithreaded wasm was shipped to the notebook site but never usable there: SharedArrayBuffer needs COOP/COEP headers, static hosting can't send them, and the usual coi-serviceworker shim can't be installed because JupyterLite's own service worker owns the scope (the contents API — and our DriveFS mount — depend on it). Fold the coi-serviceworker logic into JupyterLite's service worker instead: build.sh now runs coi_patch.py, which wraps maybeFromCache (the function serving every plain GET the worker intercepts) so responses carry COOP/COEP/CORP, and injects a bootstrap into the app pages that reloads once when the worker takes control so the document itself gets the headers. A sessionStorage guard prevents reload loops; without isolation everything falls back to single-threaded as before. The patch asserts on the pinned jupyterlite-core's service-worker internals so a version bump fails the build loudly (upstream feature request: jupyterlite/jupyterlite#1409). Verified headlessly against a plain http.server with no headers: crossOriginIsolated flips to true after one reload; basics/04 reports the KOKKOS build available and benchmarks 6.3 M atom-steps/s vs 2.6 M single-threaded (~2.4x at t 4); the DriveFS file workflows (basics/02, basics/05, md-basics/05), piplite, matplotlib and the pyodide CDN all work under require-corp; md-basics/02's "go big" cell now runs its 10k-atom KOKKOS path. Docs and the two notebooks that described the old limitation updated. Co-Authored-By: Claude Fable 5 --- NOTEBOOK_TUTORIALS.md | 19 ++-- README.md | 4 + examples/notebook/README.md | 31 +++--- examples/notebook/build.sh | 9 ++ examples/notebook/coi_patch.py | 98 +++++++++++++++++++ .../basics/04-multithreading-kokkos.ipynb | 15 ++- .../content/md-basics/02-condensation.ipynb | 11 ++- 7 files changed, 158 insertions(+), 29 deletions(-) create mode 100644 examples/notebook/coi_patch.py diff --git a/NOTEBOOK_TUTORIALS.md b/NOTEBOOK_TUTORIALS.md index 9f98a7f..b5d2d60 100644 --- a/NOTEBOOK_TUTORIALS.md +++ b/NOTEBOOK_TUTORIALS.md @@ -61,13 +61,18 @@ the kernel is the browser. step callbacks, make all `lammps.*` calls **before** the first `await` — the wasm cannot be re-entered while suspended. The Python notebooks sidestep this by running in chunks (`run 25` in a loop) from the worker. -- **No SharedArrayBuffer on GitHub Pages**: JupyterLite registers its own - service worker, so we don't install `coi-serviceworker` under `/notebook/`. - On the deployed site the KOKKOS multithreaded build is therefore unavailable - and the KOKKOS tutorial (`basics/04`) falls back to the single-threaded build - with an explanation. On a COOP/COEP-enabled host the same site is fully - multithreaded — the Pyodide kernel, piplite and CDN downloads all work under - `require-corp` (verified headlessly). +- **SharedArrayBuffer on GitHub Pages**: static hosting cannot send the + COOP/COEP headers, and JupyterLite's own service worker owns the site + scope (the contents API and our DriveFS mount need it), so the plain + `coi-serviceworker` shim can't be installed under `/notebook/`. Instead, + `build.sh` runs `coi_patch.py`, which folds the same header-injection + logic into JupyterLite's service worker and adds a one-time reload + bootstrap to the app pages. The deployed site is therefore cross-origin + isolated and the KOKKOS multithreaded build works (verified headlessly: + ~2.4× over serial at `t 4`; the Pyodide kernel, piplite and CDN downloads + all work under `require-corp`). Upstream tracking: + jupyterlite/jupyterlite#1409 — replace the patch with the supported flag + if one lands. - **Plotting**: `%pip install matplotlib` just works in the Pyodide kernel (rendered inline as rich output) — this was the main reason to go Python-first. diff --git a/README.md b/README.md index 28383c1..faf5171 100644 --- a/README.md +++ b/README.md @@ -312,6 +312,10 @@ lmp.command("run 100") x = lmp.extract_atom("x") # numpy (natoms, 3) ``` +The site is cross-origin isolated even on GitHub Pages (its service worker +injects the COOP/COEP headers), so the multithreaded KOKKOS build works in +the notebooks too — see the `basics/04` tutorial. + The site source lives in `examples/notebook/` (`content/` holds the notebooks) and deploys with the Pages workflow. See [NOTEBOOK_TUTORIALS.md](NOTEBOOK_TUTORIALS.md) for the tutorial roadmap and diff --git a/examples/notebook/README.md b/examples/notebook/README.md index c342d9a..040cb32 100644 --- a/examples/notebook/README.md +++ b/examples/notebook/README.md @@ -36,15 +36,22 @@ In CI, `.github/workflows/pages.yml` does the same and places the output at The Python bindings accept native-style KOKKOS arguments (`await lammps(cmdargs=["-k", "on", "t", "4", "-sf", "kk"])`), which load the multithreaded `lammps-kokkos.js` build. Threads need `SharedArrayBuffer`, so -the page must be cross-origin isolated: serve the site with - -``` -Cross-Origin-Opener-Policy: same-origin -Cross-Origin-Embedder-Policy: require-corp -``` - -GitHub Pages cannot send these headers and the JupyterLite service worker -occupies the site scope (so the playground's `coi-serviceworker` trick is not -available here) — on the deployed site the KOKKOS notebook falls back to the -single-threaded build. Local serving with the headers above gives real -multithreading; the Pyodide kernel and CDN downloads work fine under COOP/COEP. +the page must be cross-origin isolated (the COOP/COEP headers). + +Static hosting like GitHub Pages cannot send those headers, so `build.sh` +runs `coi_patch.py` after `jupyter lite build`: it folds the +[coi-serviceworker](https://github.com/gzuidhof/coi-serviceworker) trick into +JupyterLite's **own** service worker (only one worker can own the scope, and +the contents API — which our DriveFS mount rides on — needs it to be +JupyterLite's). The worker re-serves every response with the headers added, +and a small bootstrap in the app pages reloads once, the first time the +worker takes control. The result: the site is cross-origin isolated even on +GitHub Pages or a bare `python3 -m http.server`, and the KOKKOS build runs +truly multithreaded (~2.4× over serial at `t 4`, verified headlessly — the +Pyodide kernel, piplite and CDN downloads all work under `require-corp`). +If isolation is unavailable (e.g. the reload guard trips), notebooks fall +back to the single-threaded build gracefully. + +`coi_patch.py` asserts on the service-worker internals of the pinned +`jupyterlite-core` — a version bump that changes them fails the build loudly. +Upstream feature request: jupyterlite/jupyterlite#1409. diff --git a/examples/notebook/build.sh b/examples/notebook/build.sh index eff8c08..c3c1d98 100755 --- a/examples/notebook/build.sh +++ b/examples/notebook/build.sh @@ -39,3 +39,12 @@ assert redirect not in html html = html.replace("", "\n " + redirect, 1) open(path, "w").write(html) EOF + +# Cross-origin isolation on static hosting (the coi-serviceworker trick, +# folded into JupyterLite's own service worker — only one SW can own the +# scope, and the contents API needs it to be JupyterLite's). GitHub Pages +# cannot send COOP/COEP headers, so the service worker adds them to every +# response it serves; app pages get a bootstrap that reloads once when the +# SW takes control, so the *document* is also served with the headers. +# SharedArrayBuffer then works and the KOKKOS multithreaded wasm loads. +python3 coi_patch.py "$OUT" diff --git a/examples/notebook/coi_patch.py b/examples/notebook/coi_patch.py new file mode 100644 index 0000000..f394df3 --- /dev/null +++ b/examples/notebook/coi_patch.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""Make the built JupyterLite site cross-origin isolated on static hosting. + +GitHub Pages cannot send the COOP/COEP headers that SharedArrayBuffer (and +therefore the KOKKOS multithreaded wasm build) requires. The standard +workaround is coi-serviceworker: a service worker that re-serves every +response with the headers added, plus a one-time page reload so the document +itself goes through the service worker. JupyterLite already owns the site's +service-worker scope (its contents API — and our DriveFS mount — depend on +it), so instead of a second worker this script folds the header logic into +JupyterLite's own service worker: + +1. service-worker.js: wrap ``maybeFromCache`` — the function that answers + every plain GET the worker intercepts (documents, assets, and cross-origin + fetches alike) — so its responses carry COOP/COEP, plus CORP so the + assets remain embeddable under the policy they themselves impose. +2. every app page (lab/, notebooks/, tree/, …): inject a bootstrap that + reloads the page once, the first time the service worker takes control. + First visit: plain page registers the worker and reloads into an isolated + page. Later visits are isolated from the start. If isolation still fails + (e.g. an old worker version), the sessionStorage guard stops the loop and + the site simply runs single-threaded, as before. + +The upstream request for this is jupyterlite/jupyterlite#1409; if JupyterLite +grows a supported flag for it, this script can be replaced by that flag. +Pinned to the service-worker internals of jupyterlite-core 0.8.0 — the +asserts below fail the build loudly if a version bump changes them. +""" +import sys +from pathlib import Path + +out = Path(sys.argv[1]) + +# --- 1. teach the service worker to add the isolation headers ------------- + +sw_path = out / "service-worker.js" +sw = sw_path.read_text() + +HOOK = "async function maybeFromCache(e){" +assert sw.count(HOOK) == 1, ( + "jupyterlite service-worker.js changed shape (expected maybeFromCache); " + "update coi_patch.py for this jupyterlite-core version" +) +assert "_lammpsCoiHeaders" not in sw, "service worker already patched" + +sw = sw.replace( + HOOK, + "async function maybeFromCache(e){" + "return _lammpsCoiHeaders(await _lammpsMaybeFromCache(e))}" + "async function _lammpsMaybeFromCache(e){", + 1, +) +sw += ( + "\nfunction _lammpsCoiHeaders(r){" + "if(!r||r.status===0){return r}" + "const h=new Headers(r.headers);" + 'h.set("Cross-Origin-Embedder-Policy","require-corp");' + 'h.set("Cross-Origin-Opener-Policy","same-origin");' + 'h.set("Cross-Origin-Resource-Policy","cross-origin");' + "return new Response(r.body,{status:r.status,statusText:r.statusText,headers:h})" + "}\n" +) +sw_path.write_text(sw) + +# --- 2. reload app pages once, when the service worker takes control ------ + +BOOTSTRAP = """""" + +patched = 0 +for page in sorted(out.glob("*/index.html")): + html = page.read_text() + if "jupyter-config-data" not in html: + continue # not an app page + assert "lammps-coi-reload" not in html, f"{page} already patched" + page.write_text(html.replace("", "\n " + BOOTSTRAP, 1)) + patched += 1 + +assert patched >= 1, "no app pages found to patch" +print(f"coi_patch: service worker patched, {patched} app pages bootstrapped") diff --git a/examples/notebook/content/basics/04-multithreading-kokkos.ipynb b/examples/notebook/content/basics/04-multithreading-kokkos.ipynb index 70b4bdd..469bead 100644 --- a/examples/notebook/content/basics/04-multithreading-kokkos.ipynb +++ b/examples/notebook/content/basics/04-multithreading-kokkos.ipynb @@ -16,8 +16,11 @@ "```\n", "\n", "Threads need `SharedArrayBuffer`, which browsers only allow on\n", - "**cross-origin isolated** pages (COOP/COEP headers). The cell below checks\n", - "and falls back to the single-threaded build when isolation is off:" + "**cross-origin isolated** pages (COOP/COEP headers). This site's service\n", + "worker injects those headers itself — even on static hosting — so the\n", + "check below should say `True` (the very first visit reloads the page once\n", + "to switch it on). The cell still falls back to the single-threaded build\n", + "if isolation is unavailable in your browser:" ] }, { @@ -96,9 +99,11 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Re-run this notebook with isolation on and off (or vary `t 4`) to compare.\n", - "If this deployment is not cross-origin isolated, the same notebooks work\n", - "multithreaded when the site is served with\n", + "Try varying the thread count (`t 4`) and the system size, and compare with\n", + "the single-threaded timing from\n", + "[03 — Analysis and plotting](03-analysis-and-plotting.ipynb)'s runs. If the\n", + "isolation check printed `False`, the same notebook runs multithreaded on\n", + "any host that serves the site with\n", "\n", "```\n", "Cross-Origin-Opener-Policy: same-origin\n", diff --git a/examples/notebook/content/md-basics/02-condensation.ipynb b/examples/notebook/content/md-basics/02-condensation.ipynb index 48f45c0..26ebb42 100644 --- a/examples/notebook/content/md-basics/02-condensation.ipynb +++ b/examples/notebook/content/md-basics/02-condensation.ipynb @@ -150,11 +150,12 @@ "source": [ "## Go bigger (optional)\n", "\n", - "More atoms means more droplets and better coarsening statistics. On a\n", - "cross-origin-isolated host this cell uses the multithreaded KOKKOS build\n", - "(see [basics/04](../basics/04-multithreading-kokkos.ipynb)); on the public\n", - "site it falls back to the single-threaded engine with a smaller box — still\n", - "entirely CPU-bound WebAssembly, so be patient with big boxes:" + "More atoms means more droplets and better coarsening statistics. When the\n", + "page is cross-origin isolated — which this site arranges even on static\n", + "hosting (see [basics/04](../basics/04-multithreading-kokkos.ipynb)) — this\n", + "cell uses the multithreaded KOKKOS build on a 10,000-atom box; otherwise it\n", + "falls back to the single-threaded engine with a smaller one. Still entirely\n", + "CPU-bound WebAssembly, so be patient with big boxes:" ] }, {