perf: machine-adaptive, OOM-safe processing - #16
Merged
Conversation
- 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
lguerard
changed the title
- 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 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Optimises the processing/staging step itself, regardless of the user function, and makes it adapt to the host without OOM or CPU freeze.
Lighter graph
overlap()+map_blocks()+trim_overlap()→ one fusedda.map_overlap(..., boundary="none", trim=True): only needed halos materialised, smaller task graph.estimate_empty_tiles()gives that up front.Adapts to the machine, can't OOM or freeze
safe_worker_count()sizes staging threads and merge processes to the host:max_workers=override; a distributed client keeps managing its own concurrencyNot added (on purpose)
numba/cupy/arrow/xarray bring ~nothing here: the merge/relabel hotspots are already vectorized numpy + scipy C with no per-voxel Python loop, and the run is I/O-bound. Adding them would be deps for no gain.
Tests green (16 passed incl. boundary/overlap + new worker-count/max_workers); ruff clean.
🤖 Generated with Claude Code