Make the PyTorch install and model cache survive real-world Windows setups

ZebrafishEmbryoAnalyzer/ZebrafishEmbryoAnalyzerCore/length.py

@@ -658,12 +658,46 @@ def apply_mask(original_image, mask):
 
     return masked_image
 
+def select_torch_device(torch, probe_model=None, probe_input_shape=(1, 3, 256, 256)):
+    """Pick a usable torch device, falling back to CPU when CUDA is unusable.
+
+    torch.cuda.is_available() only confirms a CUDA runtime/driver is present,
+    not that the installed build's compiled kernels cover this GPU's compute
+    capability. On a mismatch, a real kernel launch fails with "CUDA error: no
+    kernel image is available for execution on the device".
+
+    A trivial canary op (e.g. a bare add) is not a reliable stand-in for that
+    check: observed on real hardware to succeed — after a slow one-time CUDA
+    context/JIT warmup — on a GPU where the model's own conv/batchnorm kernels
+    still failed immediately afterwards. Probing with the real model on a dummy
+    input of its expected shape exercises the same kernels the model actually
+    uses, so the failure (if any) shows up here instead of mid-analysis.
+    When probe_model is None (no model to test yet), only is_available() is
+    checked.
+    """
+    if not torch.cuda.is_available():
+        return torch.device("cpu")
+    if probe_model is None:
+        return torch.device("cuda")
+    try:
+        probe_model.to("cuda")
+        with torch.no_grad():
+            probe_model(torch.zeros(*probe_input_shape, device="cuda"))
+        return torch.device("cuda")
+    except RuntimeError:
+        probe_model.to("cpu")
+        return torch.device("cpu")
+
+
 def classification_curvature(image, mask, model, use_threshold, threshold):
     import torch
     import torch.nn.functional as F
     import torchvision.transforms as T
     import cv2  # deferred: only needed at call time
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # The model's actual device reflects whatever load_model's CUDA probe decided —
+    # recomputing availability independently here could disagree with it and feed a
+    # cuda tensor to a cpu-fallback model (or vice versa).
+    device = next(model.parameters()).device
 
     masked_image = apply_mask(image, mask)
 
@@ -714,7 +748,6 @@ def load_model(model_path: str):
     import torch.nn as nn
     import torch.nn.functional as F
     import timm
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     import logging as _logging
     _log = _logging.getLogger(__name__)
@@ -752,7 +785,9 @@ def forward(self, x):
     _log.debug("Using curvature model params: %s", best_params)
 
     try:
-        state_dict = torch.load(model_path, map_location=device, weights_only=True)
+        # Always deserialize to CPU first — the final device is decided below, after
+        # the model is built, by actually probing it (see select_torch_device).
+        state_dict = torch.load(model_path, map_location=torch.device("cpu"), weights_only=True)
     except Exception as exc:
         raise RuntimeError(
             f"Failed to load curvature model from {model_path!r} with safe loading. "
@@ -786,9 +821,11 @@ def forward(self, x):
             f"Curvature checkpoint at {model_path!r} is missing required keys: {missing}. "
             "The checkpoint may be incomplete or incompatible. Re-download the model."
         )
-    model = model_instance.to(device)
-    model.eval()
-    return model
+    # Decide the device by actually probing model_instance (see select_torch_device) —
+    # it moves the model to its final device itself, cuda or a cpu fallback.
+    select_torch_device(torch, probe_model=model_instance)
+    model_instance.eval()
+    return model_instance
 
 
 def plot_edges_with_curvature(mask, min_contour_length, window_size_ratio):

ZebrafishEmbryoAnalyzer/ZebrafishEmbryoAnalyzerLib/dependency_installer.py

@@ -130,6 +130,23 @@ def install_pytorch_extension() -> bool:
     return bool(manager.installExtensionFromServer(PYTORCH_EXTENSION_NAME))
 
 
+def _windows_torch_version_requirement() -> str | None:
+    """Version cap that keeps installTorch() off the torch 2.13 Windows regression.
+
+    torch 2.13 moved to PEP 639 license-file collection, which copies its entire
+    nested third-party tree (kineto/dynolog/prometheus-cpp/civetweb/duktape, ...)
+    into torch-<version>.dist-info/licenses/. The deepest resulting path fails
+    Windows installs with "WinError 206: filename or extension is too long"
+    whenever site-packages is longer than ~86 characters — close to what a
+    default Slicer install already uses. Confirmed as an upstream torch
+    regression, not something specific to this install path; a fix has been
+    proposed for the official Slicer PyTorch extension but is not merged yet.
+    Not needed on macOS/Linux, which have no such path-length limit.
+    """
+    import sys
+    return "<2.13.0" if sys.platform == "win32" else None
+
+
 def _install_torch() -> str:
     """Install torch and torchvision through the PyTorch extension.
 
@@ -147,7 +164,10 @@ def _install_torch() -> str:
             )
         return "restart"
 
-    if torch_logic.installTorch(askConfirmation=False) is None:
+    if torch_logic.installTorch(
+        askConfirmation=False,
+        torchVersionRequirement=_windows_torch_version_requirement(),
+    ) is None:
         raise RuntimeError("PyTorch could not be installed through the PyTorch extension.")
     return "ok"
 

ZebrafishEmbryoAnalyzer/ZebrafishEmbryoAnalyzerLib/inference_worker.py

@@ -7,9 +7,11 @@
 Exit codes:
   0 -- success
   1 -- analysis exception
-  2 -- model not cached
+  2 -- model not cached (genuinely missing; a download can fix this)
   3 -- bad/unreadable request
   4 -- result write failure
+  5 -- model preload failed for a reason other than "not cached" (e.g. a broken
+       torch install) -- retrying without fixing the underlying cause will not help
 """
 
 import json
@@ -19,7 +21,7 @@
 
 
 def run_worker(request_path: str) -> int:
-    """Execute inference from request_path. Returns exit code 0-4."""
+    """Execute inference from request_path. Returns exit code 0-5."""
     # --- 1. Read and validate request ---
     try:
         with open(request_path, "r", encoding="utf-8") as fh:
@@ -46,11 +48,17 @@ def run_worker(request_path: str) -> int:
         preload_params = dict(params)
         preload_params["model_id"] = model_id
         preload_models(preload_params)
-    except ModelNotCachedError:
+    except ModelNotCachedError as exc:
+        _write_error(result_json_path, 2, str(exc))
         return 2
     except Exception as exc:
-        print(f"preload_models failed: {exc}", file=sys.stderr)
-        return 2
+        # Anything other than a genuinely missing model — e.g. a broken torch
+        # install. Exit code 2's UI advice ("run again to trigger a download")
+        # would be actively wrong here, so this gets its own code and the real
+        # message is persisted rather than only printed to stderr, which the
+        # caller does not read for this path.
+        _write_error(result_json_path, 5, f"preload_models failed: {exc}")
+        return 5
 
     # --- 4. Run analysis ---
     n = len(image_paths)

ZebrafishEmbryoAnalyzer/ZebrafishEmbryoAnalyzerLib/logic.py

@@ -78,26 +78,28 @@ def preload_models(params: dict) -> None:
     _install_model_cache()
     from ZebrafishEmbryoAnalyzerCore.length import load_model
     from ZebrafishEmbryoAnalyzerLib.errors import ModelNotCachedError
-    from ZebrafishEmbryoAnalyzerLib.model_manifest import MODEL_SETS, get_cached_path, MODELS
+    from ZebrafishEmbryoAnalyzerLib.model_manifest import (
+        MODEL_SETS, get_cached_path, MODELS, verify_checksum,
+    )
 
     model_id = params.get("model_id", "general")
     model_set = MODEL_SETS.get(model_id, MODEL_SETS["general"])
 
     if params.get("curvature", True) and "curvature" not in _MODEL_CACHE:
         curvature_entry = MODELS["curvature"]
         curvature_path = get_cached_path(curvature_entry)
-        if not curvature_path.exists():
+        if not verify_checksum(curvature_path, curvature_entry["sha256"]):
             raise ModelNotCachedError(
-                f"{curvature_entry['label']} not found at {curvature_path}. "
+                f"{curvature_entry['label']} missing or corrupted at {curvature_path}. "
                 "Download models first."
             )
         _MODEL_CACHE["curvature"] = load_model(str(curvature_path))
 
     body_entry = model_set["body"]
     body_path = get_cached_path(body_entry)
-    if not body_path.exists():
+    if not verify_checksum(body_path, body_entry["sha256"]):
         raise ModelNotCachedError(
-            f"{body_entry['label']} not found at {body_path}. "
+            f"{body_entry['label']} missing or corrupted at {body_path}. "
             "Download models first."
         )
     _cached_load_unet(
@@ -109,9 +111,9 @@ def preload_models(params: dict) -> None:
     if params.get("eyes", False):
         eye_entry = model_set["eye"]
         eye_path = get_cached_path(eye_entry)
-        if not eye_path.exists():
+        if not verify_checksum(eye_path, eye_entry["sha256"]):
             raise ModelNotCachedError(
-                f"{eye_entry['label']} not found at {eye_path}. "
+                f"{eye_entry['label']} missing or corrupted at {eye_path}. "
                 "Download models first."
             )
         _cached_load_unet(
@@ -206,34 +208,36 @@ def analyse_images(image_paths: list, params: dict,
         compute_eye_metrics,
     )
     from ZebrafishEmbryoAnalyzerLib.errors import ModelNotCachedError
-    from ZebrafishEmbryoAnalyzerLib.model_manifest import MODEL_SETS, get_cached_path, MODELS
+    from ZebrafishEmbryoAnalyzerLib.model_manifest import (
+        MODEL_SETS, get_cached_path, MODELS, verify_checksum,
+    )
 
     um_per_px = float(params.get("um_per_px", 22.99))
     include_eyes = params.get("eyes", False)
 
     model_id = params.get("model_id", "general")
     model_set = MODEL_SETS.get(model_id, MODEL_SETS["general"])
 
-    # ---- validate required model files exist before starting ----
+    # ---- validate required model files exist and are not corrupted before starting ----
     body_entry = model_set["body"]
     body_path = get_cached_path(body_entry)
-    if not body_path.exists():
+    if not verify_checksum(body_path, body_entry["sha256"]):
         raise ModelNotCachedError(
-            f"{body_entry['label']} not found at {body_path}. Download models first."
+            f"{body_entry['label']} missing or corrupted at {body_path}. Download models first."
         )
     if include_eyes:
         eye_entry = model_set["eye"]
         eye_path = get_cached_path(eye_entry)
-        if not eye_path.exists():
+        if not verify_checksum(eye_path, eye_entry["sha256"]):
             raise ModelNotCachedError(
-                f"{eye_entry['label']} not found at {eye_path}. Download models first."
+                f"{eye_entry['label']} missing or corrupted at {eye_path}. Download models first."
             )
     if params.get("curvature", True):
         curv_entry = MODELS["curvature"]
         curv_manifest_path = get_cached_path(curv_entry)
-        if not curv_manifest_path.exists():
+        if not verify_checksum(curv_manifest_path, curv_entry["sha256"]):
             raise ModelNotCachedError(
-                f"{curv_entry['label']} not found at {curv_manifest_path}. "
+                f"{curv_entry['label']} missing or corrupted at {curv_manifest_path}. "
                 "Download models first."
             )
 

ZebrafishEmbryoAnalyzer/ZebrafishEmbryoAnalyzerLib/model_manifest.py

@@ -195,7 +195,13 @@ def collect_all_model_entries() -> dict:
 
 def get_missing_models(model_set_dict: dict) -> list:
     """
-    Return model entries whose cached path does not exist or is empty.
+    Return model entries whose cached file is missing or fails checksum verification.
+
+    A truncated or corrupted download (present on disk, non-zero size, wrong
+    content — e.g. an interrupted transfer) must be treated the same as a
+    genuinely missing file. Otherwise it looks "cached" until deserialization
+    fails deep inside analysis with a cryptic pickle error instead of the normal
+    download-prompt flow.
 
     Parameters
     ----------
@@ -205,14 +211,9 @@ def get_missing_models(model_set_dict: dict) -> list:
     Returns
     -------
     list[dict]
-        Subset of model_set_dict.values() that are not yet cached.
+        Subset of model_set_dict.values() that are not yet correctly cached.
     """
-    missing = []
-    for entry in model_set_dict.values():
-        p = get_cached_path(entry)
-        try:
-            if not p.exists() or p.stat().st_size == 0:
-                missing.append(entry)
-        except OSError:
-            missing.append(entry)
-    return missing
+    return [
+        entry for entry in model_set_dict.values()
+        if not verify_checksum(get_cached_path(entry), entry["sha256"])
+    ]

ZebrafishEmbryoAnalyzer/ZebrafishEmbryoAnalyzerLib/widget.py

@@ -1032,6 +1032,13 @@ def _categorize_inference_error(self, message, controller):
             return "Internal error: bad analysis request. Check the application log."
         if exit_code == 4:
             return "Internal error: could not write temporary results. Check disk space."
+        if exit_code == 5:
+            # Model preload failed for a reason other than "not cached" (e.g. a broken
+            # torch install) — unlike exit_code 2, retrying cannot fix this, so show the
+            # real cause instead of advice that would just repeat the same failure.
+            first_line = msg.split("\n")[0].strip() if msg else ""
+            detail = f": {first_line}" if first_line else "."
+            return f"Could not load required models{detail} This is not a missing-download issue — check the application log."
         return "Analysis failed. Check the application log."
 
     def ensure_dependencies(self, purpose="analysis"):

tests/conftest.py

@@ -30,7 +30,8 @@ def synthetic_fish_mask():
 
 @pytest.fixture
 def mock_model_paths(tmp_path):
-    """Patch get_cached_path so existence checks in logic.py pass without real models.
+    """Patch get_cached_path and checksum verification so the existence/integrity
+    checks in logic.py pass without real, correctly-hashed model files.
 
     Returns a dict mapping entry id -> Path so tests can inspect paths if needed.
     """
@@ -44,5 +45,7 @@ def fake_get_cached_path(entry):
         return p
 
     with patch("ZebrafishEmbryoAnalyzerLib.model_manifest.get_cached_path",
-               side_effect=fake_get_cached_path):
+               side_effect=fake_get_cached_path), \
+         patch("ZebrafishEmbryoAnalyzerLib.model_manifest.verify_checksum",
+               return_value=True):
         yield created

tests/test_dependency_installer.py

@@ -283,13 +283,15 @@ def test_install_torch_raises_when_extension_unavailable(monkeypatch):
 
 
 def test_install_torch_uses_pytorch_utils_when_available(monkeypatch):
-    di, _ = _reload_with_slicer(monkeypatch)
+    di, _ = _reload_with_slicer(monkeypatch, platform="linux")
     torch_logic = MagicMock()
     torch_logic.installTorch.return_value = object()
     monkeypatch.setattr(di, "_pytorch_utils_logic", lambda: torch_logic)
 
     assert di._install_torch() == "ok"
-    torch_logic.installTorch.assert_called_once_with(askConfirmation=False)
+    torch_logic.installTorch.assert_called_once_with(
+        askConfirmation=False, torchVersionRequirement=None
+    )
 
 
 def test_install_torch_raises_when_pytorch_utils_fails(monkeypatch):
@@ -302,6 +304,38 @@ def test_install_torch_raises_when_pytorch_utils_fails(monkeypatch):
         di._install_torch()
 
 
+def test_install_torch_pins_below_2_13_on_windows(monkeypatch):
+    """torch 2.13 fails on Windows with WinError 206 (PEP 639 license-file paths
+    exceeding MAX_PATH) whenever site-packages is longer than ~86 characters — close
+    to what a default Slicer install already uses. Pinning below it sidesteps the
+    trigger entirely instead of relying on a shorter install path."""
+    di, _ = _reload_with_slicer(monkeypatch, platform="win32")
+    torch_logic = MagicMock()
+    torch_logic.installTorch.return_value = object()
+    monkeypatch.setattr(di, "_pytorch_utils_logic", lambda: torch_logic)
+
+    assert di._install_torch() == "ok"
+    torch_logic.installTorch.assert_called_once_with(
+        askConfirmation=False, torchVersionRequirement="<2.13.0"
+    )
+
+
+def test_install_torch_no_version_cap_on_macos_or_linux(monkeypatch):
+    """The torch 2.13 long-path failure is Windows-specific (MAX_PATH); macOS and
+    Linux have no such limit, so an unnecessary pin would just hold back an
+    otherwise-fine current release for no reason."""
+    for platform in ("darwin", "linux"):
+        di, _ = _reload_with_slicer(monkeypatch, platform=platform)
+        torch_logic = MagicMock()
+        torch_logic.installTorch.return_value = object()
+        monkeypatch.setattr(di, "_pytorch_utils_logic", lambda: torch_logic)
+
+        di._install_torch()
+        torch_logic.installTorch.assert_called_once_with(
+            askConfirmation=False, torchVersionRequirement=None
+        )
+
+
 def test_pytorch_utils_logic_returns_none_without_extension(monkeypatch):
     di, _ = _reload_with_slicer(monkeypatch)
     monkeypatch.setitem(sys.modules, "PyTorchUtils", None)  # forces ModuleNotFoundError