patchworks 2.6.0__tar.gz → 2.6.2__tar.gz

{patchworks-2.6.0 → patchworks-2.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: patchworks
-Version: 2.6.0
+Version: 2.6.2
 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-2.6.0 → patchworks-2.6.2}/docs/guide/gpu_distributed.md

@@ -79,6 +79,11 @@ own outright. `auto_tile_shape` then sizes each tile to at most half of that
 budget; `auto_tile_shape_cellpose` uses Cellpose's own memory model instead
 (roughly 20× the raw tile bytes, plus ~2 GiB for the model).
 
+"Raw tile bytes" counts every channel a tile carries, so pass `n_channels=2`
+when feeding Cellpose a cyto+nuclei pair — see
+[Multi-channel tiles](tiling.md#multi-channel-tiles). The Snakemake workflow
+does this for you whenever `nuclei_channel` is set.
+
 The device is resolved from `CUDA_VISIBLE_DEVICES`. This matters on
 multi-GPU nodes: NVML enumerates **every** GPU regardless of `--gres=gpu:1`,
 so querying index 0 unconditionally would read a different card's free memory

{patchworks-2.6.0 → patchworks-2.6.2}/docs/guide/snakemake.md

@@ -109,6 +109,15 @@ sequential_labels: true        # renumber labels to a contiguous 1..N
     explicit z (like `[16, 1024, 1024]`) tiles in z instead. `prepare` logs
     which regime it picked.
 
+    `"auto"` also caps the tile to the **host** RAM available to the job, not
+    just VRAM — a `do_3D` tile that comfortably fits a big GPU can still be
+    too large for the SLURM/cgroup memory the job was actually granted, and
+    that shows up as a plain `SIGKILL`, not a catchable CUDA-OOM error.
+    Because `prepare` runs on a CPU node, it checks *its own* grant as a
+    stand-in for `segment`'s — keep `prepare`'s and `segment`'s `mem_mb` in
+    `profile/slurm/config.yaml` equal, or the estimate is sized against the
+    wrong job's budget.
+
 !!! tip "Use a per-axis `overlap`"
     A scalar halo is applied to every axis. On a `[16, 1024, 1024]` tile,
     `overlap: 30` reads `76 × 1084 × 1084` to keep `16 × 1024 × 1024` — 5.3×
@@ -321,14 +330,35 @@ occupancy map and the staged labels are byte-for-byte what a single-channel
 run produces, and `merge` and `label_relations` need no changes. Two things
 follow from that:
 
-- A tile holds twice the bytes, so a hand-set `tile_shape` sized to fill a GPU
-  may need halving. `tile_shape: "auto"` sizes from the single-channel array
-  and does not yet know about the pair.
+- A tile holds twice the bytes. `tile_shape: "auto"` accounts for this — it is
+  told the tile carries two channels and shrinks each spatial side by ~1/√2,
+  so the tile still fits the same VRAM *and* host-RAM budget (see the "Tile
+  size vs runtime" tip above). A **hand-set** `tile_shape` sized to fill a
+  GPU has no such protection and needs halving yourself.
 - The translation is version-specific. Cellpose 3 gets `channels: [1, 2]`
   (1-based into the channel axis, `0` = grayscale); Cellpose 4 (cpsam) dropped
   `channels` entirely and simply reads both. Either is overridable by setting
   `channels:` or `channel_axis:` in the `cellpose:` block.
 
+!!! warning "In a `multi.yaml` run, pin `tile_shape` explicitly"
+
+    `label_relations` requires its two label arrays to share a chunk layout,
+    and that layout comes from `tile_shape`. Giving *one* config a
+    `nuclei_channel` while the group uses `tile_shape: "auto"` produces a
+    **smaller** tile for that config only — so the label groups end up chunked
+    differently and the relations step fails, after every segmentation has
+    already run.
+
+    `run_multi`'s cross-config check compares the configured values, and
+    `"auto" == "auto"`, so it flags this case specifically. Fix it by setting
+    one explicit `tile_shape` in `common.yaml`, sized for the two-channel
+    config (roughly each spatial side ÷ √2 versus what you would use for a
+    single channel), so every config shares it.
+
+If you do not need that segmentation related to the others, the alternative is
+to run it on its own against the same `work_dir` and leave it out of
+`multi.yaml`.
+
 `nuclei_channel` applies to the SLURM/Snakemake path. The single-process
 `tile_process` API still takes one `channel`.
 

{patchworks-2.6.0 → patchworks-2.6.2}/docs/guide/tiling.md

@@ -57,6 +57,23 @@ tile_process("image.zarr", fn, tile_shape=tile_fn)
 The callable is called with `(shape, dtype)` at runtime, after the image is
 loaded — useful when you don't know the image shape in advance.
 
+### Multi-channel tiles
+
+Both sizers take `n_channels` (default 1) and charge it per voxel. Use it when
+a tile carries more than one channel while its *geometry* stays spatial — the
+case the Snakemake workflow's
+[`nuclei_channel`](snakemake.md#giving-cellpose-a-nuclei-channel) creates, where
+Cellpose is handed a cyto+nuclei pair but still returns one label per voxel:
+
+```python
+# Same VRAM budget, so each spatial side shrinks by ~1/√2
+auto_tile_shape_cellpose(shape, dtype, diameter=30, use_gpu=True, n_channels=2)
+```
+
+Leaving it at 1 for a 2-channel tile budgets for half the bytes the tile
+actually holds, which surfaces as an out-of-memory error in the first tile
+rather than as anything about channels.
+
 ## Overlap
 
 Methods that need spatial context (Cellpose, StarDist, U-Net) produce wrong

{patchworks-2.6.0 → patchworks-2.6.2}/src/patchworks/_chunks.py

@@ -264,6 +264,7 @@ def auto_tile_shape(
     gpu_memory: int | None = None,
     available_memory: int | None = None,
     n_workers: int | None = None,
+    n_channels: int = 1,
     verbose: bool = False,
 ) -> tuple[int, ...]:
     """Balanced tile shape for general-purpose 3-D processing.
@@ -280,13 +281,21 @@ def auto_tile_shape(
     target_bytes:
         Memory ceiling per tile. Default 64 MiB.
     use_gpu:
-        Size tiles against GPU VRAM rather than host RAM.
+        Size tiles against GPU VRAM rather than host RAM. Still capped by
+        host RAM too -- see ``available_memory``.
     gpu_memory:
         Available GPU VRAM in bytes; auto-queried when None.
     available_memory:
-        Available host RAM in bytes; auto-queried when None.
+        Available host RAM in bytes; auto-queried when None. Consulted even
+        when ``use_gpu`` is set: the tile has to be loaded into host memory
+        before (and often after) it reaches the GPU, so it must fit both
+        budgets, not just VRAM.
     n_workers:
         Number of parallel workers (divides the RAM budget).
+    n_channels:
+        Channels each tile carries (default 1). Above 1 the per-voxel cost
+        scales with it, so the tile shrinks accordingly -- e.g. the workflow's
+        ``nuclei_channel`` hands Cellpose a cyto+nuclei pair.
     verbose:
         Log the chosen shape and estimated tile size.
 
@@ -302,11 +311,22 @@ def auto_tile_shape(
     (128, 512, 512)
     """
     n_workers = n_workers or cpu_allocation()
-    itemsize = np.dtype(dtype).itemsize
+    # A tile holds n_channels planes per voxel (e.g. Cellpose's
+    # cyto+nuclei pair), so the per-voxel cost -- and every budget
+    # derived from it below -- scales with them.
+    if n_channels < 1:
+        raise ValueError(f"n_channels must be >= 1; got {n_channels!r}")
+    itemsize = np.dtype(dtype).itemsize * n_channels
     n_spatial = min(3, len(shape))
 
     if use_gpu:
-        mem = gpu_memory if gpu_memory is not None else _get_gpu_memory()
+        # A tile that fits VRAM still has to be decoded into (and often
+        # copied back out of) host RAM first, so a GPU with more memory
+        # than the job's own host allocation must not produce a tile the
+        # job itself can't hold -- take whichever budget is tighter.
+        gpu_mem = gpu_memory if gpu_memory is not None else _get_gpu_memory()
+        host_mem = available_memory or _get_available_memory()
+        mem = min(gpu_mem, host_mem)
         budget = min(target_bytes * 2, mem // 2)
     else:
         mem = available_memory or _get_available_memory()
@@ -359,6 +379,7 @@ def auto_tile_shape_cellpose(
     n_workers: int | None = None,
     model_memory_bytes: int = 2 * 1024**3,
     cellpose_memory_factor: int = 20,
+    n_channels: int = 1,
     verbose: bool = False,
 ) -> tuple[int, ...]:
     """Cellpose-optimised tile shape.
@@ -384,13 +405,27 @@ def auto_tile_shape_cellpose(
     do_3D:
         Whether Cellpose will run in 3-D mode.
     use_gpu:
-        Size tiles for GPU VRAM.
+        Size tiles for GPU VRAM. Still capped by host RAM too -- see
+        ``available_memory``.
     gpu_memory, available_memory, n_workers:
-        Memory parameters (auto-queried when None).
+        Memory parameters (auto-queried when None). ``available_memory`` is
+        consulted even when ``use_gpu`` is set: Cellpose's ``do_3D``
+        preprocessing (orthogonal-view flows, mask assembly) copies the tile
+        through host memory, so VRAM headroom alone doesn't guarantee the
+        job's own host allocation can hold it.
     model_memory_bytes:
         Memory consumed by the Cellpose model weights (default 2 GiB).
     cellpose_memory_factor:
-        Cellpose allocates roughly this multiple of raw input bytes (default 20×).
+        Cellpose allocates roughly this multiple of raw input bytes (default
+        20×). Applied against whichever of VRAM/host RAM is tighter --
+        ``ponytail:`` this reuses one factor for both budgets as an
+        approximation; if it proves off in practice, calibrate a separate
+        host-side factor from a SLURM job's peak RSS (``seff``/
+        ``/usr/bin/time -v``) and pass it alongside this one.
+    n_channels:
+        Channels each tile carries (default 1). Above 1 the per-voxel cost
+        scales with it, so the tile shrinks accordingly -- e.g. the workflow's
+        ``nuclei_channel`` hands Cellpose a cyto+nuclei pair.
     verbose:
         Log the chosen shape and memory estimates.
 
@@ -406,10 +441,21 @@ def auto_tile_shape_cellpose(
     (1, 2048, 2048)
     """
     n_workers = n_workers or cpu_allocation()
-    itemsize = np.dtype(dtype).itemsize
+    # A tile holds n_channels planes per voxel (e.g. Cellpose's
+    # cyto+nuclei pair), so the per-voxel cost -- and every budget
+    # derived from it below -- scales with them.
+    if n_channels < 1:
+        raise ValueError(f"n_channels must be >= 1; got {n_channels!r}")
+    itemsize = np.dtype(dtype).itemsize * n_channels
 
     if use_gpu:
-        total_mem = gpu_memory if gpu_memory is not None else _get_gpu_memory()
+        # A tile that fits VRAM still has to be decoded into (and often
+        # copied back out of) host RAM first, so a GPU with more memory
+        # than the job's own host allocation must not produce a tile the
+        # job itself can't hold -- take whichever budget is tighter.
+        gpu_mem = gpu_memory if gpu_memory is not None else _get_gpu_memory()
+        host_mem = available_memory or _get_available_memory()
+        total_mem = min(gpu_mem, host_mem)
     else:
         total_mem = (available_memory or _get_available_memory()) // n_workers
 

{patchworks-2.6.0 → patchworks-2.6.2}/tests/test_allocation.py

@@ -8,7 +8,11 @@ checks that would have caught it.
 
 import numpy as np
 
-from patchworks import cpu_allocation, safe_worker_count
+from patchworks import (
+    auto_tile_shape_cellpose,
+    cpu_allocation,
+    safe_worker_count,
+)
 from patchworks._chunks import _get_available_memory
 
 GIB = 1024**3
@@ -76,3 +80,40 @@ def test_worker_count_is_bounded_by_the_allocation(monkeypatch):
     n = safe_worker_count(chunk_nbytes * 40, fn_overhead=3)
     assert n < 32, "must not size itself to the core count when RAM is tighter"
     assert n >= 1
+
+
+def test_gpu_tile_sizing_is_bounded_by_the_host_allocation(monkeypatch):
+    """A big GPU must not excuse a tile the job's own host RAM can't hold.
+
+    This is the concrete failure a `do_3D` nuclei segmentation hit: an ample
+    GPU (24 GiB) sized the tile against VRAM alone, and the job -- granted
+    only 1 GiB of host RAM here -- was SIGKILLed loading it, unrelated to
+    ``nuclei_channel``. The sizer must take whichever budget is tighter.
+    """
+    monkeypatch.delenv("SLURM_MEM_PER_CPU", raising=False)
+    monkeypatch.setenv("SLURM_MEM_PER_NODE", str(1024))  # 1 GiB
+    monkeypatch.setattr(
+        "patchworks._chunks._cgroup_memory_limit", lambda: 512 * GIB
+    )
+
+    tile = auto_tile_shape_cellpose(
+        (128, 2048, 2048),
+        "uint16",
+        diameter=30,
+        do_3D=True,
+        use_gpu=True,
+        gpu_memory=24 * GIB,
+    )
+    generous = auto_tile_shape_cellpose(
+        (128, 2048, 2048),
+        "uint16",
+        diameter=30,
+        do_3D=True,
+        use_gpu=True,
+        gpu_memory=24 * GIB,
+        available_memory=64 * GIB,
+    )
+    assert np.prod(tile) < np.prod(generous), (
+        "the 1 GiB host grant must shrink the tile below what the same "
+        "24 GiB GPU would otherwise allow"
+    )

{patchworks-2.6.0 → patchworks-2.6.2}/tests/test_core.py

@@ -389,3 +389,40 @@ def test_channel_selection_respects_the_stores_axes(tmp_path):
     store4 = str(tmp_path / "czyx.zarr")
     to_ome_zarr(vol4, store4, axes="czyx", n_levels=1, progress=False)
     assert load_ome_zarr(store4, channel=1, level=0).shape == (8, 64, 64)
+
+
+def test_auto_tile_shape_charges_for_extra_channels():
+    """A 2-channel tile must fit the same byte budget, not twice it.
+
+    `nuclei_channel` doubles what a tile holds while the tile geometry stays
+    single-channel, so a sizer blind to it hands the GPU a tile needing twice
+    the VRAM it budgeted for.
+    """
+    import pytest
+
+    from patchworks import auto_tile_shape, auto_tile_shape_cellpose
+
+    shape, dtype = (128, 2048, 2048), "uint16"
+
+    one = auto_tile_shape(shape, dtype)
+    two = auto_tile_shape(shape, dtype, n_channels=2)
+    # Same bytes overall: 2 channels of roughly half the area each.
+    assert np.prod(two) * 2 <= np.prod(one)
+    assert np.prod(two) * 2 >= np.prod(one) * 0.9
+
+    kw = dict(
+        diameter=30,
+        do_3D=True,
+        use_gpu=True,
+        gpu_memory=24 * 1024**3,
+        # Generous on purpose: this test is about GPU-vs-channel scaling,
+        # not the host-RAM ceiling, so host RAM must stay non-binding here.
+        available_memory=64 * 1024**3,
+    )
+    cp_one = auto_tile_shape_cellpose(shape, dtype, **kw)
+    cp_two = auto_tile_shape_cellpose(shape, dtype, n_channels=2, **kw)
+    assert np.prod(cp_two) * 2 <= np.prod(cp_one)
+    assert cp_two[0] == cp_one[0]  # do_3D still pins z to the full extent
+
+    with pytest.raises(ValueError, match="n_channels"):
+        auto_tile_shape(shape, dtype, n_channels=0)

{patchworks-2.6.0 → patchworks-2.6.2}/tests/test_run_multi.py

@@ -158,3 +158,30 @@ def test_occupancy_is_not_rebuilt_by_the_driver():
     src = (_workflow_dir() / "scripts" / "run_multi.py").read_text()
     assert "build_occupancy_map(" not in src
     assert "occupancy.zarr" in src
+
+
+def test_auto_tile_shape_with_a_lone_nuclei_channel_is_refused():
+    """Matching `tile_shape` *values* are not enough when one config is 2-ch.
+
+    "auto" == "auto" passes the plain equality check, but the sizer charges
+    per channel, so the nuclei_channel config gets a smaller tile. The label
+    groups then disagree on chunk layout and label_relations raises -- after
+    every segmentation has already run, which is the expensive way to find out.
+    """
+    paths = [Path("a.yaml"), Path("b.yaml")]
+    base = {"work_dir": "/w", "tile_shape": "auto", "level": 0}
+
+    bad = [
+        {**base, "label_name": "a", "channel": 0, "nuclei_channel": 1},
+        {**base, "label_name": "b", "channel": 2},
+    ]
+    with pytest.raises(SystemExit):
+        _validate_configs(paths, bad)
+
+    # Same pair with one explicit shape is fine: both get that tile.
+    pinned = [{**c, "tile_shape": [16, 512, 512]} for c in bad]
+    assert _validate_configs(paths, pinned) == "/w"
+
+    # And "auto" is fine when every config carries the same channel count.
+    both = [{**bad[0]}, {**bad[1], "nuclei_channel": 3}]
+    assert _validate_configs(paths, both) == "/w"

{patchworks-2.6.0 → patchworks-2.6.2}/workflow/README.md

@@ -68,6 +68,11 @@ Edit `config/config.yaml` (input, output dir, channel, tile shape, Cellpose
 model/diameter/`do_3D`, …) and `profile/slurm/config.yaml` (partitions,
 account, GPU request).
 
+Channel indices are **0-based**. Set `nuclei_channel` alongside `channel` to
+hand Cellpose the nuclear stain as a second input, which usually improves
+cytoplasm segmentation — see
+[Giving Cellpose a nuclei channel](https://imcf.one/patchworks/guide/snakemake/#giving-cellpose-a-nuclei-channel).
+
 ## Run
 
 ```bash

{patchworks-2.6.0 → patchworks-2.6.2}/workflow/config/config_cyto.yaml

@@ -14,9 +14,11 @@ channel: 0 # cytoplasm/membrane channel
 #
 # Only the segment step reads it: the tile geometry, the occupancy map and the
 # staged labels all stay exactly as they are for a single channel -- the pair
-# is stacked on a leading axis that is carried into each tile, not tiled. It
-# does double the bytes a tile holds, so a hand-set `tile_shape` sized to the
-# GPU may need halving.
+# is stacked on a leading axis that is carried into each tile, not tiled.
+#
+# It does double the bytes a tile holds. `tile_shape: "auto"` is told about
+# that and shrinks each spatial side by ~1/sqrt(2) to fit the same VRAM
+# budget; the hand-set tile_shape in common.yaml is not, so halve it yourself.
 #
 # Cellpose 3 turns this into channels: [1, 2]; Cellpose 4 (cpsam) dropped
 # `channels` and just reads both. Override either by setting `channels:` or

{patchworks-2.6.0 → patchworks-2.6.2}/workflow/profile/slurm/config.yaml

@@ -67,6 +67,13 @@ set-resources:
   prepare:
     # Cheap now: it reduces the already-built occupancy map over each tile's
     # footprint instead of touching the image.
+    #
+    # tile_shape: "auto" now also sizes against a host-RAM budget, and it
+    # reads *this job's* grant as a stand-in for segment's -- prepare runs on
+    # a CPU node and can't query segment's GPU-node grant directly. Keep this
+    # equal to segment's mem_mb below: lower and tiles come out needlessly
+    # small; higher and the sizer budgets against more host RAM than segment
+    # will actually get, undoing the point of the host-RAM check.
     mem_mb: "attempt * 32000"
     runtime: 120
   segment:
@@ -80,7 +87,8 @@ set-resources:
     # A job now processes `tiles_per_job` tiles sequentially, so both memory
     # and runtime scale with that setting — raise it there and re-check here.
     # The old "a tile used ~1G" note predates tile_shape: "auto", which sizes
-    # tiles against the real GPU and makes them far bigger.
+    # tiles against the real GPU and makes them far bigger. If you raise
+    # this, raise prepare's mem_mb above to match (see its comment).
     mem_mb: "attempt * 32000"
     cpus_per_task: 4
     runtime: 360 # 6 hours — must match the QOS, NOT 120 (=2h → killed early)

{patchworks-2.6.0 → patchworks-2.6.2}/workflow/scripts/prepare_tiles.py

@@ -44,6 +44,10 @@ if ts == "auto":
     # "GPU memory query failed" fallback). None => the built-in 8 GiB default.
     gpu_gb = cfg.get("gpu_memory_gb")
     gpu_bytes = int(gpu_gb * 1024**3) if gpu_gb else None
+    # `image` is single-channel here (the geometry is spatial), but segment
+    # reads nuclei_channel alongside it, so a tile costs twice the bytes.
+    # Without this the sizer would hand the GPU a tile it cannot hold.
+    n_channels = 2 if cfg.get("nuclei_channel") is not None else 1
     if method == "cellpose":
         cp = cfg["cellpose"]
         sizer = partial(
@@ -52,6 +56,7 @@ if ts == "auto":
             use_gpu=cp.get("gpu", True),
             diameter=cp.get("diameter"),
             gpu_memory=gpu_bytes,
+            n_channels=n_channels,
         )
     else:
         # cfg["cellpose"] used to be read unconditionally here, so a DoG or
@@ -59,7 +64,10 @@ if ts == "auto":
         # 'cellpose'. The Cellpose estimator's memory model wouldn't apply to
         # them anyway.
         sizer = partial(
-            auto_tile_shape, use_gpu=gpu_bytes is not None, gpu_memory=gpu_bytes
+            auto_tile_shape,
+            use_gpu=gpu_bytes is not None,
+            gpu_memory=gpu_bytes,
+            n_channels=n_channels,
         )
     tile_shape = tuple(sizer(image.shape, image.dtype))
 else:

{patchworks-2.6.0 → patchworks-2.6.2}/workflow/scripts/run_multi.py

@@ -228,6 +228,24 @@ def _validate_configs(paths: list[Path], cfgs: list[dict]) -> str:
                 f"share a chunk layout; got {_spread(key)}"
             )
 
+    # `tile_shape: "auto"` is identical as a *value* across configs while
+    # producing different tiles: the sizer charges per channel, so a config
+    # with nuclei_channel gets a smaller one. The label groups then disagree on
+    # chunk layout and label_relations raises -- after every segmentation has
+    # run. Matching values are not enough here, so check the inputs that feed
+    # the sizer instead.
+    if {repr(cfg.get("tile_shape", "auto")) for cfg in cfgs} == {repr("auto")}:
+        if len({cfg.get("nuclei_channel") is not None for cfg in cfgs}) != 1:
+            problems.append(
+                'tile_shape: "auto" sizes a nuclei_channel config smaller '
+                "(a tile carries two channels), so the label groups would "
+                "not share a chunk layout and label_relations would fail "
+                f"after every segmentation had run; got "
+                f"{_spread('nuclei_channel')}. Set one explicit tile_shape in "
+                "the file `common:` points at, sized for the two-channel "
+                "config."
+            )
+
     # Phase A converts once, from the first config. Anything `convert` reads
     # out of a later config is therefore silently ignored -- someone setting
     # `shard: true` on the second config and watching a million files appear