From ded961bd2bcb7b6816c02cda2d35f96e35f9f4b6 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 4 Nov 2026 22:04:18 +0000 Subject: [PATCH] setup_colab: register autocti and install arcticpy via Heart's script (PyAutoHands#332) Adds autocti to the Colab _PROJECTS registry. Its pre-install step fetches PyAutoHeart's install_arcticpy.sh and runs it with ARCTICPY_SUDO=false and PYTHON=sys.executable, so the arcticpy pin keeps a single home in Heart. Co-Authored-By: Claude Opus 5.5 --- autonerves/setup_colab.py | 48 ++++++++++++++++++++++++++++- test_autonerves/test_setup_colab.py | 47 ++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+), 1 deletion(-) diff --git a/autonerves/setup_colab.py b/autonerves/setup_colab.py index b09374d..d2adc54 100644 --- a/autonerves/setup_colab.py +++ b/autonerves/setup_colab.py @@ -80,6 +80,18 @@ # matches AutoArray's Linux requirement. _AUTOGALAXY_STACK = _AUTOFIT_STACK + ["autoarray", "autogalaxy", "finufft>=2.4.0"] _AUTOLENS_STACK = _AUTOGALAXY_STACK + ["autolens"] +_AUTOCTI_STACK = _AUTOFIT_STACK + ["autoarray", "autocti", "finufft>=2.4.0"] + +# arcticpy (the C++ arctic clocking code autocti imports) is a source-only sdist +# that needs libgsl-dev and a no-build-isolation build, so it cannot ride the +# `--no-deps` package list. PyAutoHeart's `install_arcticpy.sh` is the single +# home of that recipe and of the arcticpy version pin (cti-upkeep ruling, +# 2026-11-04), so Colab runs that script rather than carrying a second copy. +# Colab runs as root (no sudo) and must build into the notebook's interpreter. +_ARCTICPY_INSTALLER_URL = ( + "https://raw.githubusercontent.com/PyAutoLabs/PyAutoHeart/main/" + ".github/actions/install-arcticpy/install_arcticpy.sh" +) # One entry per Colab-enabled notebook repository. ``top_package`` is the # package whose installed version names the workspace tag to clone. @@ -108,6 +120,15 @@ workspace_dir="/content/autolens_workspace", gpu_note="> 50 times faster", ), + "autocti": dict( + project_name="PyAutoCTI", + top_package="autocti", + packages=_AUTOCTI_STACK + _SHARED_EXTRAS, + workspace_repo="https://github.com/PyAutoLabs/autocti_workspace", + workspace_dir="/content/autocti_workspace", + gpu_note="no faster (arctic clocking runs on the CPU)", + pre_install="arcticpy", + ), "howtofit": dict( project_name="PyAutoFit", top_package="autofit", @@ -255,6 +276,21 @@ def _gpu_error_message(project: str, spec: dict) -> str: """ +def _install_arcticpy() -> None: + """Build arcticpy on Colab by running PyAutoHeart's installer script.""" + import os + import subprocess + import sys + + env = dict(os.environ, ARCTICPY_SUDO="false", PYTHON=sys.executable) + subprocess.check_call( + ["bash", "-c", f'curl -fsSL "{_ARCTICPY_INSTALLER_URL}" | bash'], env=env + ) + + +_PRE_INSTALLERS = {"arcticpy": _install_arcticpy} + + def _colab_setup( *, project_name: str, @@ -264,6 +300,7 @@ def _colab_setup( top_package: str, raise_error_if_not_gpu: bool, gpu_error_message: str, + pre_install: str = None, ) -> None: """ Shared Google Colab setup for the PyAuto ecosystem. @@ -308,6 +345,9 @@ def _colab_setup( print() print(f"Now Installing {project_name} and setting up Colab Environment:") + if pre_install is not None: + _PRE_INSTALLERS[pre_install]() + subprocess.check_call( [sys.executable, "-m", "pip", "install", *packages, "--no-deps"] ) @@ -347,7 +387,7 @@ def setup( Parameters ---------- project - One of: "autofit", "autogalaxy", "autolens", "howtofit", + One of: "autofit", "autogalaxy", "autolens", "autocti", "howtofit", "howtogalaxy", "howtolens". raise_error_if_not_gpu If True, raise instead of continuing when Colab has no GPU / TPU @@ -373,6 +413,7 @@ def setup( top_package=spec["top_package"], raise_error_if_not_gpu=raise_error_if_not_gpu, gpu_error_message=_gpu_error_message(project, spec), + pre_install=spec.get("pre_install"), ) @@ -391,6 +432,11 @@ def for_autolens(raise_error_if_not_gpu: bool = True) -> None: setup("autolens", raise_error_if_not_gpu=raise_error_if_not_gpu) +def for_autocti(raise_error_if_not_gpu: bool = False) -> None: + """Google Colab helper for PyAutoCTI (the autocti_workspace).""" + setup("autocti", raise_error_if_not_gpu=raise_error_if_not_gpu) + + def for_howtofit(raise_error_if_not_gpu: bool = False) -> None: """Google Colab helper for the HowToFit lectures.""" setup("howtofit", raise_error_if_not_gpu=raise_error_if_not_gpu) diff --git a/test_autonerves/test_setup_colab.py b/test_autonerves/test_setup_colab.py index d35344c..0da229c 100644 --- a/test_autonerves/test_setup_colab.py +++ b/test_autonerves/test_setup_colab.py @@ -138,6 +138,48 @@ def test_colab_installs_cpu_nufft_for_array_stacks(self, project, monkeypatch, t else: assert "finufft" not in requirements, project + @pytest.mark.parametrize("project", tuple(setup_colab._PROJECTS)) + def test_only_autocti_builds_arcticpy_and_before_the_stack( + self, project, monkeypatch, tmp_path + ): + """arcticpy is a source-only C++ build that cannot ride the --no-deps list. + + autocti runs PyAutoHeart's installer script (the single home of the + arcticpy pin) as root into this interpreter, BEFORE the package stack, + because `import autocti` needs it. No other project touches it. + """ + from autonerves import conf + + google = types.ModuleType("google") + google.__path__ = [] + google.colab = types.ModuleType("google.colab") + monkeypatch.setitem(sys.modules, "google", google) + monkeypatch.setitem(sys.modules, "google.colab", google.colab) + monkeypatch.setenv("XLA_FLAGS", "") + monkeypatch.chdir(tmp_path) + with mock.patch.object(subprocess, "check_call") as calls, \ + mock.patch.object(setup_colab, "check_jax_using_gpu", return_value=False), \ + mock.patch.object(setup_colab, "_clone_workspace"), \ + mock.patch.object(conf.instance, "push"): + setup_colab.setup(project, workspace_dir=str(tmp_path)) + + commands = [c.args[0] for c in calls.call_args_list] + builds = [c for c in calls.call_args_list if c.args[0][0] == "bash"] + if project != "autocti": + assert builds == [], project + return + assert len(builds) == 1 + script = builds[0].args[0][2] + assert setup_colab._ARCTICPY_INSTALLER_URL in script + assert "PyAutoHeart/main/.github/actions/install-arcticpy/" in script + assert "arcticpy==" not in script # the pin lives in Heart, never here + env = builds[0].kwargs["env"] + assert env["ARCTICPY_SUDO"] == "false" + assert env["PYTHON"] == sys.executable + assert commands.index(builds[0].args[0]) < len(commands) - 1 + assert commands[-1][:4] == [sys.executable, "-m", "pip", "install"] + assert "autocti" in commands[-1] + def test_all_projects_have_required_fields(self): required = { "project_name", @@ -368,6 +410,11 @@ def test_wrappers_delegate(self): "autolens", raise_error_if_not_gpu=False ) setup_mock.reset_mock() + setup_colab.for_autocti() + setup_mock.assert_called_once_with( + "autocti", raise_error_if_not_gpu=False + ) + setup_mock.reset_mock() setup_colab.for_howtofit() setup_mock.assert_called_once_with( "howtofit", raise_error_if_not_gpu=False