diff --git a/changes/4194.bugfix.md b/changes/4194.bugfix.md new file mode 100644 index 0000000000..21a4924664 --- /dev/null +++ b/changes/4194.bugfix.md @@ -0,0 +1,10 @@ +`FusedCodecPipeline` no longer runs chunk IO and codec compute on the thread +driving zarr's internal event loop. Previously each read/write executed its +synchronous fast path inline on that loop thread, and because every sync-API +call from every user thread is serviced by the same loop, concurrent +operations serialized behind each other's codec work — reported as the fused +pipeline being slower than `BatchedCodecPipeline` for zstd-compressed data +under multi-threaded (e.g. dask) access. The synchronous batch now runs on a +worker thread (one hop per batch, not per chunk), keeping the loop free. +Multi-threaded single-chunk reads of zstd data are ~4.5x faster than before +and now scale with reader threads; single-threaded performance is unchanged. diff --git a/src/zarr/core/codec_pipeline.py b/src/zarr/core/codec_pipeline.py index 2e3f1ed122..4b8831bc7b 100644 --- a/src/zarr/core/codec_pipeline.py +++ b/src/zarr/core/codec_pipeline.py @@ -1182,7 +1182,15 @@ async def read( (isinstance(first_bg, StorePath) and isinstance(first_bg.store, SupportsGetSync)) or (not isinstance(first_bg, StorePath) and isinstance(first_bg, SyncByteGetter)) ): - return self.read_sync(batch, out, drop_axes, max_workers=_resolve_max_workers()) + # One thread hop for the WHOLE batch — not per chunk, so the fused + # design's win over per-chunk async scheduling is preserved. Running + # read_sync inline here would block the event loop for the duration + # of the batch's IO+compute; every sync-API call from every user + # thread shares this one loop, so inline execution serializes + # concurrent callers behind each other's codec compute. + return await asyncio.to_thread( + self.read_sync, batch, out, drop_axes, max_workers=_resolve_max_workers() + ) # Non-sync store (e.g. ZipStore): can't use the sync fast path. But if the # array-bytes codec supports partial decoding (sharding), still route @@ -1235,7 +1243,11 @@ async def write( (isinstance(first_bs, StorePath) and isinstance(first_bs.store, SupportsSetSync)) or (not isinstance(first_bs, StorePath) and isinstance(first_bs, SyncByteSetter)) ): - self.write_sync(batch, value, drop_axes, max_workers=_resolve_max_workers()) + # One thread hop for the whole batch; see the matching comment in + # `read` for why write_sync must not run inline on the event loop. + await asyncio.to_thread( + self.write_sync, batch, value, drop_axes, max_workers=_resolve_max_workers() + ) return await _async_write_fallback(self, batch, value, drop_axes) diff --git a/tests/benchmarks/test_e2e.py b/tests/benchmarks/test_e2e.py index 9d60d9a2fb..487485e262 100644 --- a/tests/benchmarks/test_e2e.py +++ b/tests/benchmarks/test_e2e.py @@ -190,3 +190,51 @@ def setup() -> tuple[tuple[zarr.Array, EllipsisType], dict]: # type: ignore[typ return (arr, Ellipsis), {} benchmark.pedantic(getitem, setup=setup, rounds=3) # type: ignore[no-untyped-call] + + +_CONCURRENT_READ_THREADS = 8 +_concurrent_layout = Layout(shape=(64_000_000,), chunks=(4_000_000,), shards=None) + + +@pytest.mark.parametrize("pipeline", ["batched", "fused_full_threaded"], indirect=True) +@pytest.mark.parametrize("compression_name", ["zstd", None]) +@pytest.mark.parametrize("store", ["local"], indirect=["store"]) +def test_read_array_concurrent( + bench_store: Store, + compression_name: CompressorName, + pipeline: str, + benchmark: BenchmarkFixture, +) -> None: + """Dask-style access: several user threads each reading one chunk per call. + + All sync-API calls are serviced by the one global event loop, so this + measures how much of each read's IO+compute the pipeline runs while + holding the loop: anything inline serializes the readers. + """ + from concurrent.futures import ThreadPoolExecutor + + layout = _concurrent_layout + arr = create_array( + bench_store, + dtype="uint8", + shape=layout.shape, + chunks=layout.chunks, + shards=layout.shards, + compressors=compressors[compression_name], # type: ignore[arg-type] + fill_value=0, + ) + arr[:] = _data(layout.shape) + selections = [ + slice(start, start + layout.chunks[0]) + for start in range(0, layout.shape[0], layout.chunks[0]) + ] + + def read_all_chunks_concurrently() -> None: + with ThreadPoolExecutor(max_workers=_CONCURRENT_READ_THREADS) as executor: + list(executor.map(lambda sel: arr[sel], selections)) + + def setup() -> tuple[tuple[()], dict]: # type: ignore[type-arg] + clear_cache() + return (), {} + + benchmark.pedantic(read_all_chunks_concurrently, setup=setup, rounds=3) # type: ignore[no-untyped-call] diff --git a/tests/test_fused_pipeline.py b/tests/test_fused_pipeline.py index 73c2c6e1c3..09e9241c07 100644 --- a/tests/test_fused_pipeline.py +++ b/tests/test_fused_pipeline.py @@ -34,6 +34,67 @@ def test_construction(codecs: tuple[Any, ...]) -> None: assert pipeline.codecs == codecs +def test_sync_api_compute_off_event_loop(monkeypatch: pytest.MonkeyPatch) -> None: + """Codec compute must never run on the thread driving the event loop. + + Every sync-API call, from every user thread, is serviced by the one global + `zarr_io` event loop. Running decode/encode inline on that loop's thread + turns it into a mutex around codec compute: concurrent readers serialize, + and the penalty grows with codec cost (observed as "fused pipeline is + slower for zstd data" under dask-style multi-threaded single-chunk reads). + """ + import asyncio + + from zarr.core.chunk_utils import ChunkTransform + + compute_on_loop = {"decode": False, "encode": False} + calls = {"decode": 0, "encode": 0} + real_decode = ChunkTransform.decode_chunk + real_encode = ChunkTransform.encode_chunk + + def _running_loop() -> bool: + try: + asyncio.get_running_loop() + except RuntimeError: + return False + return True + + def traced_decode(self: ChunkTransform, chunk_bytes: Any, chunk_spec: Any) -> Any: + calls["decode"] += 1 + compute_on_loop["decode"] = compute_on_loop["decode"] or _running_loop() + return real_decode(self, chunk_bytes, chunk_spec) + + def traced_encode(self: ChunkTransform, chunk_array: Any, chunk_spec: Any) -> Any: + calls["encode"] += 1 + compute_on_loop["encode"] = compute_on_loop["encode"] or _running_loop() + return real_encode(self, chunk_array, chunk_spec) + + monkeypatch.setattr(ChunkTransform, "decode_chunk", traced_decode) + monkeypatch.setattr(ChunkTransform, "encode_chunk", traced_encode) + + with zarr_config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}): + arr = zarr.create_array( + MemoryStore(), + shape=(8, 8), + chunks=(4, 4), + dtype="float64", + compressors=ZstdCodec(level=1), + ) + data = np.arange(64, dtype="float64").reshape(8, 8) + arr[:4, :4] = data[:4, :4] # single-chunk write (batch of 1) + arr[:] = data # multi-chunk write + # Guard against vacuity: if the sync fast path stops triggering, the + # traced ChunkTransform methods are never called (the async fallback + # uses AsyncChunkTransform) and the on-loop flags stay trivially False. + assert calls["encode"] > 0 + assert compute_on_loop["encode"] is False + + np.testing.assert_array_equal(arr[:4, :4], data[:4, :4]) # single-chunk read + np.testing.assert_array_equal(arr[:], data) # multi-chunk read + assert calls["decode"] > 0 + assert compute_on_loop["decode"] is False + + def test_evolve_from_array_spec() -> None: """evolve_from_array_spec creates a sync transform.""" from zarr.core.array_spec import ArrayConfig, ArraySpec