perf: machine-adaptive, OOM-safe processing

docs/guide/performance.md

@@ -0,0 +1,60 @@
+# Performance & memory safety
+
+`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
+
+The staging step (running your `fn` once per tile to a temp store) and the
+merge step are sized to the host 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.
+
+## Overriding the worker count
+
+```python
+from patchworks import tile_process
+
+# let patchworks pick (recommended)
+tile_process("scan.zarr", fn)
+
+# or cap it yourself (staging threads + merge processes)
+tile_process("scan.zarr", fn, max_workers=8)
+```
+
+`max_workers` bounds both staging and merging. A running **distributed client**
+manages its own concurrency, so the override is skipped there — configure the
+cluster's memory limits instead.
+
+## Why it won't OOM or freeze
+
+| Resource | Guard |
+|----------|-------|
+| RAM | concurrent tiles × tile size × overhead ≤ 80 % of available RAM |
+| VRAM | GPU path runs one tile at a time |
+| CPU | always leaves at least one core free |
+| Disk I/O | each pyramid/stage level is streamed chunk-by-chunk; no whole volume in memory |
+
+The staging graph itself is kept small — a single fused `map_overlap`
+(halo → `fn` → trim) rather than three separate passes — and there is **no**
+extra read-back of the staged data.
+
+## Getting more speed
+
+- `tile_shape="auto"` sizes tiles to free RAM (or VRAM with `use_gpu=True`).
+- `skip_empty=True` with `estimate_empty_tiles()` skips background tiles.
+- A Dask **distributed** cluster (`make_local_cluster`) parallelises across
+  workers/GPUs; patchworks then defers concurrency to the cluster.
+
+!!! note "What doesn't help here"
+    The merge and relabel steps are already vectorised NumPy + SciPy (C-level)
+    with no per-voxel Python loop, and the pipeline is I/O-bound — so `numba`,
+    `cupy`, `arrow` and `xarray` bring essentially nothing. The real levers are
+    tile size, concurrency (above) and zarr chunking.

mkdocs.yml

@@ -38,6 +38,7 @@ nav:
   - Merging labels: guide/merging.md
   - Empty tile skipping: guide/skip_empty.md
   - GPU & distributed: guide/gpu_distributed.md
+  - Performance & memory: guide/performance.md
   - OME-ZARR & napari: guide/ome_zarr_napari.md
   - Pitfalls: guide/pitfalls.md
 - Examples:

src/patchworks/_chunks.py

@@ -57,6 +57,51 @@ def _get_available_memory() -> int:
         return 8 * 1024**3
 
 
+def safe_worker_count(
+    tile_nbytes: int,
+    *,
+    use_gpu: bool = False,
+    fn_overhead: int = 4,
+    ram_fraction: float = 0.8,
+) -> int:
+    """Concurrent tiles that fit the machine without OOM or a CPU freeze.
+
+    Bounds the threaded scheduler by two limits and takes the smaller:
+
+    * **CPU** — leaves at least one core free so the box stays responsive
+      (never pins every core).
+    * **RAM** — at most ``ram_fraction`` of available memory, assuming each
+      in-flight tile needs ``fn_overhead`` copies (halo + output + temporaries).
+
+    On GPU the answer is always 1: one evaluation at a time so concurrent
+    tiles can never exhaust VRAM. Without ``psutil`` it returns a conservative
+    default rather than guessing high.
+
+    Parameters
+    ----------
+    tile_nbytes : int
+        Size of one tile in bytes (``prod(tile_shape) * dtype.itemsize``).
+    use_gpu : bool, optional
+        Whether tiles are processed on the GPU.
+    fn_overhead : int, optional
+        Assumed peak number of tile-sized buffers alive per worker.
+    ram_fraction : float, optional
+        Fraction of available RAM the staging step may use.
+
+    Returns
+    -------
+    int
+        Worker-thread count (always >= 1).
+    """
+    cpu_cap = max(1, (os.cpu_count() or 1) - 1)
+    if use_gpu:
+        return 1
+    avail = _get_available_memory()
+    per_tile = max(1, int(tile_nbytes) * max(1, fn_overhead))
+    mem_cap = max(1, int(avail * ram_fraction) // per_tile)
+    return max(1, min(cpu_cap, mem_cap))
+
+
 def _get_gpu_memory() -> int:
     """Return free GPU VRAM in bytes. Falls back to 8 GiB default."""
     try:

src/patchworks/_core.py

@@ -11,7 +11,7 @@
 import dask.array as da
 import numpy as np
 
-from ._chunks import auto_tile_shape
+from ._chunks import auto_tile_shape, safe_worker_count
 from ._cluster import _client_is_in_process, _distributed_client
 from ._io import _auto_empty_threshold, load_ome_zarr
 from ._merge import zarr_native_merge
@@ -56,6 +56,7 @@ def tile_process(
     channel: int | None = 0,
     level: int = 0,
     use_gpu: bool = False,
+    max_workers: int | None = None,
     progress: bool = False,
     write_to: Union[str, Path, None] = None,
     output_component: str = "labels",
@@ -114,6 +115,13 @@ def tile_process(
         Pyramid level when *image* is a path (0 = full resolution).
     use_gpu:
         When ``tile_shape="auto"``, size tiles against GPU VRAM instead of RAM.
+        Also forces staging to one tile at a time (no VRAM contention).
+    max_workers:
+        Cap the worker threads/processes used for staging and merging. ``None``
+        (default) auto-sizes to the machine: bounded by available RAM (tile
+        size) and CPU (leaves one core free) so a run can neither OOM nor pin
+        every core. Ignored when a distributed client is active (it manages its
+        own concurrency).
     progress:
         Show a progress bar during the tile-writing and relabel steps.
     write_to:
@@ -283,11 +291,6 @@ def tile_process(
         for ax, c in enumerate(image.chunks)
     }
 
-    if overlap > 0:
-        # boundary="none" is required: only this boundary mode composes with
-        # trim_overlap to recover the original shape. "reflect" keeps the halo.
-        image = da.overlap.overlap(image, depth=_depth, boundary="none")
-
     # Wrap fn with optional empty-tile skipping
     _skip_thr = empty_threshold
     if skip_empty and _skip_thr is None:
@@ -303,28 +306,43 @@ def active_fn(block, block_info=None):
             logger.debug("process tile %s shape=%s", loc, block.shape)
         return fn(block)
 
-    labeled = image.map_blocks(
-        active_fn,
-        dtype=np.int32,
-        meta=np.empty((0,) * image.ndim, dtype=np.int32),
-    )
-
-    # Trim the overlap halo so staged tiles have clean boundaries for the
-    # boundary-slab scan. Without this the scan reads halo-expanded chunks and
-    # the merged output is larger than the input.
+    _meta = np.empty((0,) * image.ndim, dtype=np.int32)
     if overlap > 0:
-        labeled = da.overlap.trim_overlap(
-            labeled, depth=_depth, boundary="none"
+        # One fused pass: add the halo, run fn, trim it back off. map_overlap
+        # materialises only the halos it needs (no separate overlapped array)
+        # and keeps the task graph small. boundary="none" + trim recovers the
+        # original shape, so the boundary-slab scan reads clean tiles.
+        labeled = da.map_overlap(
+            active_fn,
+            image,
+            depth=_depth,
+            boundary="none",
+            trim=True,
+            dtype=np.int32,
+            meta=_meta,
         )
+    else:
+        labeled = image.map_blocks(active_fn, dtype=np.int32, meta=_meta)
 
-    # With no distributed client the threaded scheduler runs many tiles at
-    # once. For GPU that means several evals sharing one device → CUDA OOM.
-    # Pin to a single worker thread so evals run serially. A distributed
-    # client manages its own concurrency, so skip the override there.
+    # Bound staging concurrency to the machine so it can neither OOM nor pin
+    # every core:
+    #   - GPU → 1 eval at a time (no VRAM contention),
+    #   - CPU → as many tiles as fit RAM, leaving one core free.
+    # A distributed client manages its own concurrency, so skip the override.
     import dask as _dask
 
-    if _active is None and use_gpu:
-        _sched_ctx: Any = _dask.config.set(scheduler="threads", num_workers=1)
+    _tile_nbytes = int(np.prod(labeled.chunksize)) * labeled.dtype.itemsize
+    if _active is None:
+        _workers = (
+            max_workers
+            if max_workers is not None
+            else safe_worker_count(_tile_nbytes, use_gpu=use_gpu)
+        )
+        _workers = max(1, min(_workers, os.cpu_count() or 1))
+        logger.info("Staging with %d worker thread(s)", _workers)
+        _sched_ctx: Any = _dask.config.set(
+            scheduler="threads", num_workers=_workers
+        )
     else:
         _sched_ctx = _nullcontext()
 
@@ -347,24 +365,10 @@ def active_fn(block, block_info=None):
         _stage_to_zarr(labeled, stage_path, "staged", progress)
     labeled = da.from_zarr(stage_path, component="staged")
 
-    if skip_empty and _skip_thr is not None:
-
-        def _tile_max(block: np.ndarray) -> np.ndarray:
-            return np.full((1,) * block.ndim, int(block.max()), dtype=np.int32)
-
-        _tile_maxes = labeled.map_blocks(
-            _tile_max,
-            dtype=np.int32,
-            chunks=tuple(tuple(1 for _ in c) for c in labeled.chunks),
-        ).compute()
-        _n_skip = int((_tile_maxes == 0).sum())
-        logger.info(
-            "skip_empty: %d/%d tiles ran fn, %d skipped (max<=%.4g)",
-            int(_tile_maxes.size) - _n_skip,
-            int(_tile_maxes.size),
-            _n_skip,
-            _skip_thr,
-        )
+    # NB: no post-staging skip-count pass here — counting skipped tiles by
+    # re-reading the whole staged store off disk would double the I/O of the
+    # entire run just for a log line. Use estimate_empty_tiles() up front for
+    # that figure instead.
 
     def _cleanup_stage():
         if not keep_stage:
@@ -373,7 +377,9 @@ def _cleanup_stage():
             shutil.rmtree(stage_path, ignore_errors=True)
             logger.info("Removed stage store %s", stage_path)
 
-    _nw = min(4, os.cpu_count() or 1)
+    # Merge runs in worker processes (each holds one chunk + an mmap'd LUT);
+    # size it to RAM/CPU like staging, capped so we don't spawn a process storm.
+    _nw = max_workers or max(1, min(safe_worker_count(_tile_nbytes), 8))
 
     # Default: input is a .zarr store and no explicit write_to → labels go back
     # *into* the input store under the NGFF labels/<name>/ group with an auto

tests/test_core.py

@@ -247,3 +247,27 @@ def test_estimate_empty_tiles():
     assert info["n_tiles"] == 4
     assert info["n_occupied"] == 2
     assert info["empty_fraction"] == 0.5
+
+
+def test_safe_worker_count_bounds():
+    import os
+
+    from patchworks._chunks import safe_worker_count
+
+    # GPU → always serial (no VRAM contention)
+    assert safe_worker_count(10**6, use_gpu=True) == 1
+    # Absurdly large tile → memory-bound to 1
+    assert safe_worker_count(10**15) == 1
+    # Tiny tile → CPU-bound, leaves a core free, always >= 1
+    n = safe_worker_count(1024)
+    assert 1 <= n <= max(1, (os.cpu_count() or 1) - 1)
+
+
+def test_tile_process_max_workers():
+    import dask.array as da
+
+    from patchworks import tile_process
+
+    arr = da.from_array(_make_image((2, 32, 32)), chunks=(1, 32, 32))
+    result = tile_process(arr, _label_fn, max_workers=1).compute()
+    assert result.shape == (2, 32, 32)