diff --git a/autoconf/setup_colab.py b/autoconf/setup_colab.py index 6d1e727..5cb8c2a 100644 --- a/autoconf/setup_colab.py +++ b/autoconf/setup_colab.py @@ -1,213 +1,327 @@ -import logging -import os -from pathlib import Path -os.environ['XLA_FLAGS'] = "--xla_disable_hlo_passes=constant_folding" - -logger = logging.getLogger(__name__) - -def check_jax_using_gpu(log_gpu_warning : bool = True) -> bool: - - import jax - - devices = jax.devices() - - # Normalize device_kind for matching - for d in devices: - kind = str(d).lower() - - no_gpu = False - - if "gpu" in kind or "cuda" in kind: - print(f"GPU / CUDA detected: {d}") - elif "tpu" in kind: - print(f"TPU detected: {d}") - elif "cpu" in kind: - print(f"CPU detected: {d}") - no_gpu = True - else: - print(f"Other device detected: {d}") - no_gpu = True - - if no_gpu and log_gpu_warning: - - logger.info( - """ - JAX did not detect a GPU or TPU device and is using the CPU for computations. - - PyAutoLens runs > 50 times faster with a GPU, so it is recommended to reinstall - JAX With GPU support, if you have a GPU available, by following the JAX - installation instructions. - """) - - return no_gpu - -def _colab_setup( - *, - project_name: str, - workspace_repo: str, - workspace_dir: str, - packages: list[str], - raise_error_if_not_gpu: bool, - gpu_error_message: str, -) -> None: - """ - Shared Google Colab setup for the PyAuto* ecosystem. - - Assumes a `check_jax_using_gpu(log_gpu_warning: bool)` utility exists in scope. - """ - import os - import subprocess - import sys - - # Identical JAX performance tweak. - os.environ["XLA_FLAGS"] = "--xla_disable_hlo_passes=constant_folding" - - # --- Ensure we're in Colab and (optionally) have a GPU --- - try: - import google.colab # noqa: F401 - - no_gpu = check_jax_using_gpu(log_gpu_warning=False) - - if no_gpu and raise_error_if_not_gpu: - raise RuntimeError(gpu_error_message) - - except ImportError: - print( - f""" - You are not running in a Google Colab environment so cannot use the setup_colab() function. - - You should therefore have {project_name} installed locally in your environment already (e.g. via pip or - conda) and can run the rest of your script normally. - - You may now continue running your script or Notebook. - """ - ) - return - - # --- Install packages (no deps, consistent with your current approach) --- - print() - print(f"Now Installing {project_name} and setting up Colab Environment:") - - subprocess.check_call( - [sys.executable, "-m", "pip", "install", *packages, "--no-deps"] - ) - - # --- Clone workspace --- - try: - subprocess.run(["git", "clone", workspace_repo, workspace_dir], check=True) - except subprocess.CalledProcessError: - print("Workspace already exists so not cloning again.") - - # --- Configure autoconf paths --- - from autoconf import conf - - os.chdir(workspace_dir) - - conf.instance.push( - new_path=Path(workspace_dir) / "config", - output_path=Path(workspace_dir) / "output", - ) - - print( - f""" - ***Google Colab Setup Complete, which included:*** - - - Installation of {project_name} and other required packages. - - Cloning of the workspace GitHub repository. - - Setting up environment variables for JAX for improved performance. - - Setting the configuration paths to the workspace config and output folders suitable for Colab. - """ - ) - - -def for_autogalaxy(raise_error_if_not_gpu: bool = True) -> None: - """ - Google Colab helper for PyAutoGalaxy. - - Parallel to `for_autolens`, but installs PyAutoGalaxy only (no `autolens`) and clones the - `autogalaxy_workspace`. - """ - packages = [ - # Core stack - "autoconf", - "autofit", - "autoarray", - "autogalaxy", - # Notebook/runtime extras commonly used across workspaces - "pyvis==0.3.2", - "dill==0.4.0", - "jaxnnls", - "nautilus-sampler==1.0.4", - "timeout_decorator==0.5.0", - "anesthetic==2.8.14", - ] - - gpu_error_message = """ - No GPU detected in Google Colab. PyAutoGalaxy runs much faster with a GPU, so switch GPU - on in Colab settings. - - To do this: - - - Click "Runtime" in the top menu. - - Click "Change runtime type". - - Under "Hardware accelerator", select one of the "GPU" options available. - - You can set up Colab with a CPU (e.g. if GPUs are unavailable) by calling: - - setup_colab.for_autogalaxy(raise_error_if_not_gpu=False) - """ - - _colab_setup( - project_name="PyAutoGalaxy", - workspace_repo="https://github.com/PyAutoLabs/autogalaxy_workspace", - workspace_dir="/content/autogalaxy_workspace", - packages=packages, - raise_error_if_not_gpu=raise_error_if_not_gpu, - gpu_error_message=gpu_error_message, - ) - -def for_autolens(raise_error_if_not_gpu: bool = True) -> None: - """ - Google Colab helper for PyAutoLens. - - Installs the full PyAutoLens stack and clones the `autolens_workspace`. - """ - packages = [ - # Core stack - "autoconf", - "autofit", - "autoarray", - "autogalaxy", - "autolens", - # Notebook/runtime extras commonly used across workspaces - "pyvis==0.3.2", - "dill==0.4.0", - "jaxnnls", - "nautilus-sampler==1.0.4", - "timeout_decorator==0.5.0", - "anesthetic==2.8.14", - ] - - gpu_error_message = """ - No GPU detected in Google Colab. PyAutoLens runs > 50 times faster with a GPU, so switch GPU - on in Colab settings. - - To do this: - - - Click "Runtime" in the top menu. - - Click "Change runtime type". - - Under "Hardware accelerator", select one of the "GPU" options available. - - You can set up Colab with a CPU (e.g. if GPUs are unavailable) by calling: - - setup_colab.for_autolens(raise_error_if_not_gpu=False) - """ - - _colab_setup( - project_name="PyAutoLens", - workspace_repo="https://github.com/PyAutoLabs/autolens_workspace", - workspace_dir="/content/autolens_workspace", - packages=packages, - raise_error_if_not_gpu=raise_error_if_not_gpu, - gpu_error_message=gpu_error_message, - ) +""" +Google Colab bootstrap for the PyAuto ecosystem. + +Every notebook generated by PyAutoBuild begins with a setup cell that calls +``setup_colab.setup("")``. Outside Colab the call is a no-op (it +prints a short confirmation and returns), so the same notebook runs unchanged +locally and on Colab. + +On Colab, ``setup`` installs the project's package stack, clones the matching +workspace repository (at the tag matching the installed release where one +exists) and points the autoconf config / output paths at it. + +The per-project package lists, workspace repositories and Colab directories +live in the ``_PROJECTS`` registry below — one table, shared by the public +``setup`` entry point and the ``for_`` convenience wrappers. +""" + +import logging +from pathlib import Path + +logger = logging.getLogger(__name__) + +# JAX performance tweak applied on Colab (see the PyAuto JAX docs); applied +# inside the setup functions only, never at import time, so importing this +# module does not mutate the environment of local runs. +_XLA_FLAGS = "--xla_disable_hlo_passes=constant_folding" + +# Notebook / runtime extras installed alongside every project's core stack. +_SHARED_EXTRAS = [ + "pyvis==0.3.2", + "dill==0.4.0", + "jaxnnls", + "nautilus-sampler==1.0.4", + "timeout_decorator==0.5.0", + "anesthetic==2.8.14", +] + +_AUTOFIT_STACK = ["autoconf", "autofit"] +_AUTOGALAXY_STACK = _AUTOFIT_STACK + ["autoarray", "autogalaxy"] +_AUTOLENS_STACK = _AUTOGALAXY_STACK + ["autolens"] + +# One entry per Colab-enabled notebook repository. ``top_package`` is the +# package whose installed version names the workspace tag to clone. +_PROJECTS = { + "autofit": dict( + project_name="PyAutoFit", + top_package="autofit", + packages=_AUTOFIT_STACK + _SHARED_EXTRAS, + workspace_repo="https://github.com/PyAutoLabs/autofit_workspace", + workspace_dir="/content/autofit_workspace", + gpu_note="faster", + ), + "autogalaxy": dict( + project_name="PyAutoGalaxy", + top_package="autogalaxy", + packages=_AUTOGALAXY_STACK + _SHARED_EXTRAS, + workspace_repo="https://github.com/PyAutoLabs/autogalaxy_workspace", + workspace_dir="/content/autogalaxy_workspace", + gpu_note="much faster", + ), + "autolens": dict( + project_name="PyAutoLens", + top_package="autolens", + packages=_AUTOLENS_STACK + _SHARED_EXTRAS, + workspace_repo="https://github.com/PyAutoLabs/autolens_workspace", + workspace_dir="/content/autolens_workspace", + gpu_note="> 50 times faster", + ), + "howtofit": dict( + project_name="PyAutoFit", + top_package="autofit", + packages=_AUTOFIT_STACK + _SHARED_EXTRAS, + workspace_repo="https://github.com/PyAutoLabs/HowToFit", + workspace_dir="/content/HowToFit", + gpu_note="faster", + ), + "howtogalaxy": dict( + project_name="PyAutoGalaxy", + top_package="autogalaxy", + packages=_AUTOGALAXY_STACK + _SHARED_EXTRAS, + workspace_repo="https://github.com/PyAutoLabs/HowToGalaxy", + workspace_dir="/content/HowToGalaxy", + gpu_note="much faster", + ), + "howtolens": dict( + project_name="PyAutoLens", + top_package="autolens", + packages=_AUTOLENS_STACK + _SHARED_EXTRAS, + workspace_repo="https://github.com/PyAutoLabs/HowToLens", + workspace_dir="/content/HowToLens", + gpu_note="> 50 times faster", + ), +} + + +def check_jax_using_gpu(log_gpu_warning: bool = True) -> bool: + """ + Report the devices JAX sees and return True if no GPU / TPU is available. + + Considers every device (not just the last one enumerated); an empty device + list counts as no accelerator. + """ + import jax + + devices = jax.devices() + + has_accelerator = False + + for d in devices: + kind = str(d).lower() + + if "gpu" in kind or "cuda" in kind: + print(f"GPU / CUDA detected: {d}") + has_accelerator = True + elif "tpu" in kind: + print(f"TPU detected: {d}") + has_accelerator = True + elif "cpu" in kind: + print(f"CPU detected: {d}") + else: + print(f"Other device detected: {d}") + + no_gpu = not has_accelerator + + if no_gpu and log_gpu_warning: + + logger.info( + """ + JAX did not detect a GPU or TPU device and is using the CPU for computations. + + PyAuto runs significantly faster with a GPU, so it is recommended to reinstall + JAX with GPU support, if you have a GPU available, by following the JAX + installation instructions. + """ + ) + + return no_gpu + + +def _installed_version(top_package: str): + """The installed version string of ``top_package``, or None if unknown.""" + try: + from importlib.metadata import version + + return version(top_package) + except Exception: + return None + + +def _clone_workspace(workspace_repo: str, workspace_dir: str, top_package: str) -> None: + """ + Shallow-clone the workspace, preferring the tag matching the installed + release of ``top_package`` so the cloned configs / datasets match the + pip-installed code; fall back to the default branch when no such tag + exists (e.g. a dev install). + """ + import subprocess + + if Path(workspace_dir).exists(): + print("Workspace already exists so not cloning again.") + return + + tag = _installed_version(top_package) + + if tag is not None: + result = subprocess.run( + ["git", "clone", "--depth", "1", "--branch", tag, workspace_repo, workspace_dir] + ) + if result.returncode == 0: + return + print(f"No workspace tag '{tag}' found; cloning the default branch instead.") + + subprocess.run( + ["git", "clone", "--depth", "1", workspace_repo, workspace_dir], check=True + ) + + +def _gpu_error_message(project: str, spec: dict) -> str: + return f""" + No GPU detected in Google Colab. {spec["project_name"]} runs {spec["gpu_note"]} with a GPU, + so switch GPU on in Colab settings. + + To do this: + + - Click "Runtime" in the top menu. + - Click "Change runtime type". + - Under "Hardware accelerator", select one of the "GPU" options available. + + You can set up Colab with a CPU (e.g. if GPUs are unavailable) by calling: + + setup_colab.setup("{project}", raise_error_if_not_gpu=False) + """ + + +def _colab_setup( + *, + project_name: str, + workspace_repo: str, + workspace_dir: str, + packages: list, + top_package: str, + raise_error_if_not_gpu: bool, + gpu_error_message: str, +) -> None: + """ + Shared Google Colab setup for the PyAuto ecosystem. + + A no-op (bar a short message) when not running inside Google Colab. + """ + import os + import subprocess + import sys + + try: + import google.colab # noqa: F401 + except ImportError: + print( + f""" + You are not running in a Google Colab environment so cannot use the setup_colab() function. + + You should therefore have {project_name} installed locally in your environment already (e.g. via pip or + conda) and can run the rest of your script normally. + + You may now continue running your script or Notebook. + """ + ) + return + + # JAX performance tweak, applied only on Colab. + os.environ["XLA_FLAGS"] = _XLA_FLAGS + + no_gpu = check_jax_using_gpu(log_gpu_warning=False) + + if no_gpu and raise_error_if_not_gpu: + raise RuntimeError(gpu_error_message) + + # --- Install packages (no deps — Colab ships the scientific stack) --- + print() + print(f"Now Installing {project_name} and setting up Colab Environment:") + + subprocess.check_call( + [sys.executable, "-m", "pip", "install", *packages, "--no-deps"] + ) + + _clone_workspace(workspace_repo, workspace_dir, top_package) + + # --- Configure autoconf paths --- + from autoconf import conf + + os.chdir(workspace_dir) + + conf.instance.push( + new_path=Path(workspace_dir) / "config", + output_path=Path(workspace_dir) / "output", + ) + + print( + f""" + ***Google Colab Setup Complete, which included:*** + + - Installation of {project_name} and other required packages. + - Cloning of the workspace GitHub repository. + - Setting up environment variables for JAX for improved performance. + - Setting the configuration paths to the workspace config and output folders suitable for Colab. + """ + ) + + +def setup(project: str, raise_error_if_not_gpu: bool = False) -> None: + """ + Set up Google Colab for a PyAuto notebook repository. + + Parameters + ---------- + project + One of: "autofit", "autogalaxy", "autolens", "howtofit", + "howtogalaxy", "howtolens". + raise_error_if_not_gpu + If True, raise instead of continuing when Colab has no GPU / TPU + runtime enabled. + """ + try: + spec = _PROJECTS[project] + except KeyError: + raise KeyError( + f"Unknown Colab project '{project}'. Choose from: {sorted(_PROJECTS)}" + ) + + _colab_setup( + project_name=spec["project_name"], + workspace_repo=spec["workspace_repo"], + workspace_dir=spec["workspace_dir"], + packages=spec["packages"], + top_package=spec["top_package"], + raise_error_if_not_gpu=raise_error_if_not_gpu, + gpu_error_message=_gpu_error_message(project, spec), + ) + + +def for_autofit(raise_error_if_not_gpu: bool = False) -> None: + """Google Colab helper for PyAutoFit (the autofit_workspace).""" + setup("autofit", raise_error_if_not_gpu=raise_error_if_not_gpu) + + +def for_autogalaxy(raise_error_if_not_gpu: bool = True) -> None: + """Google Colab helper for PyAutoGalaxy (the autogalaxy_workspace).""" + setup("autogalaxy", raise_error_if_not_gpu=raise_error_if_not_gpu) + + +def for_autolens(raise_error_if_not_gpu: bool = True) -> None: + """Google Colab helper for PyAutoLens (the autolens_workspace).""" + setup("autolens", 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) + + +def for_howtogalaxy(raise_error_if_not_gpu: bool = False) -> None: + """Google Colab helper for the HowToGalaxy lectures.""" + setup("howtogalaxy", raise_error_if_not_gpu=raise_error_if_not_gpu) + + +def for_howtolens(raise_error_if_not_gpu: bool = False) -> None: + """Google Colab helper for the HowToLens lectures.""" + setup("howtolens", raise_error_if_not_gpu=raise_error_if_not_gpu) diff --git a/test_autoconf/test_setup_colab.py b/test_autoconf/test_setup_colab.py new file mode 100644 index 0000000..ba7b7e7 --- /dev/null +++ b/test_autoconf/test_setup_colab.py @@ -0,0 +1,161 @@ +import subprocess +import sys +import types +from unittest import mock + +import pytest + +from autoconf import setup_colab + + +class FakeDevice: + def __init__(self, kind): + self.kind = kind + + def __str__(self): + return self.kind + + +@pytest.fixture(name="fake_jax") +def make_fake_jax(monkeypatch): + """ + Install a stub ``jax`` module (library unit tests never import real JAX) + whose device list the test controls. + """ + module = types.ModuleType("jax") + module.devices = lambda: [] + monkeypatch.setitem(sys.modules, "jax", module) + return module + + +class TestCheckJaxUsingGpu: + def test_gpu_detected(self, fake_jax): + fake_jax.devices = lambda: [FakeDevice("cuda:0")] + assert setup_colab.check_jax_using_gpu() is False + + def test_tpu_detected(self, fake_jax): + fake_jax.devices = lambda: [FakeDevice("TPU_0")] + assert setup_colab.check_jax_using_gpu() is False + + def test_cpu_only(self, fake_jax): + fake_jax.devices = lambda: [FakeDevice("TFRT_CPU_0")] + assert setup_colab.check_jax_using_gpu() is True + + def test_any_accelerator_counts_regardless_of_order(self, fake_jax): + # Regression: the old implementation kept only the last device's + # status, so [gpu, cpu] wrongly reported no accelerator. + fake_jax.devices = lambda: [FakeDevice("cuda:0"), FakeDevice("TFRT_CPU_0")] + assert setup_colab.check_jax_using_gpu() is False + + def test_empty_device_list(self, fake_jax): + # Regression: the old implementation raised UnboundLocalError here. + fake_jax.devices = lambda: [] + assert setup_colab.check_jax_using_gpu() is True + + +class TestRegistry: + def test_all_projects_have_required_fields(self): + required = { + "project_name", + "top_package", + "packages", + "workspace_repo", + "workspace_dir", + "gpu_note", + } + for project, spec in setup_colab._PROJECTS.items(): + assert required <= set(spec), project + assert spec["packages"][0] == "autoconf", project + assert spec["workspace_repo"].startswith( + "https://github.com/PyAutoLabs/" + ), project + + def test_every_project_has_a_wrapper(self): + for project in setup_colab._PROJECTS: + assert hasattr(setup_colab, f"for_{project}"), project + + def test_unknown_project_raises_with_choices(self): + with pytest.raises(KeyError, match="autogalaxy"): + setup_colab.setup("not_a_project") + + +class TestNoImportSideEffects: + def test_import_does_not_set_xla_flags(self): + # Regression: the module used to set XLA_FLAGS at import time, + # clobbering user environments on every `import autoconf.setup_colab`. + import ast + import inspect + + tree = ast.parse(inspect.getsource(setup_colab)) + for node in tree.body: + assert not isinstance( + node, (ast.Assign, ast.Expr) + ) or "environ" not in ast.dump(node), "module-level os.environ mutation" + + +class TestOutsideColab: + def test_setup_is_a_noop(self, capsys): + # google.colab is not importable here, so setup() must return cleanly + # without installing or cloning anything. + with mock.patch.object(subprocess, "check_call") as check_call: + setup_colab.setup("autolens") + check_call.assert_not_called() + assert "not running in a Google Colab" in capsys.readouterr().out + + def test_wrappers_delegate(self): + with mock.patch.object(setup_colab, "setup") as setup_mock: + setup_colab.for_autolens(raise_error_if_not_gpu=False) + setup_mock.assert_called_once_with( + "autolens", 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 + ) + + +class TestCloneWorkspace: + def test_existing_dir_skips_clone(self, tmp_path, capsys): + with mock.patch.object(subprocess, "run") as run: + setup_colab._clone_workspace( + "https://github.com/PyAutoLabs/autolens_workspace", + str(tmp_path), + "autolens", + ) + run.assert_not_called() + assert "not cloning again" in capsys.readouterr().out + + def test_clones_release_tag_when_available(self, tmp_path): + target = str(tmp_path / "ws") + with mock.patch.object(setup_colab, "_installed_version", return_value="2026.7.6.649"): + with mock.patch.object( + subprocess, "run", return_value=mock.Mock(returncode=0) + ) as run: + setup_colab._clone_workspace("repo_url", target, "autolens") + run.assert_called_once() + args = run.call_args[0][0] + assert args[:4] == ["git", "clone", "--depth", "1"] + assert ["--branch", "2026.7.6.649"] == args[4:6] + + def test_falls_back_to_default_branch_when_tag_missing(self, tmp_path): + target = str(tmp_path / "ws") + with mock.patch.object(setup_colab, "_installed_version", return_value="1.2.3"): + with mock.patch.object( + subprocess, "run", return_value=mock.Mock(returncode=1) + ) as run: + setup_colab._clone_workspace("repo_url", target, "autolens") + assert run.call_count == 2 + fallback = run.call_args_list[1] + assert "--branch" not in fallback[0][0] + assert fallback[1] == {"check": True} + + def test_falls_back_when_version_unknown(self, tmp_path): + target = str(tmp_path / "ws") + with mock.patch.object(setup_colab, "_installed_version", return_value=None): + with mock.patch.object( + subprocess, "run", return_value=mock.Mock(returncode=0) + ) as run: + setup_colab._clone_workspace("repo_url", target, "autolens") + run.assert_called_once() + assert "--branch" not in run.call_args[0][0]