patchworks 2.6.1__tar.gz → 2.6.3__tar.gz

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: patchworks
-Version: 2.6.1
+Version: 2.6.3
 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.3}/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.3}/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.3}/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.3}/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.3}/src/patchworks/_notify.py

@@ -159,7 +159,7 @@ def slurm_mail_extra(
 
 
 def failing_step(
-    snakemake_log: Union[str, Path, None],
+    snakemake_log: "Union[str, Path, list, tuple, None]",
 ) -> "tuple[Union[str, None], Union[str, None]]":
     """Find which rule failed, and its log, from Snakemake's own log file.
 
@@ -171,9 +171,12 @@ def failing_step(
 
     Parameters
     ----------
-    snakemake_log : str or Path or None
-        Path to Snakemake's own log (the ``log`` variable inside an
-        ``onerror`` handler).
+    snakemake_log : str, Path, list, tuple, or None
+        Snakemake's own log (the ``log`` variable inside an ``onerror``
+        handler). Snakemake 8+ passes this as a list of paths (its
+        ``LoggerManager.get_logfile()`` returns ``List[str]``, even though
+        there is normally just one) rather than a single string -- take the
+        first entry.
 
     Returns
     -------
@@ -181,6 +184,8 @@ def failing_step(
         ``(rule_name, log_path)``, either of which may be None when the log
         is unreadable or records no rule error.
     """
+    if isinstance(snakemake_log, (list, tuple)):
+        snakemake_log = snakemake_log[0] if snakemake_log else None
     try:
         text = Path(snakemake_log).read_text(errors="replace")
     except (OSError, TypeError, ValueError):

{patchworks-2.6.1 → patchworks-2.6.3}/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.3}/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.3}/tests/test_notify.py

@@ -99,6 +99,30 @@ def test_failing_step_reads_the_rule_from_snakemakes_log(tmp_path):
     )
 
 
+def test_failing_step_accepts_snakemakes_list_log(tmp_path):
+    """Snakemake 8+ passes `log` as a list, not a bare path.
+
+    ``LoggerManager.get_logfile()`` returns ``List[str]``, and that's what
+    the `log` variable inside `onerror:` actually is. `Path(a_list)` raises
+    `TypeError`, which the old code caught and turned into `(None, None)` --
+    so this path was *always* silently falling back to the mtime guess it
+    was written to replace, on every real failure.
+    """
+    from patchworks._notify import failing_step
+
+    log = tmp_path / "sm.log"
+    log.write_text(
+        "Error in rule segment:\n"
+        "    jobid: 69\n"
+        "    output: seg/61.done\n"
+        "    log: /w/nuclei_labels/logs/segment/61.log (check log file(s))\n"
+    )
+    expected = ("segment", "/w/nuclei_labels/logs/segment/61.log")
+    assert failing_step([str(log)]) == expected
+    assert failing_step((str(log),)) == expected
+    assert failing_step([]) == (None, None)
+
+
 def test_failing_step_degrades_quietly(tmp_path):
     """It runs inside an error handler, so it must never raise itself."""
     from patchworks._notify import failing_step

{patchworks-2.6.1 → patchworks-2.6.3}/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.3}/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.3}/workflow/profile/slurm/config.yaml

@@ -67,7 +67,14 @@ set-resources:
   prepare:
     # Cheap now: it reduces the already-built occupancy map over each tile's
     # footprint instead of touching the image.
-    mem_mb: "attempt * 32000"
+    #
+    # 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 * 128000"
     runtime: 120
   segment:
     # one GPU per job — this is what spreads Cellpose across GPUs.
@@ -76,14 +83,35 @@ set-resources:
     # leaves the job without a device.)
     slurm_partition: "rtx4090"
     gres: "gpu:1"
-    qos: "rtx4090-6hours" # scicore: <partition>-<duration> QOS
+    # rtx4090-6hours wasn't enough: a do_3D whole-z tile (126 planes, full
+    # footprint) can run past 6h even though most same-sized tiles finish in
+    # 2.5-3.5h -- retrying it bought nothing since `runtime` didn't scale with
+    # `attempt` the way `mem_mb` does below, so it hit the identical wall on
+    # every attempt and burned all `retries` for a job that was never going to
+    # finish in the window. rtx4090-1day is scicore's next tier up (see
+    # `sacctmgr show qos format=name,maxwall | grep rtx4090`); 24h is wide
+    # margin over anything observed so far. If a tile still times out at 24h,
+    # that's not a slow tile anymore -- it's stuck, and worth profiling on its
+    # own rather than reaching for rtx4090-1week.
+    qos: "rtx4090-1day" # scicore: <partition>-<duration> QOS
     # 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.
-    mem_mb: "attempt * 32000"
+    # 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).
+    #
+    # 128 GB base (256/384 GB on retry) is a wide safety margin, not a tight
+    # estimate: a do_3D tile that `auto_tile_shape_cellpose`'s 20x
+    # cellpose_memory_factor judged safe within a 24 GiB GPU budget still hit
+    # a real 32 GB host OOM, so that heuristic underestimates do_3D's actual
+    # host-RAM use by more than expected. rtx4090 nodes have ~1 TB RAM
+    # (~800 GB usable) and the QOS caps at 4 TB account-wide, so this has
+    # plenty of room -- it buys time until do_3D gets a properly measured
+    # memory factor (e.g. from `seff` on a job that completes) instead of a
+    # guess. ``ponytail:`` tighten this once real peak-RSS numbers exist.
+    mem_mb: "attempt * 128000"
     cpus_per_task: 4
-    runtime: 360 # 6 hours — must match the QOS, NOT 120 (=2h → killed early)
+    runtime: 1400 # ~23.3h — must stay under the rtx4090-1day QOS cap (24h)
   merge:
     # merge.py sizes its worker pool from the cgroup/SLURM budget rather than
     # the node's core count, and the pyramid is written zarr-natively (one

{patchworks-2.6.1 → patchworks-2.6.3}/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