Re-lock every notebook for Colab's Python 3.13 runtime

.github/docker/Dockerfile

@@ -13,7 +13,7 @@
 #
 # The server runs as the non-root user jovyan (uid 1000).
 
-ARG BASE_IMAGE=python:3.12-slim@sha256:2c941e860699f878900b0edc2403613c234d4b32eda3cc9fa7036991a2a63c4a
+ARG BASE_IMAGE=python:3.13-slim@sha256:ffb752e139c0a19692a43af8d8523b274222dd68eebad5d583b45c2201c6e30a
 FROM ${BASE_IMAGE}
 
 # libgl1/libglib2.0-0/libxcb1 are runtime needs of opencv-python and similar

.github/notebook-colab-exclusions.txt

@@ -90,3 +90,14 @@ tutorials/cosyne_2023/advanced_asset_search.ipynb
 
 # Validates every zarr chunk in the dandiset; >1 h in Colab too.
 000108/chunglab/demo/validate_lev6.ipynb
+
+# =============================================================================
+# Data not yet on DANDI.
+# =============================================================================
+
+# The only session published in Dandiset 001712 (sub-FD-28) has no
+# lab_meta_data["atlas_registration"], which these two notebooks need for
+# their atlas-overlay sections. The notebooks are ahead of the data; remove
+# once the registered files are uploaded.
+001712/IBL-Widefield/public_demo/processed_widefield.ipynb
+001712/IBL-Widefield/public_demo/anatomical_localization_widefield.ipynb

.github/notebook-test-exclusions.txt

@@ -93,6 +93,13 @@ tutorials/cosyne_2023/advanced_asset_search.ipynb
 000055/BruntonLab/peterson21/Table_coarse_labels.ipynb
 000055/BruntonLab/peterson21/Table_part_characteristics.ipynb
 
+# The only session published in Dandiset 001712 (sub-FD-28) has no
+# lab_meta_data["atlas_registration"], which these two notebooks need for
+# their atlas-overlay sections. The notebooks are ahead of the data; remove
+# once the registered files are uploaded.
+001712/IBL-Widefield/public_demo/processed_widefield.ipynb
+001712/IBL-Widefield/public_demo/anatomical_localization_widefield.ipynb
+
 # =============================================================================
 # Long-running notebooks — exceed the CI per-notebook timeout.
 # =============================================================================

.github/scripts/backfill_requirements.py

@@ -0,0 +1,196 @@
+"""Derive requirements.in files for notebook groups that predate the contract.
+
+For each image group without a requirements file, collect the top-level
+modules its notebooks import, then resolve those to distribution names using
+`importlib.metadata.packages_distributions()` inside the group's published
+container image (the exact environment the notebook was verified in, so the
+mapping is not a guess). Standard-library modules and the group's own helper
+files are skipped; `pkg @ git+...` pins from the install cell are direct
+dependencies by definition and are carried over verbatim.
+
+Known incompatibilities are encoded as bounds: nwbwidgets only imports on the
+pre-2024 stack, so groups that pin it get pynwb<3, hdmf<4, zarr<3.
+
+Usage:
+    python .github/scripts/backfill_requirements.py [--dry-run] [--filter REGEX]
+
+Requires docker and nbformat.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+import nbformat
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from build_notebook_image import DEFAULT_IMAGE_PREFIX, collect_groups, matches  # noqa: E402
+from lock_notebook import requirements_for  # noqa: E402
+
+IMPORT_RE = re.compile(r"^\s*(?:from\s+([A-Za-z_][\w]*)|import\s+([A-Za-z_][\w]*))", re.M)
+NWBWIDGETS_BOUNDS = ["pynwb<3", "hdmf<4", "zarr<3"]
+# Import names whose distribution is not discoverable by name matching.
+IMPORT_ALIASES = {"skimage": "scikit-image", "cv2": "opencv-python", "PIL": "pillow",
+                  "yaml": "pyyaml", "sklearn": "scikit-learn"}
+
+# Runs inside the image: maps import names to distributions, and scans the
+# helper modules baked into /work (local .py files, packages, and files the
+# install cell fetched with curl) for their own imports, since a notebook's
+# direct dependencies include whatever its helpers import.
+PROBE = r"""
+import importlib.metadata as m, json, re, sys
+from pathlib import Path
+IMPORT_RE = re.compile(r"^\s*(?:from\s+([A-Za-z_]\w*)|import\s+([A-Za-z_]\w*))", re.M)
+work = Path("/work")
+helper_modules, helper_imports = set(), set()
+for py in work.rglob("*.py"):
+    rel = py.relative_to(work)
+    helper_modules.add(rel.parts[0].removesuffix(".py"))
+    try:
+        for mm in IMPORT_RE.finditer(py.read_text(errors="ignore")):
+            helper_imports.add(mm.group(1) or mm.group(2))
+    except OSError:
+        pass
+print(json.dumps({
+    "dists": m.packages_distributions(),
+    "stdlib": sorted(sys.stdlib_module_names),
+    "helper_modules": sorted(helper_modules),
+    "helper_imports": sorted(helper_imports),
+}))
+"""
+
+
+def notebook_imports(path: Path) -> tuple[set[str], bool]:
+    """Top-level imports, plus whether the notebook drives the dandi CLI from shell cells."""
+    nb = nbformat.read(path, as_version=4)
+    mods: set[str] = set()
+    uses_dandi_cli = False
+    for cell in nb.cells:
+        if cell.cell_type != "code":
+            continue
+        for line in cell.source.splitlines():
+            stripped = line.lstrip()
+            if stripped.startswith("!dandi ") or stripped.startswith("!dandi\t"):
+                uses_dandi_cli = True
+            if stripped.startswith(("!", "%")):
+                continue
+            for m in IMPORT_RE.finditer(line):
+                mods.add(m.group(1) or m.group(2))
+    return mods, uses_dandi_cli
+
+
+def probe_image(image: str) -> dict:
+    subprocess.run(["docker", "pull", "-q", image], check=True, capture_output=True)
+    r = subprocess.run(["docker", "run", "--rm", "-i", image, "python", "-"],
+                       input=PROBE, capture_output=True, text=True, check=True)
+    return json.loads(r.stdout.strip().splitlines()[-1])
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--dry-run", action="store_true")
+    parser.add_argument("--filter", default="")
+    parser.add_argument("--force", action="store_true", help="rewrite existing requirements files too")
+    args = parser.parse_args()
+
+    repo = Path(__file__).resolve().parents[2]
+    for group in collect_groups():
+        if not matches(group, args.filter):
+            continue
+        nb_paths = [repo / group.directory / n for n in group.notebooks]
+        try:
+            requirements_for(nb_paths[0])
+            if not args.force:
+                continue  # already has a requirements file
+        except FileNotFoundError:
+            pass
+
+        imports: set[str] = set()
+        uses_dandi_cli = False
+        for p in nb_paths:
+            mods, cli = notebook_imports(p)
+            imports |= mods
+            uses_dandi_cli = uses_dandi_cli or cli
+
+        # Probe this group's image, or a sibling group's in the same directory
+        # when this group has never published (same helpers, near-identical env).
+        siblings = [g for g in collect_groups() if g.directory == group.directory]
+        info = None
+        for candidate in [group] + [g for g in siblings if g.name != group.name]:
+            try:
+                info = probe_image(f"{DEFAULT_IMAGE_PREFIX}/{candidate.name}:latest")
+                break
+            except subprocess.CalledProcessError:
+                continue
+        if info is None:
+            print(f"[{group.name}] no image to probe in {group.directory}; skipping", file=sys.stderr)
+            continue
+        dists, stdlib = info["dists"], set(info["stdlib"])
+        local_modules = set(info["helper_modules"])
+        imports |= set(info["helper_imports"])
+        if uses_dandi_cli:
+            direct_extra = {"dandi"}
+        else:
+            direct_extra = set()
+
+        direct: set[str] = set(direct_extra)
+        unmapped: list[str] = []
+        for mod in sorted(imports):
+            if mod in stdlib or mod in local_modules:
+                continue
+            names = dists.get(mod)
+            if names:
+                direct.update(names)
+                continue
+            # Not installed in the probed image (a sibling's): fall back to the
+            # group's own pins, by alias or by name with _ and - interchangeable.
+            wanted = IMPORT_ALIASES.get(mod, mod).lower().replace("_", "-")
+            pinned = [pin.split("==")[0] for pin in group.pins
+                      if pin.split("==")[0].lower().replace("_", "-") == wanted]
+            if pinned:
+                direct.add(pinned[0])
+            else:
+                unmapped.append(mod)
+
+        # NWB extension packages are used through a file's cached schema and
+        # never imported, so the import scan cannot see them; they are only
+        # ever installed deliberately, so carry them over from the old lock.
+        for pin in group.pins:
+            name = pin.split("==")[0].split(" @ ")[0].strip()
+            if name.lower().startswith("ndx-") and " @ " not in pin:
+                direct.add(name)
+
+        git_pins = [p for p in group.pins if " @ " in p]
+        for gp in git_pins:
+            name = gp.split(" @ ")[0].strip()
+            direct.discard(name)
+        lines = sorted(direct, key=str.lower) + git_pins
+        if any(p.startswith("nwbwidgets==") for p in group.pins):
+            lines += NWBWIDGETS_BOUNDS
+        elif "pynwb" in direct or "hdmf" in direct:
+            # Hold the NWB stack at the majors the fleet was last verified on;
+            # a Colab-driven re-lock should not double as a pynwb major upgrade.
+            lines += ["pynwb<4", "hdmf<5"]
+
+        # One requirements.in per directory when the directory holds a single
+        # group; otherwise a per-notebook file, keyed on the first notebook.
+        if len(siblings) == 1:
+            target = repo / group.directory / "requirements.in"
+        else:
+            target = nb_paths[0].with_name(f"{nb_paths[0].stem}.requirements.in")
+
+        rel = target.relative_to(repo)
+        print(f"[{group.name}] -> {rel}: {', '.join(lines)}"
+              + (f"  (UNMAPPED imports: {', '.join(unmapped)})" if unmapped else ""))
+        if not args.dry_run:
+            target.write_text("\n".join(lines) + "\n")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

.github/scripts/lock_notebook.py

@@ -23,7 +23,9 @@
 from __future__ import annotations
 
 import argparse
+import re
 import subprocess
+import tempfile
 import sys
 from pathlib import Path
 
@@ -56,7 +58,7 @@
 )
 INSTALL_HEADER = (
     '#@title Installing requirements (click ▶ to run) { display-mode: "form" }\n'
-    "# Colab provides Python 3.12. We install with `uv --system` because Colab's\n"
+    "# Colab provides Python 3.13. We install with `uv --system` because Colab's\n"
     "# kernel runs outside a virtualenv. All versions (direct + transitive) are\n"
     "# pinned below so the notebook is reproducible regardless of resolver drift.\n"
     "!pip install -q uv\n"
@@ -77,15 +79,52 @@ def requirements_for(nb_path: Path) -> Path:
     )
 
 
+OVERRIDE_PREFIX = "# override:"
+
+
+def overrides_in(requirements: Path) -> list[str]:
+    """Constraint overrides declared as `# override: <spec>` lines.
+
+    Colab's preinstalled versions are applied as constraints so installs stay
+    fast and compatible with the runtime. Occasionally a notebook's verified
+    stack needs an older version of one of those packages (e.g. the dandi
+    release that still reads a file needs click<8.2); an override line replaces
+    the Colab pin for that package alone, at the cost of Colab downgrading
+    it in the install cell.
+    """
+    return [
+        line[len(OVERRIDE_PREFIX):].strip()
+        for line in requirements.read_text().splitlines()
+        if line.strip().startswith(OVERRIDE_PREFIX) and line[len(OVERRIDE_PREFIX):].strip()
+    ]
+
+
 def compile_pins(requirements: Path) -> list[str]:
+    overrides = overrides_in(requirements)
+    constraint = CONSTRAINT
+    if overrides:
+        # Replace the Colab pin for each overridden package with the override.
+        overridden = {re.split(r"[<>=!~\s\[]", spec, 1)[0].lower().replace("_", "-")
+                      for spec in overrides}
+        kept = [line for line in CONSTRAINT.read_text().splitlines()
+                if line.strip() and not line.startswith("#")
+                and re.split(r"[<>=!~\s\[]", line.strip(), 1)[0].lower().replace("_", "-")
+                not in overridden]
+        with tempfile.NamedTemporaryFile("w", suffix=".constraints.txt", delete=False) as f:
+            f.write("\n".join(kept + overrides) + "\n")
+            constraint = Path(f.name)
     cmd = [
         "uv", "pip", "compile", str(requirements),
-        "--python-version", "3.12",
+        "--python-version", "3.13",
         "--python-platform", "linux",
-        "--constraint", str(CONSTRAINT),
+        "--constraint", str(constraint),
         "--no-header", "--no-annotate",
     ]
-    r = subprocess.run(cmd, capture_output=True, text=True, stdin=subprocess.DEVNULL)
+    try:
+        r = subprocess.run(cmd, capture_output=True, text=True, stdin=subprocess.DEVNULL)
+    finally:
+        if constraint is not CONSTRAINT:
+            constraint.unlink(missing_ok=True)
     if r.returncode != 0:
         raise RuntimeError(f"uv pip compile failed for {requirements}:\n{r.stderr}")
     pins = [