From ec84a163e4453386b98911393967c6ec3499f73b Mon Sep 17 00:00:00 2001 From: Laurent Guerard Date: Tue, 23 Jun 2026 09:30:03 +0200 Subject: [PATCH 1/3] =?UTF-8?q?perf:=20=E2=9A=A1=20slim=20the=20staging=20?= =?UTF-8?q?graph=20and=20drop=20redundant=20disk=20re-read?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - replace overlap() + map_blocks() + trim_overlap() with a single fused da.map_overlap(..., trim=True): only the needed halos are materialised and the task graph is smaller (no separate overlapped array) - remove the post-staging skip-count pass, which re-read the *entire* staged store off disk just to log how many tiles were skipped — doubling the run's read I/O for a log line. estimate_empty_tiles() already reports that up front. Both are independent of the user fn, so every run gets cheaper. Co-Authored-By: Claude Opus 4.8 --- src/patchworks/_core.py | 53 +++++++++++++++-------------------------- 1 file changed, 19 insertions(+), 34 deletions(-) diff --git a/src/patchworks/_core.py b/src/patchworks/_core.py index 56cfa3c..0007b83 100644 --- a/src/patchworks/_core.py +++ b/src/patchworks/_core.py @@ -283,11 +283,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,19 +298,23 @@ 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. @@ -347,24 +346,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: From 708674dc709945f597970948c552578c0e4e136f Mon Sep 17 00:00:00 2001 From: Laurent Guerard Date: Tue, 23 Jun 2026 09:37:41 +0200 Subject: [PATCH 2/3] =?UTF-8?q?perf:=20=E2=9A=A1=20bound=20staging/merge?= =?UTF-8?q?=20concurrency=20to=20the=20machine=20(no=20OOM/freeze)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add safe_worker_count() and use it to size both the staging threaded scheduler and the merge worker pool to the host: - GPU → 1 tile at a time (no VRAM contention) - CPU → as many tiles as fit available RAM (tile size × overhead), capped to leave one core free so the box never fully freezes - new max_workers= knob to override; a distributed client keeps managing its own concurrency So a run adapts to whatever machine it lands on and can't blow up RAM, VRAM, or peg every core. Co-Authored-By: Claude Opus 4.8 --- src/patchworks/_chunks.py | 45 +++++++++++++++++++++++++++++++++++++++ src/patchworks/_core.py | 37 +++++++++++++++++++++++++------- tests/test_core.py | 24 +++++++++++++++++++++ 3 files changed, 98 insertions(+), 8 deletions(-) diff --git a/src/patchworks/_chunks.py b/src/patchworks/_chunks.py index 56ef2e3..febf054 100644 --- a/src/patchworks/_chunks.py +++ b/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.6, +) -> 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: diff --git a/src/patchworks/_core.py b/src/patchworks/_core.py index 0007b83..cf60fdf 100644 --- a/src/patchworks/_core.py +++ b/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: @@ -316,14 +324,25 @@ def active_fn(block, block_info=None): 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() @@ -358,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// group with an auto diff --git a/tests/test_core.py b/tests/test_core.py index c3cb26c..8958c30 100644 --- a/tests/test_core.py +++ b/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) From 9d0f85b09a3d186ae5c3e74597e43bf0aca10f98 Mon Sep 17 00:00:00 2001 From: Laurent Guerard Date: Tue, 23 Jun 2026 09:42:01 +0200 Subject: [PATCH 3/3] =?UTF-8?q?docs:=20=F0=9F=93=9D=20add=20performance/me?= =?UTF-8?q?mory=20guide;=20use=2080%=20RAM=20budget?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - raise the staging RAM budget from 60% to 80% of available memory - add a Performance & memory guide (auto worker sizing, max_workers, no-OOM/freeze guarantees, what doesn't help) + nav entry Co-Authored-By: Claude Opus 4.8 --- docs/guide/performance.md | 60 +++++++++++++++++++++++++++++++++++++++ mkdocs.yml | 1 + src/patchworks/_chunks.py | 2 +- 3 files changed, 62 insertions(+), 1 deletion(-) create mode 100644 docs/guide/performance.md diff --git a/docs/guide/performance.md b/docs/guide/performance.md new file mode 100644 index 0000000..5963233 --- /dev/null +++ b/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. diff --git a/mkdocs.yml b/mkdocs.yml index 4559e37..8567e10 100644 --- a/mkdocs.yml +++ b/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: diff --git a/src/patchworks/_chunks.py b/src/patchworks/_chunks.py index febf054..0659c3a 100644 --- a/src/patchworks/_chunks.py +++ b/src/patchworks/_chunks.py @@ -62,7 +62,7 @@ def safe_worker_count( *, use_gpu: bool = False, fn_overhead: int = 4, - ram_fraction: float = 0.6, + ram_fraction: float = 0.8, ) -> int: """Concurrent tiles that fit the machine without OOM or a CPU freeze.