Make the notebook site cross-origin isolated on GitHub Pages (KOKKOS works deployed)

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.

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

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.

examples/notebook/build.sh

@@ -39,3 +39,12 @@ assert redirect not in html
 html = html.replace("<head>", "<head>\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"

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 = """<script>
+    (() => {
+      if (window.crossOriginIsolated || !('serviceWorker' in navigator)) return;
+      const KEY = 'lammps-coi-reload';
+      if (sessionStorage.getItem(KEY)) return; // never loop
+      const reload = () => {
+        sessionStorage.setItem(KEY, '1');
+        window.location.reload();
+      };
+      if (navigator.serviceWorker.controller) {
+        // Controlled but not isolated: an older worker without the header
+        // patch is running; one reload picks up the updated worker.
+        navigator.serviceWorker.ready.then(reload);
+      } else {
+        // First visit: JupyterLite registers the worker during boot; as
+        // soon as it claims this page, reload into an isolated document.
+        navigator.serviceWorker.addEventListener('controllerchange', reload);
+      }
+    })();
+    </script>"""
+
+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("<head>", "<head>\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")

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",

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:"
    ]
   },
   {