Merge dev: pipeline overhaul (v2.0.0)

Cuts GPU time roughly 3x, makes tile skipping exact rather than sampled, and
removes the merge OOM by construction rather than by tuning.

Performance
- per-axis overlap: a scalar halo on a [16,1024,1024] tile read 5.3x the
  voxels it kept, nearly all wasted z; [4,30,30] brings that to ~1.7x
- tiles_per_job batches tiles into one SLURM job, so CUDA init and the
  Cellpose weight load are paid once per batch, not once per tile
- six full passes over the label volume become two: per-tile counts replace
  the global-uniqueness rewrite, the sequential renumber folds into the
  merge's LUT, and the merge writes into the label group and relabels in place
- pyramid levels stream one source chunk per task instead of rechunking
  upward through dask, which is what OOM-killed the merge
- chunks holding no labels are skipped in the scan and the relabel, so
  background costs neither I/O nor disk
- the boundary scan runs in parallel; multi-config runs run concurrently

Correctness
- tile skipping used a centred window covering 6.25% of a tile, so objects in
  the outer ring were never segmented; an exact max-pooled occupancy map
  replaces it
- worker counts and memory budgets read the SLURM/cgroup allocation instead
  of the node's totals
- cupy OOM is retried like torch's; NVML reads the granted GPU, not index 0
- an in-place merge refuses to run twice, which would re-offset global ids
- NGFF metadata now matches the zarr version being written (0.5 on v3)
- config mistakes fail in prepare, not in the first GPU job

BREAKING CHANGE: stage_tile returns the number of labels it wrote rather than
the tile index. The workflow's markers move from seg/<tile>.done to
seg/<batch>.done and carry per-tile label counts as JSON. Stores are written
with NGFF 0.5 metadata on zarr v3; both layouts are still read.

Author: lguerard

README.md

@@ -286,9 +286,12 @@ tiles where the dask-image approach stalls.
 | ------------------------------ | ----------------------------------------- | ------------------------------------------------------------- |
 | In-process Dask client         | `FutureCancelledError: lost dependencies` | Detected at startup, raises immediately with fix instructions |
 | 3-4× fn recompute during merge | Cellpose runs 3× per tile                 | Staging writes labels once, merge reads from disk             |
-| O(n²) sequential relabelling   | Graph construction hangs at 1000+ tiles   | Linear post-pass O(voxels) via `np.unique` + LUT              |
+| O(n²) sequential relabelling   | Graph construction hangs at 1000+ tiles   | Folded into the merge's own LUT — no extra pass over the volume |
 | Wrong overlap boundary         | Output shape mismatch                     | Always uses `boundary="none"`                                 |
 | Persisting large arrays        | Worker OOM                                | Never persists; keeps dask graph lazy and streams             |
+| Sizing work to the whole node  | Job OOM-killed on a shared cluster node   | Reads the SLURM/cgroup allocation, not `os.cpu_count()`       |
+| Rechunking a pyramid level     | Threaded scheduler stockpiles intermediates | Levels stream one source chunk per task, bounded by construction |
+| Isotropic halo on flat tiles   | 5× the voxels read and segmented, then trimmed | `overlap` takes one width per axis                       |
 
 ---
 

docs/api/chunks.md

@@ -3,3 +3,17 @@
 ::: patchworks.auto_tile_shape
 
 ::: patchworks.auto_tile_shape_cellpose
+
+::: patchworks.auto_overlap
+
+::: patchworks.normalize_overlap
+
+## Cluster resource detection
+
+On a shared node the machine's core count and free RAM say nothing about what
+this job was granted. These read the allocation instead, and everything that
+sizes a worker pool goes through them.
+
+::: patchworks.cpu_allocation
+
+::: patchworks.safe_worker_count

docs/api/io.md

@@ -2,4 +2,18 @@
 
 ::: patchworks.load_ome_zarr
 
+## Deciding which tiles hold signal
+
+`estimate_empty_tiles` is a fast **preview** — it samples a centred window per
+tile, so it can miss signal at a tile's edge. `build_occupancy_map` +
+`tile_occupancy` are **exact**: a brick maximum exceeds the threshold exactly
+when some voxel in that brick does. Use the latter pair when the result is
+used as a skip list. See [Skipping empty tiles](../guide/skip_empty.md).
+
 ::: patchworks.estimate_empty_tiles
+
+::: patchworks.build_occupancy_map
+
+::: patchworks.tile_occupancy
+
+::: patchworks.auto_empty_threshold

docs/api/plugins/ome_zarr.md

@@ -24,3 +24,15 @@ resolution, matching anisotropic microscopy stacks.
 ## register_labels
 
 ::: patchworks.plugins.ome_zarr.register_labels
+
+## NGFF metadata layout
+
+NGFF 0.4 is defined over zarr v2 and puts its keys at the top level; 0.5 is
+the zarr-v3 revision and nests them under `ome`. patchworks writes whichever
+matches the store, and reads both.
+
+::: patchworks.plugins.ome_zarr.ngff_version
+
+::: patchworks.plugins.ome_zarr.read_ngff_attr
+
+::: patchworks.plugins.ome_zarr.write_ngff_attrs

docs/guide/custom_segmentation.md

@@ -118,6 +118,18 @@ def segment(tile: np.ndarray, prob_thresh: float = 0.5) -> np.ndarray:
 
 Using a GPU? Just let your framework see it — nothing extra needed.
 
+!!! tip "You do not need to modify patchworks"
+    `method: "custom"` imports any `(tile) -> labels` callable, so a new
+    method is a module of your own and a config block — nothing in the
+    package changes. It then inherits everything the pipeline does: empty
+    tiles are skipped, tiles are batched per job, labels are stitched across
+    boundaries and renumbered, and the result is written as a calibrated
+    pyramid.
+
+    (The `KNOWN_METHODS` list in `workflow/scripts/_pw.py` is only for
+    *built-in* shortcuts like `"cellpose"`. Adding to it is for methods that
+    ship with patchworks, not for your own.)
+
 ## Test it before you submit
 
 Run your function on one real tile first — it catches shape/dtype bugs in
@@ -148,6 +160,9 @@ custom:
   module: "my_seg"        # import name
   function: "segment"     # default is "segment"
   kwargs:                 # optional — forwarded as segment(tile, **kwargs)
+  #                         checked against your function's signature during
+  #                         `prepare`, so a typo fails on a cheap CPU job
+  #                         rather than in the first GPU job hours later
     sigma: 1.5
 ```
 

docs/guide/gpu_distributed.md

@@ -73,12 +73,35 @@ FutureCancelledError: lost dependencies
 ## GPU memory sizing
 
 When `use_gpu=True`, patchworks queries free GPU VRAM via `nvidia-ml-py`
-(install: `pip install "patchworks[gpu]"`). Tile size is set so each tile
-uses at most half the available VRAM.
+(install: `pip install "patchworks[gpu]"`) and keeps 20% of it as headroom,
+because `info.free` is a point-in-time reading of a device you usually do not
+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).
+
+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
+than the one your job was granted.
 
 Without `nvidia-ml-py`, a conservative 8 GiB default is used with a warning.
 Install it for accurate sizing on large-VRAM cards (A100, H100):
 
 ```bash
 pip install "patchworks[gpu]"
 ```
+
+!!! tip "Sizing for a GPU you can't see"
+    The Snakemake workflow plans tiles in `prepare`, which runs on a CPU
+    node, so no GPU can be queried there. Set `gpu_memory_gb` in the config
+    to the segment GPU's VRAM (e.g. `24` for an RTX 4090) instead of relying
+    on the fallback.
+
+## Surviving a shared GPU
+
+An out-of-memory error on a shared device is often transient — a co-tenant
+job grew, not a tile that doesn't fit. patchworks retries on the GPU with a
+backoff rather than falling back to the CPU, where a single tile can take
+over an hour. Before each backoff it releases its own device memory
+(including any cached Cellpose model), since holding that is exactly what
+starves the other job. Both torch and cupy OOM errors are recognised.

docs/guide/merging.md

@@ -28,37 +28,86 @@ This is the same approach used by
 
 ### Step 1: stage
 
-Each tile's labels are written to a temporary zarr once. This is critical:
-without staging, any downstream operation that reads the label array re-runs
-your segmentation function. The merge internally reads labels multiple times.
+Each tile's labels are written to zarr once. This is critical: without
+staging, any downstream operation that reads the label array re-runs your
+segmentation function. The merge internally reads labels multiple times.
 
 ```text
 tile_process calls fn once per tile → staged zarr
                                          │
                          merge reads from staged zarr (no fn calls)
 ```
 
-### Step 2: boundary scan
+The Snakemake workflow goes further and stages **directly into**
+`image.zarr/labels/<name>/0`, then has the merge rewrite that array in place
+— saving a whole extra write of the volume plus the scratch store's disk. It
+falls back to a separate store when the tile is larger than the label chunk
+cap, since in place the chunking cannot be changed and level 0 has to stay
+pageable for a viewer.
+
+### Step 2: make the ids globally unique
+
+Tiles write local `1..n`, which collide, so the boundary scan could not
+otherwise tell two different objects apart. If each tile's label **count** is
+known, this is just an exclusive cumulative sum — global id is
+`offset[tile] + local`, computed in `O(n_tiles)` with no read of the volume
+at all. `stage_tile` returns that count for exactly this purpose; pass the
+counts as `label_counts=`.
+
+Without counts, the merge falls back to streaming every chunk and renumbering
+it in place — correct, but a full read **and write** of the volume.
+
+### Step 3: boundary scan
 
 Only the two voxels on either side of each tile boundary are read. For any
-pair of touching non-zero labels `(a, b)`, they must be the same object.
+pair of touching non-zero labels `(a, b)`, they must be the same object. The
+per-tile offsets are applied here, on the fly.
 
-I/O cost: `O(n_boundaries × face_area)`, not `O(full_volume)`.
+I/O cost: `O(n_boundaries × face_area)`, not `O(full_volume)`. The columns
+are read in parallel, and a boundary next to a chunk that holds no labels is
+skipped outright — a pair needs a non-zero label on *both* sides, so it could
+never produce one.
 
-### Step 3: connected components
+### Step 4: connected components
 
 scipy sparse connected components on the touching pairs produces a relabeling
 lookup table. All labels that transitively touch each other are mapped to the
 same canonical label.
 
 Cost: `O(n_touching_pairs)`.
 
-### Step 4: parallel relabel
+With `sequential_labels=True` the contiguous renumbering is folded into this
+same LUT. Because the id domain is dense by construction, the surviving ids
+are exactly the distinct LUT values — a `np.unique` over an array the length
+of the object count, with no scan of the volume.
+
+### Step 5: parallel relabel
 
 The LUT is applied to every tile in parallel via `multiprocessing.Pool`. The
 LUT is shared via process initializer to avoid re-pickling it for every chunk
 (LUTs can be hundreds of MB for dense label volumes).
 
+Chunks whose tile wrote no labels are skipped entirely — not read, and not
+written. Zarr never materialises an unwritten chunk and reads it back as the
+fill value, so background regions cost neither I/O nor disk. On a sparse
+image that is most of the volume.
+
+When the merge's output *is* its input, this pass rewrites the array in
+place. That is safe because the boundary scan (step 3) has already finished,
+so nothing still needs the original ids.
+
+Because an in-place merge destroys its own input, it records how far it got
+on the array itself, and refuses to guess on a re-run:
+
+| State found | What happens |
+| --- | --- |
+| nothing recorded | fresh tile-local ids — merge normally |
+| `running` | a previous attempt died mid-relabel, so the array is part local and part global. **Refuses**: re-segment to rebuild it. |
+| `done` | already merged — a no-op, so a failure *after* the relabel (the pyramid, say) can simply be retried |
+
+Without that, a second pass would add each tile's offset to ids that are
+already global, which can land two unrelated objects on the same id.
+
 ## Using the merge step standalone
 
 You can call the merge step directly on any existing label array or zarr:
@@ -86,17 +135,20 @@ merged = merge_tile_labels(
 
 ## Sequential label numbering
 
-By default, merged labels are globally unique but may be **gappy**
-(block-encoded IDs like 1, 2, 500001, 500002, …). This is fine for
-counting, `regionprops`, and measurement — the IDs just aren't consecutive.
+By default, merged labels are globally unique but may be **gappy** — boundary
+merging fuses ids, leaving holes where the absorbed ones were. This is fine
+for counting, `regionprops`, and measurement — the IDs just aren't
+consecutive.
 
 For contiguous 1..N numbering, use `sequential_labels=True`:
 
 ```python
 tile_process("image.zarr", fn, write_to="labels.zarr", sequential_labels=True)
 ```
 
-This runs a cheap linear post-pass: `np.unique` + lookup-table remap, O(voxels).
+This is free: it composes into the relabel LUT the merge already applies, so
+it costs a `np.unique` over the object count rather than another pass over
+the volume.
 
 !!! warning "Do not use dask's built-in sequential relabel"
     `dask_image.ndmeasure.merge_labels_across_chunk_boundaries` has a

docs/guide/ome_zarr_napari.md

@@ -22,9 +22,23 @@ natively. Pass `write_to="other.zarr"` to instead write a separate
 single-resolution label store, or `output_component="cells"` to name the label
 image.
 
-The label pyramid is built lazily (`da.to_zarr`, streamed chunk by chunk), so
-it stays OOM-safe even for terabyte volumes. Control it with `pyramid_levels`
-and `pyramid_downscale`.
+The label pyramid is streamed chunk by chunk, so it stays OOM-safe even for
+terabyte volumes. Control it with `pyramid_levels` and `pyramid_downscale`.
+
+!!! note "NGFF version"
+    The metadata layout follows the zarr format being written: **NGFF 0.5**
+    (keys nested under an `ome` attribute) on zarr v3, **0.4** (keys at the
+    top level) on v2. Writing v3 data with 0.4's top-level layout — which
+    earlier versions did — matches neither revision, and a strict 0.5 reader
+    finds nothing there.
+
+    Reading accepts **both** layouts, so stores written by any patchworks
+    version still load. Use `read_ngff_attr()` if you need to inspect the
+    metadata yourself rather than indexing `attrs["multiscales"]` directly.
+
+    patchworks' own hints (`n_objects`, `sequential_labels`) are not NGFF
+    keys, so they stay at the top level where a consumer finds them without
+    knowing the layout.
 
 ## Why a pyramid?
 
@@ -107,6 +121,21 @@ input — see [Cluster usage](snakemake.md).
     to_ome_zarr("scan.ims", "scan.zarr", reuse_pyramid=True)
     ```
 
+### Axis names
+
+For a file, axes come from the reader's own dimension metadata (bioio's
+`dims.order`, Imaris, or the `sequence_pattern` groups) — not guessed. Only
+non-spatial singleton axes are dropped, so a genuine `z=1` survives.
+
+For a **bare array** there is nothing to read, so axes are inferred from the
+number of dimensions (`yx`, `zyx`, `czyx`, `tczyx`). That is unambiguous up to
+3-D; from 4-D the leading axis could be channel or time, and patchworks logs a
+warning saying what it guessed. Pass `axes=` to settle it:
+
+```python
+to_ome_zarr(arr, "out.zarr", axes="tzyx")   # not the czyx it would assume
+```
+
 ### Pixel calibration
 
 The physical voxel size is read from the input — bioio's `physical_pixel_sizes`,

docs/guide/performance.md

@@ -3,19 +3,27 @@
 `tile_process` is built so a run **adapts to whatever machine it lands on** and
 can't run out of RAM/VRAM or freeze the box — without you tuning anything.
 
-## Automatic, machine-aware concurrency
+## Automatic, allocation-aware concurrency
 
 The staging step (running your `fn` once per tile to a temp store) and the
-merge step are sized to the host automatically:
+merge step are sized automatically:
 
 - **GPU** (`use_gpu=True`) → **one tile at a time**, so concurrent evaluations
   can never exhaust VRAM.
 - **CPU** → as many tiles in flight as fit **80 % of available RAM** (estimated
   from the tile size), and always **leaving one core free** so the machine
   stays responsive — it never pins every core.
 
-The RAM figure is read live via `psutil`; without it, a conservative default is
-used instead of guessing high.
+"Available" means available **to this process**, not to the machine. On a
+shared cluster node those differ wildly — a 32-core, 128 GB job on a 128-core,
+512 GB node would otherwise size itself for the whole box and get OOM-killed
+while its own accounting said it had room. patchworks takes the smallest of
+`SLURM_MEM_PER_NODE`, `SLURM_MEM_PER_CPU × cpus`, the cgroup limit (the one
+that actually triggers the kill) and `psutil`'s free RAM, and reads the core
+count from `SLURM_CPUS_PER_TASK` or the process' CPU affinity mask.
+
+Without any of those signals, a conservative default is used instead of
+guessing high.
 
 ## Live progress dashboard (GPU runs)