fix(impact-reg): show IMPACT's messages, explain an unseen GPU

apps/impact_reg/impact_reg_konfai/models/elastix_engine.py

@@ -69,6 +69,18 @@ def _displacement_on(fixed: sitk.Image, transform: sitk.Transform) -> np.ndarray
     return dvf_np
 
 
+def _cuda_hint(captured: list[str], root: Path) -> str:
+    """What to do about an install that cannot see the GPU it was asked to use. A CPU build answers ``-h``
+    and passes for a valid install, so the failure only comes once IMPACT asks for the device."""
+    if not any("CUDA is not available" in line for line in captured):
+        return ""
+    return (
+        f"\nThe elastix-IMPACT install at '{root}' cannot use CUDA with this environment's torch "
+        f"({torch.__version__}): it is a CPU build, or one made for another CUDA. Point KONFAI_ELASTIX_DIR at "
+        "an elastix-IMPACT built against this torch, or run on the CPU (--cpu)."
+    )
+
+
 class ElastixEngine:
     """Run the elastix-IMPACT binary on a fixed/moving pair; return the displacement field on the fixed grid.
 
@@ -345,6 +357,7 @@ def register(
             # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
             # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
             captured: list[str] = []
+            told: set[str] = set()
             iteration_line = re.compile(r"^\d+\s")
             budget = None if self._max_iterations > 0 else (self._iterations or None)
             progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
@@ -358,6 +371,11 @@ def register(
                         resolution = int(stripped.split(":", 1)[1])
                     except ValueError:
                         pass
+                elif stripped.startswith("IMPACT:") and stripped not in told:
+                    # IMPACT says when it ran out of device memory and went on with smaller patches: the run
+                    # is slower for it, and nothing else tells why.
+                    told.add(stripped)
+                    progress.write(stripped)
                 elif iteration_line.match(line):
                     progress.update(1)
                     columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
@@ -371,7 +389,9 @@ def register(
             progress.close()
             returncode = proc.wait()
             if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
+                raise RuntimeError(
+                    f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}{_cuda_hint(captured, self._elastix_root)}"
+                )
 
             transforms = sorted(
                 work.glob("TransformParameters.*-Composite.itk.txt"),

apps/impact_reg/impact_reg_konfai/models/elastix_install.py

@@ -168,6 +168,21 @@ def extract_archive(archive: Path, dst_dir: Path) -> None:
     archive.unlink()
 
 
+_NO_CUDA_ASSET = (
+    "The CUDA asset links the CUDA 12 runtime, which this environment's torch does not ship: it could not "
+    "load. For the GPU, point KONFAI_ELASTIX_DIR at an elastix-IMPACT built against this torch."
+)
+
+
+def torch_cuda_flavor() -> str | None:
+    """The CUDA asset the environment's torch can load, ``None`` for a CPU torch or one built for another
+    CUDA: no asset bundles LibTorch, so the binary finds the CUDA runtime where torch keeps its own."""
+    import torch
+
+    cuda = torch.version.cuda
+    return "cu128" if cuda is not None and cuda.split(".")[0] == "12" else None
+
+
 def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool):
     os_name = platform.system()
     arch = normalize_arch(platform.machine())
@@ -185,13 +200,13 @@ def install_elastix_impact(install_path: Path, force_cuda: bool, force_cpu: bool
             raise NameError(
                 "CUDA forced but NVIDIA driver/GPU not suitable. Detected: has_nvidia={has_nvidia}, driver={drv}"
             )
-
+        if torch_cuda_flavor() is None:
+            raise NameError(_NO_CUDA_ASSET)
         flavor = "cu128"
-    elif not force_cpu:
-        if has_nvidia and driver_ok_for_cuda(os_name, drv):
-            flavor = "cu128"
-        else:
-            flavor = "cpu"
+    elif not force_cpu and has_nvidia and driver_ok_for_cuda(os_name, drv):
+        flavor = torch_cuda_flavor() or "cpu"
+        if flavor == "cpu":
+            print(f"{_NO_CUDA_ASSET} Installing the CPU asset.", flush=True)
 
     print(f"System: {os_name} {arch}", flush=True)
     print(f"NVIDIA: {has_nvidia}, driver={drv}", flush=True)

apps/impact_reg/tests/unit/test_elastix_install.py

@@ -49,3 +49,43 @@ def fake_download(url: str, dst: Path) -> None:
         "lib/libANNlib.so",
     ]
     assert (install / "bin" / "elastix").read_bytes() == b"new"
+
+
+@pytest.mark.parametrize(("torch_cuda", "asset"), [("12.8", "cu128"), ("13.0", "cpu"), (None, "cpu")])
+def test_the_cuda_asset_is_fetched_only_for_a_torch_that_can_load_it(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, torch_cuda: str | None, asset: str
+) -> None:
+    # The CUDA asset links the CUDA 12 runtime from where torch keeps its own. Under a torch built for
+    # CUDA 13 it could not load at all, and the install failed where the CPU asset would have run.
+    import torch
+
+    fetched: list[str] = []
+
+    def fake_download(url: str, dst: Path) -> None:
+        fetched.append(url)
+        with zipfile.ZipFile(dst, "w") as archive:
+            archive.writestr("bin/elastix", b"new")
+
+    monkeypatch.setattr(elastix_install, "download_file", fake_download)
+    monkeypatch.setattr(elastix_install, "detect_nvidia_driver", lambda: (True, (595, 84)))
+    monkeypatch.setattr(elastix_install.platform, "system", lambda: "Linux")
+    monkeypatch.setattr(elastix_install.platform, "machine", lambda: "x86_64")
+    monkeypatch.setattr(torch.version, "cuda", torch_cuda)
+
+    elastix_install.install_elastix_impact(tmp_path / "elastix-impact", force_cuda=False, force_cpu=False)
+
+    assert fetched[0].endswith(f"-{asset}.zip")
+
+
+def test_a_forced_cuda_install_refuses_a_torch_that_cannot_load_the_asset(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    import torch
+
+    monkeypatch.setattr(elastix_install, "detect_nvidia_driver", lambda: (True, (595, 84)))
+    monkeypatch.setattr(elastix_install.platform, "system", lambda: "Linux")
+    monkeypatch.setattr(elastix_install.platform, "machine", lambda: "x86_64")
+    monkeypatch.setattr(torch.version, "cuda", "13.0")
+
+    with pytest.raises(NameError, match="KONFAI_ELASTIX_DIR"):
+        elastix_install.install_elastix_impact(tmp_path / "elastix-impact", force_cuda=True, force_cpu=False)

apps/impact_reg/tests/unit/test_engine_contracts.py

@@ -79,6 +79,51 @@ def elastix_must_not_run(*args, **kwargs):
     assert not field.any()
 
 
+def _elastix_run(monkeypatch, lines: list[str], code: int):
+    """``ElastixEngine.register`` over a subprocess that prints ``lines`` and exits with ``code``."""
+
+    class Process:
+        stdout = iter(lines)
+
+        def wait(self) -> int:
+            return code
+
+    monkeypatch.setattr(elastix_engine_module.subprocess, "Popen", lambda *args, **kwargs: Process())
+    monkeypatch.setattr(elastix_engine_module, "loader_env", lambda root: {})
+    engine = SimpleNamespace(
+        _local_models=[],
+        _elastix_bin="elastix",
+        _elastix_root=Path("/opt/elastix-impact"),
+        _stage_parameter_maps=lambda work, device_index: [],
+        _max_iterations=0,
+        _iterations=None,
+    )
+    fixed = sitk.Image([6, 5, 4], sitk.sitkFloat32)
+    return ElastixEngine.register(engine, fixed, sitk.Image(fixed), 0)
+
+
+def test_what_impact_says_about_device_memory_is_shown_once(monkeypatch, capsys) -> None:
+    # IMPACT goes on with smaller patches when the device runs out of memory: the run is slower for it, and
+    # elastix's output is otherwise shown on a failure only.
+    retry = "IMPACT: the model ran out of device memory on the whole image; retrying with a patch of (96 160 192).\n"
+    with pytest.raises(FileNotFoundError, match="no composite transform"):
+        _elastix_run(monkeypatch, ["Resolution: 0\n", retry, retry, "1 -0.25 3.0\n"], 0)
+
+    assert capsys.readouterr().out.count("retrying with a patch of (96 160 192)") == 1
+
+
+def test_a_gpu_the_install_cannot_see_says_what_to_do(monkeypatch) -> None:
+    # A CPU build answers `-h` and passes for a valid install: the failure only comes mid-registration.
+    unseen = "Description: ITK ERROR: ImpactMetric(0x5e): CUDA is not available. Please check your CUDA installation.\n"
+    with pytest.raises(RuntimeError, match="KONFAI_ELASTIX_DIR") as raised:
+        _elastix_run(monkeypatch, [unseen], 1)
+    assert str(Path("/opt/elastix-impact")) in str(raised.value)
+
+    with pytest.raises(RuntimeError) as other:
+        _elastix_run(monkeypatch, ["no such parameter file\n"], 1)
+    assert "KONFAI_ELASTIX_DIR" not in str(other.value)
+
+
 def test_elastix_engine_refuses_an_empty_parameter_map_list() -> None:
     # 'resolutions' rewrites a template's resolution-dependent lines; it never creates one. Without a
     # map elastix would launch with no -p and die in a cryptic subprocess error.