perf: use sync methods for chunk encoding / decoding#3885
Conversation
`PreparedWrite` models a set of per-chunk changes that would be applied to a stored chunk. `SupportsChunkPacking` is a protocol for array -> bytes codecs that can use `PreparedWrite` objects to update an existing chunk.
…into perf/prepared-write-v2
…into perf/prepared-write-v2
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #3885 +/- ##
==========================================
+ Coverage 93.62% 93.87% +0.25%
==========================================
Files 90 91 +1
Lines 11936 12560 +624
==========================================
+ Hits 11175 11791 +616
- Misses 761 769 +8
🚀 New features to boost your workflow:
|
…into perf/prepared-write-v2
|
@TomAugspurger how would this design work with CUDA codecs? |
5d3064e to
b67a5a0
Compare
a84a15a to
68a7cdc
Compare
| # Phase 1: fetch all chunks (IO, sequential) | ||
| raw_buffers: list[Buffer | None] = [ | ||
| bg.get_sync(prototype=cs.prototype) # type: ignore[attr-defined] | ||
| for bg, cs, *_ in batch | ||
| ] | ||
|
|
||
| # Phase 2: decode (compute, optionally threaded) | ||
| def _decode_one(raw: Buffer | None, chunk_spec: ArraySpec) -> NDBuffer | None: | ||
| if raw is None: | ||
| return None | ||
| return transform.decode_chunk(raw, chunk_spec) | ||
|
|
||
| specs = [cs for _, cs, *_ in batch] | ||
| if n_workers > 0 and len(batch) > 1: | ||
| with ThreadPoolExecutor(max_workers=n_workers) as pool: | ||
| decoded_list = list(pool.map(_decode_one, raw_buffers, specs)) | ||
| else: | ||
| decoded_list = [ | ||
| _decode_one(raw, spec) for raw, spec in zip(raw_buffers, specs, strict=True) | ||
| ] |
There was a problem hiding this comment.
Why isn't this all multi-threaded i.e., the I/O as well?
There was a problem hiding this comment.
I should benchmark this, but my expectation was that IO against memory storage and local storage is not compute-limited, and so threads wouldn't remove a real bottleneck. for memory storage i'm sure this is true, not sure about local storage though
Adds a SupportsSetRange protocol to zarr.abc.store for stores that allow overwriting a byte range within an existing value. Implementations are added for LocalStore (using file-handle seek+write) and MemoryStore (in-memory bytearray slice assignment). This is the prerequisite for the partial-shard write fast path in ShardingCodec, which can patch individual inner-chunk slots without rewriting the entire shard blob when the inner codec chain is fixed-size. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
V2Codec, BytesCodec, BloscCodec, etc. previously only implemented the async _decode_single / _encode_single methods. Add their sync counterparts (_decode_sync / _encode_sync) so that the upcoming SyncCodecPipeline can dispatch through them without spinning up an event loop. For codecs that wrap external compressors (numcodecs.Zstd, numcodecs.Blosc, the V2 fallback chain), the sync versions just call the underlying compressor's blocking API directly instead of routing through asyncio.to_thread. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…arallelism
Adds SyncCodecPipeline alongside BatchedCodecPipeline. The new pipeline
runs codecs through their sync entry points (_decode_sync / _encode_sync)
and dispatches per-chunk work to a module-level thread pool sized by
the codec_pipeline.max_workers config (default = os.cpu_count()).
Each chunk's full lifecycle (fetch + decode + scatter for reads;
get-existing + merge + encode + set/delete for writes) runs as one
pool task — overlapping IO of one chunk with compute of another.
Scatter into the shared output buffer is thread-safe because chunks
have non-overlapping output selections.
The async wrappers (read/write) detect SupportsGetSync/SupportsSetSync
stores and dispatch to the sync fast path, passing the configured
max_workers. Other stores fall through to the async path, which still
uses asyncio.concurrent_map at async.concurrency.
Notes on perf:
- Default (None → cpu_count) is tuned for chunks ≥ ~512 KB.
- Small chunks (≤ 64 KB) regress 1.5-3x because pool dispatch overhead
(~30-50 µs/task) dominates per-chunk work. Workaround:
zarr.config.set({"codec_pipeline.max_workers": 1}).
- For large chunks on local/memory stores, IO+compute parallelism
yields 1.7-2.5x over BatchedCodecPipeline on direct-API reads and
~2.5x on roundtrip.
ChunkTransform encapsulates the sync codec chain. It caches resolved
ArraySpecs across calls with the same chunk_spec — combined with the
constant-ArraySpec optimization in indexing, hot-path overhead is
minimized.
Includes test scaffolding for the new pipeline (test_sync_codec_pipeline)
and config plumbing for the max_workers key.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds _encode_partial_sync and _decode_partial_sync to ShardingCodec.
For fixed-size inner codec chains and stores that implement
SupportsSetRange, partial writes patch individual inner-chunk slots
in-place instead of rewriting the whole shard:
- Reads existing shard index (one byte-range get).
- For each affected inner chunk: decodes the slot, merges the new
region, re-encodes.
- Writes each modified slot at its deterministic byte offset, then
rewrites just the index.
For variable-size inner codecs (e.g. with compression) or stores that
don't support byte-range writes, falls through to a full-shard rewrite
matching BatchedCodecPipeline semantics.
The partial-decode path computes a ReadPlan from the shard index and
issues one byte-range get per overlapping chunk, decoding only what
the read selection touches.
Both paths are dispatched from SyncCodecPipeline via the existing
supports_partial_decode / supports_partial_encode protocol checks.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Two new test files:
test_codec_invariants — asserts contract-level properties that every
codec / shard / buffer combination must satisfy: round-trip exactness,
prototype propagation, fill-value handling, all-empty shard handling.
test_pipeline_parity — exhaustive matrix asserting that
SyncCodecPipeline and BatchedCodecPipeline produce semantically
identical results across codec configs, layouts (including
nested sharding), write sequences, and write_empty_chunks settings.
Three checks per cell:
1. Same array contents on read.
2. Same set of store keys after writes.
3. Each pipeline reads the other's output identically (catches
layout-divergence bugs).
These tests pinned the design throughout the SyncCodecPipeline +
partial-shard development.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds .gitignore entries for .claude/, CLAUDE.md, and docs/superpowers/ so local IDE/agent planning artifacts don't get committed by accident. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
aa111a2 to
1be5563
Compare
| selected = decoded[chunk_selection] | ||
| if drop_axes: | ||
| selected = selected.squeeze(axis=drop_axes) | ||
| out[out_selection] = selected |
There was a problem hiding this comment.
It might be worth experimenting with moving this setting operation out[out_selection] = selected outside the threadpool execution since, IIRC, it holds the GIL and is probably non-trivial time-wise.
There was a problem hiding this comment.
The memory usage will probably go up a bit though....
|
@ilan-gold some changes for you to note:
|
ilan-gold
left a comment
There was a problem hiding this comment.
Just little things at this point!
| Threading is off by default because it is not universally a win: it slows small chunks (where the | ||
| per-task overhead outweighs the work), and on many-core nodes a pool sized to `cpu_count` can | ||
| oversubscribe workloads that already parallelize at a higher level (e.g. Dask). `max_workers` only | ||
| affects `FusedCodecPipeline`; the default batched pipeline ignores it. |
There was a problem hiding this comment.
I could do some work on this. I'll have to check the benchmarks to understand more. But any guidance here would be appreciated
Co-authored-by: Ilan Gold <ilanbassgold@gmail.com>
Co-authored-by: Ilan Gold <ilanbassgold@gmail.com>
Co-authored-by: Ilan Gold <ilanbassgold@gmail.com>
Co-authored-by: Ilan Gold <ilanbassgold@gmail.com>
Co-authored-by: Ilan Gold <ilanbassgold@gmail.com>
|
I have found Lachi's https://github.com/zarrs/zarr_benchmarks to be the most reflective benchmark of real-world data. Here are the most relevant results IMO on my 16-core, SSD-backed linux VM from https://www.denbi.de/: Read-onlysingle-threaded, new pipeline: | Image | Time (s)<br>zarrs<br>rust | <br>zarr<br>python | <br>zarrs<br>python | Memory (GB)<br>zarrs<br>rust | <br>zarr<br>python | <br>zarrs<br>python |
|:-----------------------------------|----------------------------:|---------------------:|----------------------:|-------------------------------:|---------------------:|----------------------:|
| data/benchmark_compress_shard.zarr | 0.98 | 3.88 | 1.32 | 2.16 | 2.20 | 2.23 |fully-threaded, new pipeline: | Image | Time (s)<br>zarrs<br>rust | <br>zarr<br>python | <br>zarrs<br>python | Memory (GB)<br>zarrs<br>rust | <br>zarr<br>python | <br>zarrs<br>python |
|:-----------------------------------|----------------------------:|---------------------:|----------------------:|-------------------------------:|---------------------:|----------------------:|
| data/benchmark_compress_shard.zarr | 1.00 | 1.60 | 1.37 | 2.16 | 2.76 | 2.22 |
| Image | Time (s)<br>zarrs<br>rust | <br>zarr<br>python | <br>zarrs<br>python | Memory (GB)<br>zarrs<br>rust | <br>zarr<br>python | <br>zarrs<br>python |
|:-----------------------------------|----------------------------:|---------------------:|----------------------:|-------------------------------:|---------------------:|----------------------:|
| data/benchmark_compress_shard.zarr | 0.98 | 4.00 | 1.37 | 2.16 | 3.97 | 2.24 |Roundtripsingle-threaded, new pipeline: | Image | Time (s)<br>zarrs<br>rust | <br>zarr<br>python | <br>zarrs<br>python | Memory (GB)<br>zarrs<br>rust | <br>zarr<br>python | <br>zarrs<br>python |
|:-----------------------------------|----------------------------:|---------------------:|----------------------:|-------------------------------:|---------------------:|----------------------:|
| data/benchmark_compress_shard.zarr | 1.49 | 16.32 | 1.94 | 0.22 | 2.20 | 2.24 |fully-threaded, new pipeline: | Image | Time (s)<br>zarrs<br>rust | <br>zarr<br>python | <br>zarrs<br>python | Memory (GB)<br>zarrs<br>rust | <br>zarr<br>python | <br>zarrs<br>python |
|:-----------------------------------|----------------------------:|---------------------:|----------------------:|-------------------------------:|---------------------:|----------------------:|
| data/benchmark_compress_shard.zarr | 1.52 | 6.07 | 2.06 | 0.22 | 2.79 | 2.24 |
| Image | Time (s)<br>zarrs<br>rust | <br>zarr<br>python | <br>zarrs<br>python | Memory (GB)<br>zarrs<br>rust | <br>zarr<br>python | <br>zarrs<br>python |
|:-----------------------------------|----------------------------:|---------------------:|----------------------:|-------------------------------:|---------------------:|----------------------:|
| data/benchmark_compress_shard.zarr | 1.44 | 14.55 | 1.97 | 0.22 | 5.35 | 2.24 |Smaller subchunks than Presented Here (i.e.,
|
…r-python into perf/prepared-write-v2
* refactor codecs * threading + benchmark cleanups
…r-python into perf/prepared-write-v2
Soften the overstated 'default pipeline is unchanged' claim to note it holds for standard configurations, and add a bugfix fragment for the sharding inner-codec spec-evolution fix (zarr-developers#2179) that also runs on the default BatchedCodecPipeline. Assisted-by: ClaudeCode:claude-opus-4.8
…r-python into perf/prepared-write-v2
ilan-gold
left a comment
There was a problem hiding this comment.
This PR should contain no breaking changes and simply add a new pipeline option that is more performant to the current stack. That pipeline includes a bunch of desired features that have been sitting around in our heads for a while and are now materialized, centered around hte notion that "sync should be sync". It has been tested downstream in the CI here and in scverse/anndata#2529 where everything is passing that would be expected to pass (we have a flag that detects pre-releases so that's what's failing there).
In short, I am confident this PR is ready for prime-time!
|
thanks for your help here ilan! this is getting merged. I will follow-up with some docs / QOL changes in separate PRs |
perf: synchronous codec pipeline (
FusedCodecPipeline), opt-inSummary
This PR adds
FusedCodecPipeline, an opt-in codec pipeline. It removes the per-chunk async scheduling overhead of the defaultBatchedCodecPipeline— roughly one coroutine per chunk, which on real workloads dominates the actual codec work — by running codec compute synchronously and in bulk.The default pipeline is unchanged for standard configurations: existing users see no behavior change unless they opt in. One caveat — sharded arrays whose inner codec chain changes the dtype now evolve that chain correctly (the #2179 fix under Notable internal changes); that is a bugfix and it also runs on the default pipeline. Standard inner chains (
[BytesCodec],[BytesCodec, ZstdCodec/BloscCodec], transpose + bytes) are byte-identical to before. We want to move slowly here — let early adopters exercise the new pipeline before considering it as a default.What it does
SupportsSyncCodec, aChunkTransformruns the codec chain synchronously — no event loop, no per-chunk coroutine — optionally across a thread pool for CPU-heavy decode/encode.Store.get_ranges_sync, and dense, fixed-size, uncompressed shards take a vectorized whole-shard bulk decode. Sharded writes go through the codec's synchronous full-shard-rewrite path.ZipStore), or a codec isn't sync-capable, the pipeline falls back to the async path — behavior-equivalent toBatchedCodecPipeline.IO ownership is unchanged: the sharding codec still holds the byte getter/setter and reads/writes storage directly (the same model as before; this PR does not move toward storage-free codecs).
Performance
Large speedups for sharded arrays — up to ~24× writes / ~14× reads on many-chunks-per-shard layouts, more with compression — and no regressions on compute-bound workloads.
Opting in
codec_pipeline.pathis unchanged (BatchedCodecPipeline). Enable the new pipeline with:codec_pipeline.max_workersdefaults toNone, which meansFusedCodecPipelineruns fully threaded (os.cpu_count()workers). It is the only pipeline that reads this setting —BatchedCodecPipelineignores it entirely. Set it to1for sequential, single-threaded execution (often faster for memory-backed stores, where store reads are GIL-locked and thread-pool overhead dominates), or to a fixed integer to cap the pool.These config additions are non-breaking: the new pipeline and its threading only take effect once you opt in via
codec_pipeline.path. The default pipeline never readsmax_workersand never spins up a thread pool.Public API additions
SyncByteGetter/SyncByteSetter— runtime-checkable protocols letting custom byte getters/setters opt into the sync fast path for in-memory IO (used by the sharding codec for its inner chunks).Store.get_ranges_sync— synchronous, coalescing counterpart ofget_ranges, with the same coalescing policy.Correctness
Both pipelines are held to the same behavior, with new equivalence/parity suites asserting that each fast path produces results identical to the general path it bypasses:
tests/test_pipeline_parity.py— Fused vs Batched parity across sharded/unsharded, compressed/uncompressed, read/write/overwrite scenarios.tests/test_codec_pipeline_suite.py— a shared suite both pipelines must pass (incl. zarr v2 + nested sharding).tests/test_fastpath_equivalence.py— property tests pinning each fast path (complete-chunk merge view, bulk shard decode, scalar-broadcast write, byte-range coalescing) equal to the general path.tests/test_fused_pipeline.py— Fused-specific behavior incl. thread-pool and async-fallback paths.Three fast-path correctness issues were found, reproduced, and fixed with regression tests: bulk shard decode could serve a reordering
vindex/oindexselection in natural order on uncompressed crc-free shards (now gated onsel_shape); theChunkTransformspec cache was not torn-read-safe undermax_workers > 1(now a single atomically-assigned entry); and sharded reads with variable-length / numcodecs inner codecs crashed becauseCodec.is_fixed_sizehad no default (now defaults toFalse).Notable internal changes
src/zarr/core/chunk_utils.py(merge/encode, decode/scatter, empty-chunk normalization, spec evolution) used by both pipelines and the sharding codec, so the sync and async paths can't drift.V2Codecand the numcodecs wrappers gained_decode_sync/_encode_synccores; the async methods are now thinasyncio.to_threadwrappers.codec_pipeline.path(restoring Narrow JSON type, ensure thatto_dictalways returns a dict, and v2 filter / compressor parsing #2179 behavior) and is evolved/memoized per spec.Lineage
Builds on prior work in #3719. Depends on changes from #3907 and #3908. Mostly authored with Claude.