patchworks 1.4.2__tar.gz → 1.4.4__tar.gz

{patchworks-1.4.2 → patchworks-1.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: patchworks
-Version: 1.4.2
+Version: 1.4.4
 Summary: Tiled processing of arbitrarily large images with globally consistent labels
 Project-URL: Homepage, https://github.com/imcf/patchworks
 Project-URL: Issues, https://github.com/imcf/patchworks/issues

{patchworks-1.4.2 → patchworks-1.4.4}/src/patchworks/plugins/cellpose.py

@@ -196,6 +196,82 @@ 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``.
+
+    Returns
+    -------
+    bool
+        True if the exception is a CUDA OOM.
+    """
+    return "out of memory" in str(exc).lower()
+
+
+_OOM_RETRIES = 4
+_OOM_BACKOFF_SECONDS = 30
+
+
+def _eval_with_oom_fallback(
+    model: Any, img: np.ndarray, kwargs: dict[str, Any], cellpose_dict: dict[str, Any]
+) -> np.ndarray:
+    """Run ``model.eval``, surviving transient GPU contention.
+
+    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.
+
+    Parameters
+    ----------
+    model : Any
+        Cellpose model instance.
+    img : np.ndarray
+        Image to segment.
+    kwargs : dict
+        Keyword arguments for ``model.eval``.
+    cellpose_dict : dict
+        Configuration from :func:`_make_config`.
+
+    Returns
+    -------
+    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)
+
+
 def _run(block: np.ndarray, cellpose_dict: dict[str, Any]) -> np.ndarray:
     """Segment one tile with a cached Cellpose model.
 
@@ -231,12 +307,15 @@ def _run(block: np.ndarray, cellpose_dict: dict[str, Any]) -> np.ndarray:
 
     if do_3D:
         kwargs["z_axis"] = 0
-        return model.eval(block, **kwargs)[0].astype("int32")
+        masks = _eval_with_oom_fallback(model, 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
         img = block[0] if squeeze else block
-        masks = model.eval(img, **kwargs)[0].astype("int32")
+        masks = _eval_with_oom_fallback(model, img, kwargs, cellpose_dict).astype(
+            "int32"
+        )
         return masks[np.newaxis] if squeeze else masks
 
 

{patchworks-1.4.2 → patchworks-1.4.4}/src/patchworks/plugins/ome_zarr.py

@@ -727,19 +727,23 @@ def _open_tiff_sequence(
     ts = tifffile.TiffSequence(pattern, pattern=sequence_pattern)
 
     # Each file should be a single 2D (y, x) plane, but tifffile.imread can
-    # report extra leading singleton dims (e.g. a page/sample axis) even for
-    # a plain single-page TIFF — drop those so one file maps to exactly (y, x).
+    # report extra singleton dims — leading, trailing, or both (e.g. a stray
+    # page or samples-per-pixel axis) — even for a plain single-page TIFF.
+    # Drop whichever ones are singleton so one file maps to exactly (y, x).
     chunkshape = tifffile.imread(ts[0]).shape
-    lead = len(chunkshape) - 2
-    if lead < 0 or any(s != 1 for s in chunkshape[:lead]):
+    keep_chunk = [i for i, s in enumerate(chunkshape) if s != 1]
+    if len(keep_chunk) != 2:
         raise ValueError(
             f"expected single 2D (y, x) planes per file, got shape "
             f"{chunkshape} for {ts[0]!r}"
         )
 
     arr = da.from_zarr(zarr.open(store=ts.aszarr()))
-    if lead:
-        arr = arr[(slice(None),) * len(ts.axes) + (0,) * lead]
+    n_seq = len(ts.axes)
+    index = tuple(
+        slice(None) if i in keep_chunk else 0 for i in range(len(chunkshape))
+    )
+    arr = arr[(slice(None),) * n_seq + index]
     full_axes = ts.axes.lower() + "yx"
 
     # Canonicalise to patchworks' tczyx axis order regardless of the order