fix: install every lazily-imported dependency in the Colab bootstrap; gate setup_colab.py on PRs

.github/workflows/colab-gate.yml

@@ -0,0 +1,88 @@
+# Runs PyAutoHeart's Colab gate (verify_install check F) against THIS branch's
+# setup_colab.py: a venv holding the package set Google actually ships, the
+# injected setup cell's real `--no-deps` bootstrap, then an audit of every
+# import and declared dependency that bootstrap left unmet. So a _SHARED_EXTRAS
+# change is proven against Google's manifest before it merges — the PR-time
+# complement of the release-time verify_install_release job in PyAutoHeart.
+name: Colab gate
+
+on:
+  pull_request:
+    paths:
+      - autonerves/setup_colab.py
+      - .github/workflows/colab-gate.yml
+  workflow_dispatch:
+
+# A push to the PR supersedes the run before it: the gate is a ~15 minute
+# network-bound job and only the tip commit's verdict is wanted.
+concurrency:
+  group: colab-gate-${{ github.ref }}
+  cancel-in-progress: true
+
+permissions:
+  contents: read
+
+jobs:
+  colab-gate:
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    steps:
+      # Unshallow: setuptools-scm reads the tags to version the overlay wheel.
+      - name: Checkout PyAutoNerves
+        uses: actions/checkout@v4
+        with:
+          path: PyAutoNerves
+          fetch-depth: 0
+      # The gate itself lives in Heart, which owns all readiness checking.
+      # Pinned to main: a PR here is gated by the gate as it currently ships.
+      - name: Checkout PyAutoHeart (the gate)
+        uses: actions/checkout@v4
+        with:
+          repository: PyAutoLabs/PyAutoHeart
+          ref: main
+          path: PyAutoHeart
+      # Colab runs Python 3.12 and check F invokes the `python3.12` binary by
+      # name — a different interpreter would seed Colab's pins against the
+      # wrong wheels, so the check FAILs rather than SKIPs when it is absent.
+      - name: Set up Python 3.12 (the interpreter Colab runs)
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Confirm python3.12 resolves by name
+        run: python3.12 --version
+      # COLAB_GATE_AUTONERVES_SRC overlays this checkout's setup_colab.py on
+      # top of the released `autonerves` the setup cell bootstraps. Heart's own
+      # release run must NEVER set it (a release gate reads the wheels about to
+      # ship); here the unreleased file IS the thing under test.
+      - name: Run the Colab gate (check F) against this branch's setup_colab
+        env:
+          COLAB_GATE_AUTONERVES_SRC: ${{ github.workspace }}/PyAutoNerves
+          HEART_STATE_DIR: ${{ runner.temp }}/heart-state
+        run: |
+          bash PyAutoHeart/heart/checks/verify_install.sh F \
+            --report-json "$RUNNER_TEMP/verify_install.json"
+      - name: Upload the gate report
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: colab-gate-report
+          path: ${{ runner.temp }}/verify_install.json
+          retention-days: 30
+      # verify_install already exits non-zero on a failed check, so the step
+      # above is the gate. This re-reads the verdict the sidecar recorded, so a
+      # report that says `ready: false` can never ride through on a zero exit.
+      - name: Fail on a non-ready report
+        if: always()
+        run: |
+          python3 - "$RUNNER_TEMP/verify_install.json" <<'PY'
+          import json, sys
+
+          with open(sys.argv[1]) as handle:
+              report = json.load(handle)
+
+          for check in report["checks"]:
+              print(f"{check['check']}|{check['status']}|{check['detail']}")
+
+          if not report["ready"]:
+              sys.exit("Colab gate not ready — see the RESULTS rows above")
+          PY

autonerves/setup_colab.py

@@ -64,6 +64,14 @@
     # package rather than adding a missing one.
     "xxhash<=3.4.1",
     "blackjax>=1.6.2",
+    # `jax_zero_contour` is an autogalaxy BASE dependency, reached by the
+    # critical-curve code in `autogalaxy/operate/lens_calc.py`. Declared behind
+    # the same marker as `optax` above, dropped here for the same reason.
+    "jax_zero_contour>=2.0.0,<3.0.0",
+    # `zeus` is imported inside the `af.Zeus` search, which the autofit
+    # workspace notebooks run (cookbooks/search, searches/mcmc, the zeus
+    # plotter); the distribution is `zeus-mcmc`, autofit's `optional` extra.
+    "zeus-mcmc==2.5.4",
 ]
 
 _AUTOFIT_STACK = ["autonerves", "autofit"]

test_autonerves/test_setup_colab.py

@@ -130,10 +130,12 @@ def test_every_project_installs_every_sampler(self):
         # ModuleNotFoundError (HowToFit chapter 1 tutorials 4, 5 and 6 on
         # Colab). Widened past the samplers to every autofit dependency that is
         # imported lazily, inside a function, and so survives `import autofit`:
-        # `corner` (the reported failure), `optax`, `xxhash` and `blackjax`.
-        # Match on the package name only, so a future re-pin of any of them
-        # does not break this test — drift is `TestSpecifiersTrackAutofit`'s
-        # job, this one guards the "missing entirely" class.
+        # `corner` (the reported failure), `optax`, `xxhash`, `blackjax`,
+        # `jax_zero_contour` (autogalaxy's critical-curve code) and
+        # `zeus-mcmc` (`af.Zeus`). Match on the package name only, so a future
+        # re-pin of any of them does not break this test — drift is
+        # `TestSpecifiersTrackAutofit`'s job, this one guards the "missing
+        # entirely" class.
         required = {
             "dynesty",
             "emcee",
@@ -142,6 +144,8 @@ def test_every_project_installs_every_sampler(self):
             "optax",
             "xxhash",
             "blackjax",
+            "jax_zero_contour",
+            "zeus-mcmc",
         }
         for project, spec in setup_colab._PROJECTS.items():
             names = {
@@ -186,8 +190,8 @@ def test_shared_extras_match_autofits_declared_specifiers(self):
         Do not "simplify" this back to plain string equality on both arms: that
         is what `dill` fails, and it fails for no good reason.
 
-        The expectations are DERIVED from PyAutoFit's pyproject.toml at run
-        time, never restated here. Two entries had already drifted from the
+        The expectations are DERIVED from PyAutoFit's and PyAutoGalaxy's
+        pyproject.toml at run time, never restated here. Two entries had already drifted from the
         file the list's own comment claims to track (`nautilus-sampler` a patch
         behind autofit's pin, `anesthetic` pinned BELOW autofit's floor)
         precisely because the list repeats literals nobody re-checks. Copying
@@ -199,26 +203,46 @@ def test_shared_extras_match_autofits_declared_specifiers(self):
         except ImportError:  # pragma: no cover - `packaging` ships with pip
             pytest.skip("`packaging` is not importable, so specifiers cannot be compared")
 
-        pyproject = Path(__file__).parents[2] / "PyAutoFit" / "pyproject.toml"
+        siblings = Path(__file__).parents[2]
+        pyprojects = {
+            name: siblings / name / "pyproject.toml"
+            for name in ("PyAutoFit", "PyAutoGalaxy")
+        }
+
+        missing = sorted(
+            str(path) for path in pyprojects.values() if not path.is_file()
+        )
 
-        if not pyproject.is_file():
-            # PyAutoNerves CI may run with no sibling PyAutoFit checkout; there
+        if missing:
+            # PyAutoNerves CI may run with no sibling library checkouts; there
             # is nothing to compare against, and that is not a failure.
             pytest.skip(
-                f"no sibling PyAutoFit checkout at {pyproject} to read "
+                f"no sibling checkout at {', '.join(missing)} to read "
                 "declared specifiers from"
             )
 
-        with open(pyproject, "rb") as f:
-            project = tomllib.load(f)["project"]
-
-        # `blackjax` and `nautilus-sampler` are declared in the `optional`
-        # extra rather than the base dependencies.
+        projects = {}
+        for name, path in pyprojects.items():
+            with open(path, "rb") as f:
+                projects[name] = tomllib.load(f)["project"]
+
+        def requirements(project):
+            # `blackjax`, `nautilus-sampler` and `zeus-mcmc` are declared in
+            # the `optional` extra rather than the base dependencies.
+            return project["dependencies"] + project["optional-dependencies"][
+                "optional"
+            ]
+
+        # PyAutoGalaxy first, then PyAutoFit over the top: where both declare a
+        # name, autofit's specifier is the one this test has always compared
+        # against, and it stays the authority. PyAutoGalaxy is read because
+        # `jax_zero_contour` is a base dependency of THAT file, not autofit's,
+        # and the Colab list has to track it just as closely.
         declared = dict(
             _split_requirement(requirement)
             for requirement in (
-                project["dependencies"]
-                + project["optional-dependencies"]["optional"]
+                requirements(projects["PyAutoGalaxy"])
+                + requirements(projects["PyAutoFit"])
             )
         )