perf(ome_zarr): ⚡ stream pyramid levels instead of rechunking through dask

This is what OOM-killed the merge job. _write_pyramid strided each level then
rechunked it back *upward* -- (16,512,512) to (16,1024,1024) -- so every
output chunk pulled four source chunks, and dask's threaded scheduler holds
those intermediates with no backpressure and no cap tied to the SLURM
allocation.

Decimation needs no halo, so a level's chunks can be chosen as
ceil(src_chunk / stride): one task then reads exactly one source chunk and
writes exactly one output chunk, nothing shared, peak memory
n_workers x one chunk however large the image is. Levels are written
zarr-natively via a bounded thread pool (zarr's codecs release the GIL).
Explicit chunks= or sharding still take the dask path, which genuinely needs
the rechunk.

The merge also stops writing the volume twice. merge_tile_labels gained
output_chunks, so it writes straight into labels/<name>/0 -- register_labels
already assumes that level exists -- instead of filling a scratch
_merged.zarr that write_labels then copied across chunk by chunk. The cap has
to divide the tile shape or concurrent workers would read-modify-write a
shared chunk, so capped_output_chunks picks the largest divisor at or below
the cap.

A test asserts the streaming and dask routes produce byte-identical levels,
including the ragged trailing chunks.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: lguerard

src/patchworks/__init__.py

@@ -39,7 +39,7 @@
     stage_tile,
 )
 from ._io import auto_empty_threshold, estimate_empty_tiles, load_ome_zarr
-from ._merge import merge_tile_labels
+from ._merge import capped_output_chunks, merge_tile_labels
 from ._occupancy import build_occupancy_map, occupancy_path, tile_occupancy
 from ._postprocess import dilate_labels
 from ._relabel import relabel_sequential_array, relabel_sequential_zarr
@@ -52,6 +52,7 @@
 __all__ = [
     "tile_process",
     "merge_tile_labels",
+    "capped_output_chunks",
     "auto_overlap",
     "auto_tile_shape",
     "auto_tile_shape_cellpose",

src/patchworks/_merge.py

@@ -343,6 +343,46 @@ def _make_globally_unique(arr, shape: tuple, chunk_shape: tuple) -> int:
     return base
 
 
+def capped_output_chunks(
+    chunk_shape: Sequence[int], caps: Sequence[int]
+) -> tuple[int, ...]:
+    """Shrink each chunk to at most *cap*, staying an exact divisor.
+
+    The merge's workers write one staged chunk at a time. As long as the output
+    chunking divides the staged chunking, each write still covers whole output
+    chunks, so concurrent workers never read-modify-write a chunk they share.
+    A non-divisor cap would break that, so the largest divisor at or below the
+    cap is used instead.
+
+    Parameters
+    ----------
+    chunk_shape : sequence of int
+        Staged chunk (= tile) shape.
+    caps : sequence of int
+        Maximum chunk size per axis.
+
+    Returns
+    -------
+    tuple of int
+        Chunking for the merged output.
+
+    Examples
+    --------
+    >>> capped_output_chunks((16, 2048, 2048), (16, 1024, 1024))
+    (16, 1024, 1024)
+    >>> capped_output_chunks((16, 1024, 1024), (16, 1024, 1024))
+    (16, 1024, 1024)
+    """
+    out = []
+    for c, cap in zip(chunk_shape, caps):
+        c, cap = int(c), int(cap)
+        if c <= cap:
+            out.append(c)
+            continue
+        out.append(next(d for d in range(cap, 0, -1) if c % d == 0))
+    return tuple(out)
+
+
 def _offsets_from_counts(
     counts: Mapping[int, int] | Sequence[int], n_chunks: int
 ) -> np.ndarray:
@@ -397,6 +437,7 @@ def zarr_native_merge(
     show_progress: bool = False,
     label_counts: "Mapping[int, int] | Sequence[int] | None" = None,
     sequential: bool = False,
+    output_chunks: "Sequence[int] | None" = None,
 ) -> "int | None":
     """Zarr-native label merge: boundary scan → scipy CC → parallel relabel.
 
@@ -429,6 +470,10 @@ def zarr_native_merge(
         domain is dense by construction, so the compaction is a ``np.unique``
         over the LUT (length = object count) and folds into the same lookup,
         costing no extra pass over the volume.
+    output_chunks : sequence of int, optional
+        Chunking for the merged store. Must divide the staged chunk shape, so
+        each worker's write still covers whole chunks. ``None`` mirrors the
+        staged chunking. Use :func:`capped_output_chunks` to derive it.
 
     Returns
     -------
@@ -497,10 +542,23 @@ def zarr_native_merge(
         )
 
     out_root = zarr.open_group(out_path, mode="a")
+    out_chunks = tuple(chunk_shape)
+    if output_chunks is not None:
+        out_chunks = tuple(int(c) for c in output_chunks)
+        bad = [
+            (i, c, o)
+            for i, (c, o) in enumerate(zip(chunk_shape, out_chunks))
+            if o <= 0 or c % o
+        ]
+        if bad:
+            raise ValueError(
+                "output_chunks must divide the staged chunk shape so workers "
+                f"write whole chunks; axis/staged/output mismatches: {bad}"
+            )
     # Match the staged dtype: ids are already compact (dense by construction,
     # and compacted again above when sequential), so nothing needs a wider one.
     _create_zarr_label_array(
-        out_root, out_component, shape, chunk_shape, dtype=arr.dtype
+        out_root, out_component, shape, out_chunks, dtype=arr.dtype
     )
 
     # Row-major, matching spatial_tiles' order -- so chunk i is tile i and the
@@ -585,6 +643,7 @@ def merge_tile_labels(
     progress: bool = False,
     return_count: bool = False,
     label_counts: "Mapping[int, int] | Sequence[int] | None" = None,
+    output_chunks: "Sequence[int] | None" = None,
 ) -> Union["da.Array", tuple["da.Array", Union[int, None]]]:
     """Merge per-tile labels into a globally consistent label array.
 
@@ -627,6 +686,11 @@ def merge_tile_labels(
         merge derive every tile's global id range arithmetically instead of
         streaming the whole store to renumber it. ``None`` keeps the
         renumber pass.
+    output_chunks:
+        Chunking for the merged store; must divide the staged chunk shape.
+        Lets the merge write straight into a store you will keep (e.g. an
+        OME-ZARR label group's level 0) instead of a scratch store that then
+        has to be copied. See :func:`capped_output_chunks`.
     n_workers:
         Parallel workers for the relabel step. Default ``min(4, cpu_count)``.
     stage_dir:
@@ -741,6 +805,7 @@ def merge_tile_labels(
         show_progress=progress,
         label_counts=label_counts,
         sequential=sequential_labels,
+        output_chunks=output_chunks,
     )
 
     # -- Cleanup temp stage (only when we created it) --

src/patchworks/plugins/ome_zarr.py

@@ -25,8 +25,13 @@
 Downsampling uses strided (nearest-neighbour) subsampling — the correct,
 label-preserving choice — and only on **X and Y**; ``z`` (and channel/time)
 stay at full resolution. Every level is built by reading the *previous level
-back from disk* and streaming the downsampled result out through dask with
-bounded chunks, so the pyramid never materialises a whole volume in RAM.
+back from disk*, so the pyramid never materialises a whole volume in RAM.
+Because decimation needs no halo, levels are written zarr-natively: each
+level's chunks are chosen as ``ceil(src_chunk / stride)`` so one task reads
+exactly one source chunk and writes exactly one output chunk, making peak
+memory ``n_workers × one chunk`` by construction. Explicit ``chunks=`` or
+sharding fall back to dask, which must rechunk and therefore reads several
+source chunks per output chunk.
 
 This module also exposes :func:`add_pyramid` (add levels to an existing store)
 and :func:`write_labels` (store a label image under the NGFF ``labels/`` group).
@@ -42,6 +47,8 @@
 
 import logging
 import math
+from concurrent.futures import ThreadPoolExecutor
+from itertools import product as _iproduct
 from pathlib import Path
 from typing import Union
 
@@ -273,6 +280,97 @@ def _shard_for(
     return _effective_shard(tuple(shard), chunks, shape)
 
 
+# Don't let deep pyramid levels fragment into thousands of tiny chunks.
+_CHUNK_FLOOR = {"y": 128, "x": 128}
+
+
+def _level_chunks(
+    src_chunks: tuple[int, ...],
+    strides: tuple[int, ...],
+    axes: str,
+    level_shape: tuple[int, ...],
+) -> tuple[int, ...]:
+    """Chunking for a strided level that keeps one source chunk per task.
+
+    Choosing ``ceil(src_chunk / stride)`` makes each output chunk the exact
+    image of one source chunk, so a task reads one chunk and writes one chunk
+    with nothing shared between tasks. Capped by :data:`_CHUNK_CAP` and floored
+    so deep levels don't fragment.
+    """
+    out = []
+    for c, st, a, s in zip(src_chunks, strides, axes, level_shape):
+        v = -(-c // st)
+        v = min(v, _CHUNK_CAP.get(a, v))
+        v = max(v, _CHUNK_FLOOR.get(a, 1))
+        out.append(max(1, min(v, s)))
+    return tuple(out)
+
+
+def _create_level_array(
+    group: "zarr.Group",
+    name: str,
+    shape: tuple[int, ...],
+    chunks: tuple[int, ...],
+    dtype,
+) -> "zarr.Array":
+    """Create (replacing any existing) a pyramid level array in *group*."""
+    if name in group:
+        del group[name]
+    if _ZARR_V3:
+        return group.create_array(name, shape=shape, chunks=chunks, dtype=dtype)
+    return group.zeros(
+        name, shape=shape, chunks=chunks, dtype=dtype, overwrite=True
+    )
+
+
+def _stream_strided_level(
+    src: "zarr.Array",
+    dst: "zarr.Array",
+    strides: tuple[int, ...],
+    n_workers: int = 4,
+) -> None:
+    """Write *dst* as the strided subsample of *src*, one chunk at a time.
+
+    Labels are downsampled by plain decimation, which needs no halo, so each
+    output chunk depends only on the source region that maps onto it. Peak
+    memory is therefore ``n_workers x (one source region + its subsample)``
+    however large the image is -- unlike a dask ``rechunk``, whose threaded
+    scheduler holds intermediates with no backpressure.
+
+    Parameters
+    ----------
+    src, dst : zarr.Array
+        Source level and the (already created) destination level.
+    strides : tuple of int
+        Per-axis decimation step.
+    n_workers : int
+        Threads used for the copy. zarr's codecs release the GIL, so threads
+        are enough and avoid pickling a worker payload.
+    """
+    grid = [-(-s // c) for s, c in zip(dst.shape, dst.chunks)]
+    take = tuple(slice(None, None, st) for st in strides)
+
+    def _one(idx: tuple[int, ...]) -> None:
+        out_sl = tuple(
+            slice(i * c, min((i + 1) * c, s))
+            for i, c, s in zip(idx, dst.chunks, dst.shape)
+        )
+        src_sl = tuple(
+            slice(o.start * st, min(o.stop * st, s))
+            for o, st, s in zip(out_sl, strides, src.shape)
+        )
+        dst[out_sl] = np.asarray(src[src_sl])[take]
+
+    indices = list(_iproduct(*[range(g) for g in grid]))
+    if n_workers <= 1:
+        for idx in indices:
+            _one(idx)
+        return
+    with ThreadPoolExecutor(max_workers=n_workers) as pool:
+        for _ in pool.map(_one, indices):
+            pass
+
+
 def _progress_ctx(progress: bool, label: str):
     """Return a progress-bar context manager.
 
@@ -484,19 +582,51 @@ def _write_pyramid(
     prev_name = base_name
     prev_shape = arr.shape
     for i in range(1, n_levels):
-        next_shape = tuple(s // st for s, st in zip(prev_shape, strides))
+        next_shape = tuple(-(-s // st) for s, st in zip(prev_shape, strides))
         if min(next_shape) < 1:
             logger.info("stopping pyramid at level %d (next too small)", i)
             break
-        src = da.from_zarr(group_path, component=prev_name)
-        nxt = src[tuple(slice(None, None, st) for st in strides)]
-        nxt = nxt.rechunk(chunks or _default_chunks(nxt.shape, axes))
-        _to_zarr_level(nxt, group_path, str(i), shard, progress)
+        if shard or chunks:
+            # Sharding needs whole shards written atomically, and an explicit
+            # chunking may not line up with the source chunks -- both want the
+            # dask path.
+            src = da.from_zarr(group_path, component=prev_name)
+            nxt = src[tuple(slice(None, None, st) for st in strides)]
+            nxt = nxt.rechunk(chunks or _default_chunks(nxt.shape, axes))
+            _to_zarr_level(nxt, group_path, str(i), shard, progress)
+            next_shape = nxt.shape
+        else:
+            # Stream it: decimation needs no halo, so with chunks chosen as
+            # ceil(src_chunk / stride) each task reads exactly one source chunk
+            # and writes exactly one output chunk. Bounded by construction --
+            # the dask route rechunks *upward* here (e.g. (16,512,512) back to
+            # (16,1024,1024)), pulling four source chunks per output chunk and
+            # letting the threaded scheduler stockpile the intermediates.
+            grp = zarr.open_group(group_path, mode="a")
+            src_arr = grp[prev_name]
+            out_chunks = _level_chunks(
+                tuple(src_arr.chunks), strides, axes, next_shape
+            )
+            dst_arr = _create_level_array(
+                grp, str(i), next_shape, out_chunks, src_arr.dtype
+            )
+            if progress:
+                logger.info(
+                    "writing %s/%s (streaming %d chunks)",
+                    Path(group_path).name,
+                    i,
+                    int(
+                        np.prod(
+                            [-(-s // c) for s, c in zip(next_shape, out_chunks)]
+                        )
+                    ),
+                )
+            _stream_strided_level(src_arr, dst_arr, strides)
         scale = [base_scale[k] * (strides[k] ** i) for k in range(len(axes))]
         datasets.append(_dataset(str(i), scale))
-        logger.info("pyramid level %d: shape=%s", i, nxt.shape)
+        logger.info("pyramid level %d: shape=%s", i, next_shape)
         prev_name = str(i)
-        prev_shape = nxt.shape
+        prev_shape = next_shape
     return datasets
 
 

tests/test_ome_zarr.py

@@ -14,6 +14,58 @@
 )
 
 
+def test_streaming_pyramid_matches_dask(tmp_path):
+    """The streaming levels must be byte-identical to the dask-built ones.
+
+    The streaming path exists to bound memory (one source chunk per task); it
+    must not change a single voxel of the result.
+    """
+    rng = np.random.default_rng(0)
+    data = rng.integers(0, 500, size=(9, 100, 100), dtype="uint16")
+
+    streamed = tmp_path / "streamed.zarr"
+    to_ome_zarr(data, streamed, axes="zyx", n_levels=4, progress=False)
+    # chunks= forces the dask route through the same public entry point.
+    dasked = tmp_path / "dasked.zarr"
+    to_ome_zarr(
+        data,
+        dasked,
+        axes="zyx",
+        n_levels=4,
+        chunks=(16, 1024, 1024),
+        progress=False,
+    )
+
+    a_root = zarr.open_group(str(streamed), mode="r")
+    b_root = zarr.open_group(str(dasked), mode="r")
+    a_levels = [d["path"] for d in a_root.attrs["multiscales"][0]["datasets"]]
+    b_levels = [d["path"] for d in b_root.attrs["multiscales"][0]["datasets"]]
+    assert a_levels == b_levels, "both routes must produce the same levels"
+    for lvl in a_levels:
+        assert np.array_equal(
+            np.asarray(a_root[lvl]), np.asarray(b_root[lvl])
+        ), f"level {lvl} differs between the streaming and dask routes"
+
+
+def test_streaming_level_reads_one_source_chunk_per_task(tmp_path):
+    """Each output chunk must map onto exactly one source chunk.
+
+    That alignment is what makes peak memory ``n_workers x one chunk`` and so
+    is the property the OOM fix rests on -- not an incidental detail.
+    """
+    from patchworks.plugins.ome_zarr import _level_chunks
+
+    # A (16, 1024, 1024) source at stride (1, 2, 2) -> (16, 512, 512).
+    out = _level_chunks((16, 1024, 1024), (1, 2, 2), "zyx", (16, 4096, 4096))
+    assert out == (16, 512, 512)
+    # Deep levels get floored rather than fragmenting into tiny chunks.
+    floored = _level_chunks((16, 128, 128), (1, 2, 2), "zyx", (16, 256, 256))
+    assert floored == (16, 128, 128)
+    # Never larger than the level itself.
+    clamped = _level_chunks((16, 1024, 1024), (1, 2, 2), "zyx", (4, 100, 100))
+    assert clamped == (4, 100, 100)
+
+
 def _level_scale(store, level):
     root = zarr.open_group(str(store), mode="r")
     ds = root.attrs["multiscales"][0]["datasets"][level]