diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index a26f5f5..1d97eb9 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -8,7 +8,21 @@ on: jobs: test: + name: test (${{ matrix.install }}) runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + include: + # The full install, with the optional XLB backend, runs every test. + - install: with xlb extra + sync: uv sync --group dev --extra xlb + without_xlb: "0" + # The default install leaves XLB, JAX and Warp out; the XLB tests skip + # and the missing-extra error is checked against a real environment. + - install: without xlb extra + sync: uv sync --group dev + without_xlb: "1" steps: - uses: actions/checkout@v6 @@ -24,12 +38,15 @@ jobs: python-version: "3.12" - name: Install dependencies - run: uv sync --group dev + run: ${{ matrix.sync }} - name: Run tests - run: uv run pytest tests/ -v --cov=kabs --cov-report=xml + env: + KABS_TEST_WITHOUT_XLB: ${{ matrix.without_xlb }} + run: uv run --no-sync pytest tests/ -v --cov=kabs --cov-report=xml - name: Upload coverage to Codecov + if: matrix.without_xlb == '0' uses: codecov/codecov-action@v5 with: token: ${{ secrets.CODECOV_TOKEN }} diff --git a/README.md b/README.md index eb1fcfc..fbedd7b 100644 --- a/README.md +++ b/README.md @@ -22,11 +22,22 @@ cd kabs uv sync ``` -Alternatively, install with `pip install -e .`. This installs both solvers: -Taichi and XLB, with `warp-lang==1.10.0` pinned for compatibility with XLB. +Alternatively, install with `pip install -e .`. This installs the default +Taichi solver. The XLB solver, with its JAX and Warp compute backends, is the +optional `xlb` extra, which adds XLB, JAX and `warp-lang==1.10.0` (pinned for +compatibility with XLB): -Key dependencies: `taichi` (GPU/CPU acceleration), `xlb`, `numpy`, `pyevtk`. -Optional: `porespy` (used in the examples below to generate synthetic images). +```bash +uv sync --extra xlb # in a clone +pip install -e ".[xlb]" # or with pip +``` + +Selecting `backend="xlb"` without the extra raises an `ImportError` that names +it. + +Key dependencies: `taichi` (GPU/CPU acceleration), `numpy`, `pyevtk`. +Optional: `xlb` and `warp-lang` (the `xlb` extra), and `porespy` (used in the +examples below to generate synthetic images). --- @@ -71,8 +82,9 @@ ti.init(arch=ti.metal) # Apple Silicon GPU ### Solver backends `solve_flow` uses the Taichi implementation by default. Select XLB through the -same API; its default JAX compute backend works on CPU and supports multi-GPU -execution on supported accelerator platforms: +same API once the `xlb` extra is installed; its default JAX compute backend +works on CPU and supports multi-GPU execution on supported accelerator +platforms: ```python # Default: solve_flow(im, backend="taichi") diff --git a/docs/warp-backend-readiness.md b/docs/warp-backend-readiness.md index 6354806..abc3807 100644 --- a/docs/warp-backend-readiness.md +++ b/docs/warp-backend-readiness.md @@ -222,7 +222,7 @@ fresh environment: ```powershell git fetch origin git switch finalizing-warp-option -uv sync --group dev +uv sync --group dev --extra xlb ``` Verify that Warp detects CUDA: diff --git a/kabs/_solve_flow_xlb.py b/kabs/_solve_flow_xlb.py index 354ac1c..3198b13 100644 --- a/kabs/_solve_flow_xlb.py +++ b/kabs/_solve_flow_xlb.py @@ -224,10 +224,13 @@ def solve_flow_xlb( call it with custom settings before invoking this function, set ``compute_backend`` to match your prior ``xlb.init()`` call. - **Installation**:: + **Installation**: XLB and Warp are not installed with kabs by default. + They come with the optional ``xlb`` extra:: - pip install "xlb>=0.3.1" "warp-lang==1.10.0" # CPU/JAX - pip install "xlb[cuda]>=0.3.1" "warp-lang==1.10.0" # NVIDIA GPU via JAX + pip install "kabs[xlb]" # CPU/JAX and Warp + pip install "kabs[xlb]" "xlb[cuda]>=0.3.1" # NVIDIA GPU via JAX + + Without it, this function raises ``ImportError`` naming the extra. """ if not isinstance(direction, str) or direction.lower() not in _FACE_NAMES: raise ValueError(f"direction must be 'x', 'y', or 'z', got {direction!r}") @@ -271,8 +274,11 @@ def solve_flow_xlb( from xlb.operator.macroscopic import Macroscopic except ImportError as exc: raise ImportError( - "XLB is required for solve_flow_xlb. Install compatible dependencies with:\n" - " pip install 'xlb>=0.3.1' 'warp-lang==1.10.0'\n" + "The XLB backend needs kabs's optional 'xlb' extra, which installs " + "XLB and Warp. Install it with:\n" + " pip install 'kabs[xlb]'\n" + "or, in a clone of the kabs repository:\n" + " uv sync --extra xlb\n" ) from exc _backend_map = {"jax": ComputeBackend.JAX, "warp": ComputeBackend.WARP} diff --git a/pyproject.toml b/pyproject.toml index bd8df7b..e1c66c9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,10 @@ dependencies = [ "trame>=3.12.0", "trame-vtk>=2.11.3", "trame-vuetify>=3.2.1", +] + +[project.optional-dependencies] +xlb = [ "xlb>=0.3.1", "warp-lang==1.10.0", ] diff --git a/tests/test_backend_import_isolation.py b/tests/test_backend_import_isolation.py index eabc04c..af9a346 100644 --- a/tests/test_backend_import_isolation.py +++ b/tests/test_backend_import_isolation.py @@ -33,6 +33,21 @@ def test_import_kabs_does_not_import_taichi(): assert completed.returncode == 0, completed.stdout + completed.stderr +def test_import_kabs_does_not_import_the_xlb_stack(): + completed = _run_fresh_python( + """ + import sys + import kabs + + loaded = {"xlb", "jax", "warp"} & {name.split(".")[0] for name in sys.modules} + assert not loaded, sorted(loaded) + assert "kabs._warp_convergence" not in sys.modules + """ + ) + + assert completed.returncode == 0, completed.stdout + completed.stderr + + def test_public_single_phase_solver_export_remains_lazy(): completed = _run_fresh_python( """ diff --git a/tests/test_convergence_xlb.py b/tests/test_convergence_xlb.py index 9aecd00..01312f4 100644 --- a/tests/test_convergence_xlb.py +++ b/tests/test_convergence_xlb.py @@ -3,14 +3,17 @@ import importlib import os -import jax -import jax.numpy as jnp import numpy as np import pytest -from kabs import compute_permeability, solve_flow_xlb -from kabs._convergence import ConvergenceConfig -from kabs._solve_flow_xlb import ( +pytest.importorskip("xlb", reason="xlb is not installed; skipping XLB tests") + +import jax # noqa: E402 +import jax.numpy as jnp # noqa: E402 + +from kabs import compute_permeability, solve_flow_xlb # noqa: E402 +from kabs._convergence import ConvergenceConfig # noqa: E402 +from kabs._solve_flow_xlb import ( # noqa: E402 _convergence_sums_to_host, _get_jax_convergence_reducer, ) diff --git a/tests/test_solve_flow_xlb_validation.py b/tests/test_solve_flow_xlb_validation.py index f76131d..c7ac1da 100644 --- a/tests/test_solve_flow_xlb_validation.py +++ b/tests/test_solve_flow_xlb_validation.py @@ -1,11 +1,13 @@ """Public validation for the optional XLB solver.""" import builtins +import importlib.util +import os import numpy as np import pytest -from kabs import solve_flow_xlb +from kabs import solve_flow, solve_flow_xlb _IMAGE = np.ones((3, 3, 3), dtype=np.int8) @@ -22,6 +24,18 @@ def guarded_import(name, *args, **kwargs): monkeypatch.setattr(builtins, "__import__", guarded_import) +def _hide_xlb(monkeypatch): + """Make ``import xlb`` fail as it does when the ``xlb`` extra is not installed.""" + original_import = builtins.__import__ + + def import_without_xlb(name, *args, **kwargs): + if name == "xlb" or name.startswith("xlb."): + raise ModuleNotFoundError(f"No module named {name!r}", name=name) + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", import_without_xlb) + + @pytest.mark.parametrize("direction", ["invalid", None]) def test_invalid_direction_is_reported_without_importing_xlb(direction, monkeypatch): _block_xlb_import(monkeypatch) @@ -36,3 +50,22 @@ def test_invalid_compute_backend_is_reported_without_importing_xlb( _block_xlb_import(monkeypatch) with pytest.raises(ValueError, match="compute_backend"): solve_flow_xlb(_IMAGE, compute_backend=compute_backend) + + +@pytest.mark.parametrize("compute_backend", ["jax", "warp"]) +def test_missing_xlb_extra_is_named_in_the_error(compute_backend, monkeypatch): + _hide_xlb(monkeypatch) + with pytest.raises(ImportError, match=r"kabs\[xlb\]") as excinfo: + solve_flow(_IMAGE, backend="xlb", compute_backend=compute_backend) + assert isinstance(excinfo.value.__cause__, ModuleNotFoundError) + + +@pytest.mark.skipif( + os.environ.get("KABS_TEST_WITHOUT_XLB") != "1", + reason="set KABS_TEST_WITHOUT_XLB=1 in an environment installed without the xlb extra", +) +def test_install_without_the_extra_omits_the_xlb_stack(): + for module in ("xlb", "warp", "jax"): + assert importlib.util.find_spec(module) is None, module + with pytest.raises(ImportError, match=r"kabs\[xlb\]"): + solve_flow(_IMAGE, backend="xlb") diff --git a/uv.lock b/uv.lock index 932059e..94fa041 100644 --- a/uv.lock +++ b/uv.lock @@ -890,6 +890,10 @@ dependencies = [ { name = "trame" }, { name = "trame-vtk" }, { name = "trame-vuetify" }, +] + +[package.optional-dependencies] +xlb = [ { name = "warp-lang" }, { name = "xlb" }, ] @@ -915,9 +919,10 @@ requires-dist = [ { name = "trame", specifier = ">=3.12.0" }, { name = "trame-vtk", specifier = ">=2.11.3" }, { name = "trame-vuetify", specifier = ">=3.2.1" }, - { name = "warp-lang", specifier = "==1.10.0" }, - { name = "xlb", specifier = ">=0.3.1" }, + { name = "warp-lang", marker = "extra == 'xlb'", specifier = "==1.10.0" }, + { name = "xlb", marker = "extra == 'xlb'", specifier = ">=0.3.1" }, ] +provides-extras = ["xlb"] [package.metadata.requires-dev] dev = [