fix(gpu): 🐛 catch cupy OOM, read the granted GPU, validate configs in prepare

Five defects around running Cellpose and DoG on shared cluster GPUs.

_is_cuda_oom matched a substring on RuntimeError only, but
cupy.cuda.memory.OutOfMemoryError derives from Exception -- so the DoG path
and dilate_labels(use_gpu=True) had no retry at all while Cellpose had four.
The retry now lives in patchworks._gpu, classifies by exception class first
(cupy says "Out of memory allocating", torch says "CUDA out of memory"), and
both plugins use it.

The Cellpose retry also slept up to 300s while keeping its model pinned in
VRAM -- exactly the memory the contending job needed. It now evicts the model
cache before backing off and reloads on retry, seconds against a backoff
measured in minutes.

_get_gpu_memory asked NVML for index 0 unconditionally. NVML enumerates every
GPU on the node regardless of --gres=gpu:1, so on a multi-GPU node this read
a different device's free memory -- and auto tile sizing was built on that
number. It now resolves the device from CUDA_VISIBLE_DEVICES (index or UUID)
and keeps 20% headroom instead of sizing tiles to fill every free byte of a
device we don't own outright.

The DoG plugin never released cupy's pool, which does not return blocks to
the driver on its own, so a long-lived worker's footprint only grew. It now
drops its references and frees the pool per tile.

Config mistakes now fail in prepare, a cheap CPU job, instead of in the first
GPU job after convert and prepare have already run: gpu_memory_gb: "auto" was
truthy so "auto" * 1024**3 built a ~4 GB string; a non-"auto" tile_shape
string became a character tuple; and unknown cellpose:/custom.kwargs keys were
forwarded blind. prepare also stopped reading cfg["cellpose"] unconditionally
for tile_shape: "auto", which raised KeyError for any DoG or threshold config.

Also: the Cellpose version parse blew up at import on a suffixed release like
"4.0rc1" (only ImportError was caught), and do_3D: false with a z>1 tile
silently handed Cellpose a stack it would read as channels -- now an error.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: lguerard

src/patchworks/_chunks.py

@@ -213,23 +213,40 @@ def safe_worker_count(
     return max(1, min(cpu_cap, mem_cap))
 
 
+# Leave room for a co-tenant: info.free is a point-in-time reading of a device
+# we usually do not own outright, and sizing a tile to fill all of it is what
+# turns another job's growth into our OOM.
+_GPU_HEADROOM = 0.8
+
+
 def _get_gpu_memory() -> int:
-    """Return free GPU VRAM in bytes.
+    """Return usable GPU VRAM in bytes for the device we were granted.
+
+    NVML enumerates every GPU on the node regardless of ``--gres=gpu:1``, so
+    asking for index 0 unconditionally reads the *wrong* device's free memory
+    whenever SLURM granted anything but the first one -- and tile sizing was
+    built on that number. Resolve the device from ``CUDA_VISIBLE_DEVICES``,
+    then keep a headroom fraction rather than claiming every free byte.
 
     Returns
     -------
     int
-        Free VRAM of GPU 0 via ``nvidia-ml-py``, or an 8 GiB fallback if the
-        query fails.
+        Usable VRAM, or an 8 GiB fallback if the query fails.
     """
     try:
         import pynvml
 
+        from ._gpu import visible_device_index, visible_device_uuid
+
         pynvml.nvmlInit()
-        handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+        uuid = visible_device_uuid()
+        if uuid is not None:
+            handle = pynvml.nvmlDeviceGetHandleByUUID(uuid.encode())
+        else:
+            handle = pynvml.nvmlDeviceGetHandleByIndex(visible_device_index())
         info = pynvml.nvmlDeviceGetMemoryInfo(handle)
         pynvml.nvmlShutdown()
-        return int(info.free)
+        return int(int(info.free) * _GPU_HEADROOM)
     except Exception:
         logger.warning(
             "GPU memory query failed (nvidia-ml-py not installed?); "

src/patchworks/_gpu.py

@@ -0,0 +1,158 @@
+"""Shared GPU helpers: device selection and surviving transient OOM.
+
+Cluster GPUs are shared. A co-tenant job's footprint can grow mid-run and push
+an otherwise-fine tile over the edge, so an out-of-memory error is often
+transient rather than a sign the work does not fit. Retrying on the GPU beats
+falling back to the CPU, where a single tile can take well over an hour.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+from typing import Any, Callable
+
+logger = logging.getLogger(__name__)
+
+OOM_RETRIES = 4
+OOM_BACKOFF_SECONDS = 30
+
+
+def is_oom(exc: BaseException) -> bool:
+    """Return whether *exc* is an out-of-memory error from any GPU stack.
+
+    Matching only ``RuntimeError`` missed cupy entirely:
+    ``cupy.cuda.memory.OutOfMemoryError`` derives from ``Exception``, not
+    ``RuntimeError``, so the DoG plugin and ``dilate_labels(use_gpu=True)``
+    got no retry at all. Class name is checked first because cupy's message
+    reads "Out of memory allocating ..." while torch's reads "CUDA out of
+    memory", and a future backend may word it differently again.
+
+    Parameters
+    ----------
+    exc : BaseException
+        The exception to classify.
+
+    Returns
+    -------
+    bool
+        True when the exception reports exhausted device memory.
+    """
+    if type(exc).__name__ in ("OutOfMemoryError", "CUDAOutOfMemoryError"):
+        return True
+    return "out of memory" in str(exc).lower()
+
+
+def free_gpu_caches() -> None:
+    """Release cached (but unused) device memory from torch and cupy.
+
+    Both keep their own allocator pool and neither releases it on its own, so
+    a process that has used either holds that memory for its lifetime. Freeing
+    both also defragments, which is often what actually lets a retry succeed.
+    Each is optional and skipped when not imported.
+    """
+    torch = __import__("sys").modules.get("torch")
+    if torch is not None and getattr(torch, "cuda", None) is not None:
+        try:
+            torch.cuda.empty_cache()
+        except Exception:  # pragma: no cover - defensive
+            logger.debug("torch.cuda.empty_cache() failed", exc_info=True)
+
+    cupy = __import__("sys").modules.get("cupy")
+    if cupy is not None:
+        try:
+            cupy.get_default_memory_pool().free_all_blocks()
+            cupy.get_default_pinned_memory_pool().free_all_blocks()
+        except Exception:  # pragma: no cover - defensive
+            logger.debug("cupy pool release failed", exc_info=True)
+
+
+def retry_on_oom(
+    call: Callable[[], Any],
+    *,
+    enabled: bool = True,
+    on_release: Callable[[], None] | None = None,
+    retries: int = OOM_RETRIES,
+    backoff: int = OOM_BACKOFF_SECONDS,
+) -> Any:
+    """Run *call*, retrying with a backoff while the GPU is out of memory.
+
+    Parameters
+    ----------
+    call : callable
+        Zero-argument callable doing the GPU work.
+    enabled : bool
+        Set False on a CPU path so real errors surface immediately.
+    on_release : callable, optional
+        Called before each backoff, to drop anything big this process is
+        holding on the device. Sleeping while still pinning a cached model is
+        self-defeating -- that memory is exactly what the co-tenant needs.
+    retries : int
+        Retries after the first attempt.
+    backoff : int
+        Base seconds; the wait grows linearly with the attempt number.
+
+    Returns
+    -------
+    Any
+        Whatever *call* returns.
+    """
+    for attempt in range(retries + 1):
+        try:
+            return call()
+        except Exception as exc:
+            if not enabled or not is_oom(exc):
+                raise
+            if on_release is not None:
+                on_release()
+            free_gpu_caches()
+            if attempt == retries:
+                logger.error(
+                    "GPU OOM persisted after %d retries; giving up.", retries
+                )
+                raise
+            wait = backoff * (attempt + 1)
+            logger.warning(
+                "GPU OOM (likely contention on a shared device); released "
+                "caches, retrying in %ds (attempt %d/%d).",
+                wait,
+                attempt + 1,
+                retries,
+            )
+            time.sleep(wait)
+
+
+def visible_device_index() -> int:
+    """NVML index of the device this process is actually allowed to use.
+
+    ``CUDA_VISIBLE_DEVICES`` remaps indices for the CUDA runtime but not for
+    NVML, which always enumerates every GPU on the node. Querying NVML index 0
+    unconditionally therefore reads a *different* GPU's free memory whenever
+    SLURM granted anything other than the first one -- and tile sizing was
+    built on that number.
+
+    Returns
+    -------
+    int
+        NVML device index, defaulting to 0 when nothing constrains us.
+    """
+    visible = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
+    if not visible:
+        return 0
+    first = visible.split(",")[0].strip()
+    if not first or first.startswith("GPU-") or first.startswith("MIG-"):
+        # A UUID form: NVML can resolve it directly, so leave index lookup
+        # alone and let the caller fall back.
+        return 0
+    try:
+        return max(0, int(first))
+    except ValueError:
+        return 0
+
+
+def visible_device_uuid() -> "str | None":
+    """UUID of the allowed device, when ``CUDA_VISIBLE_DEVICES`` uses one."""
+    visible = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
+    first = visible.split(",")[0].strip() if visible else ""
+    return first if first.startswith(("GPU-", "MIG-")) else None

src/patchworks/plugins/cellpose.py

@@ -21,16 +21,36 @@
 
 import numpy as np
 
+from .._gpu import retry_on_oom
+
 logger = logging.getLogger(__name__)
 
+
+def _parse_version(raw: str) -> tuple[int, ...]:
+    """Leading numeric components of a version string.
+
+    ``int(x)`` over the raw segments blows up on any suffixed release
+    (``"4.0rc1"``, ``"4.0.1.dev0"``) — and at *import* time, where only
+    ImportError was being caught, so the package became unimportable rather
+    than degrading.
+    """
+    parts = []
+    for segment in raw.split(".")[:2]:
+        digits = "".join(c for c in segment if c.isdigit())
+        if not digits:
+            break
+        parts.append(int(digits))
+    return tuple(parts) or (0,)
+
+
 try:
     from cellpose import models as _cellpose_models
 
-    _CELLPOSE_VERSION: tuple[int, ...] = tuple(
-        int(x) for x in importlib.metadata.version("cellpose").split(".")[:2]
+    _CELLPOSE_VERSION: tuple[int, ...] = _parse_version(
+        importlib.metadata.version("cellpose")
     )
     _CELLPOSE_V4 = _CELLPOSE_VERSION[0] >= 4
-except ImportError as _e:
+except ImportError:
     _cellpose_models = None  # type: ignore[assignment]
     _CELLPOSE_VERSION = (0, 0)
     _CELLPOSE_V4 = False
@@ -196,28 +216,23 @@ def _get_model(cellpose_dict: dict[str, Any]) -> Any:
     return _model_cache[key]
 
 
-def _is_cuda_oom(exc: Exception) -> bool:
-    """Return whether *exc* looks like a CUDA out-of-memory error.
-
-    Parameters
-    ----------
-    exc : Exception
-        Exception raised by ``model.eval``.
+def _drop_cached_models() -> None:
+    """Evict the per-process model cache so its VRAM can be reclaimed.
 
-    Returns
-    -------
-    bool
-        True if the exception is a CUDA OOM.
+    Waiting out an OOM while still holding a loaded model pinned on the device
+    is self-defeating: that memory is exactly what the contending job needs.
+    The model is reloaded from the (already downloaded) weights on the next
+    call, which costs seconds against a backoff measured in minutes.
     """
-    return "out of memory" in str(exc).lower()
-
-
-_OOM_RETRIES = 4
-_OOM_BACKOFF_SECONDS = 30
+    if _model_cache:
+        logger.info(
+            "releasing %d cached Cellpose model(s) before backing off",
+            len(_model_cache),
+        )
+        _model_cache.clear()
 
 
 def _eval_with_oom_fallback(
-    model: Any,
     img: np.ndarray,
     kwargs: dict[str, Any],
     cellpose_dict: dict[str, Any],
@@ -227,14 +242,13 @@ def _eval_with_oom_fallback(
     Cluster GPUs are often shared: another job's memory footprint can grow
     mid-run and push an otherwise-fine tile size over the edge. Staying on
     GPU matters — this tile can take well over an hour, so falling back to
-    CPU would be far worse than waiting. Instead this clears the allocator
-    cache (fixes fragmentation) and retries with a backoff, giving the
-    cotenant job time to free memory, before finally giving up.
+    CPU would be far worse than waiting. See :func:`patchworks._gpu.retry_on_oom`.
+
+    The model is fetched inside the retry, so an eviction during backoff is
+    followed by a reload rather than a use-after-free of the cached handle.
 
     Parameters
     ----------
-    model : Any
-        Cellpose model instance.
     img : np.ndarray
         Image to segment.
     kwargs : dict
@@ -247,32 +261,11 @@ def _eval_with_oom_fallback(
     np.ndarray
         Label array from ``model.eval``.
     """
-    import time
-
-    for attempt in range(_OOM_RETRIES + 1):
-        try:
-            return model.eval(img, **kwargs)[0]
-        except RuntimeError as exc:
-            if not cellpose_dict.get("gpu", False) or not _is_cuda_oom(exc):
-                raise
-            import torch
-
-            torch.cuda.empty_cache()
-            if attempt == _OOM_RETRIES:
-                logger.error(
-                    "CUDA OOM persisted after %d retries; giving up on tile.",
-                    _OOM_RETRIES,
-                )
-                raise
-            wait = _OOM_BACKOFF_SECONDS * (attempt + 1)
-            logger.warning(
-                "CUDA OOM on tile (likely GPU contention); cleared cache, "
-                "retrying in %ds (attempt %d/%d).",
-                wait,
-                attempt + 1,
-                _OOM_RETRIES,
-            )
-            time.sleep(wait)
+    return retry_on_oom(
+        lambda: _get_model(cellpose_dict).eval(img, **kwargs)[0],
+        enabled=bool(cellpose_dict.get("gpu", False)),
+        on_release=_drop_cached_models,
+    )
 
 
 def _run(block: np.ndarray, cellpose_dict: dict[str, Any]) -> np.ndarray:
@@ -290,7 +283,6 @@ def _run(block: np.ndarray, cellpose_dict: dict[str, Any]) -> np.ndarray:
     np.ndarray
         Integer (``int32``) label array of the same spatial shape.
     """
-    model = _get_model(cellpose_dict)
     do_3D = cellpose_dict["do_3D"]
 
     if _CELLPOSE_V4:
@@ -310,13 +302,20 @@ def _run(block: np.ndarray, cellpose_dict: dict[str, Any]) -> np.ndarray:
 
     if do_3D:
         kwargs["z_axis"] = 0
-        masks = _eval_with_oom_fallback(model, block, kwargs, cellpose_dict)
+        masks = _eval_with_oom_fallback(block, kwargs, cellpose_dict)
         return masks.astype("int32")
     else:
         # Squeeze singleton z so Cellpose gets a clean 2-D image
         squeeze = block.ndim == 3 and block.shape[0] == 1
+        if block.ndim == 3 and not squeeze:
+            raise ValueError(
+                f"do_3D is False but this tile has {block.shape[0]} z-planes. "
+                "Cellpose would receive the stack with no z_axis and treat "
+                "the leading axis as channels. Set do_3D: true, or tile with "
+                "z=1 to segment plane by plane."
+            )
         img = block[0] if squeeze else block
-        masks = _eval_with_oom_fallback(model, img, kwargs, cellpose_dict)
+        masks = _eval_with_oom_fallback(img, kwargs, cellpose_dict)
         masks = masks.astype("int32")
         return masks[np.newaxis] if squeeze else masks