patchworks 2.6.1__tar.gz → 2.6.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: patchworks
-Version: 2.6.1
+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.1 → 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.1 → 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×
@@ -323,13 +332,33 @@ follow from that:
 
 - 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 budget. A **hand-set** `tile_shape`
-  sized to fill a GPU has no such protection and needs halving yourself.
+  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.1 → 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.1 → patchworks-2.6.2}/src/patchworks/_chunks.py

@@ -281,11 +281,15 @@ 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:
@@ -316,7 +320,13 @@ def auto_tile_shape(
     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()
@@ -395,13 +405,23 @@ 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
@@ -429,7 +449,13 @@ def auto_tile_shape_cellpose(
     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.1 → 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.1 → patchworks-2.6.2}/tests/test_core.py

@@ -410,7 +410,15 @@ def test_auto_tile_shape_charges_for_extra_channels():
     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)
+    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)

{patchworks-2.6.1 → 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.1 → 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.1 → 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.1 → 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