From eefa424df0e7cfd1683ca7fb27996384141fea72 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 18:11:51 +0200 Subject: [PATCH 1/3] chore(deps): bump the actions group with 3 updates (#4181) Bumps the actions group with 3 updates: [actions/labeler](https://github.com/actions/labeler), [actions/attest](https://github.com/actions/attest) and [zizmorcore/zizmor-action](https://github.com/zizmorcore/zizmor-action). Updates `actions/labeler` from 6.2.0 to 7.0.0 - [Release notes](https://github.com/actions/labeler/releases) - [Commits](https://github.com/actions/labeler/compare/b8dd2d9be0f68b860e7dae5dae7d772984eacd6d...bf12e9b00b37c5c0ca2b87b79b2daf7891dbda13) Updates `actions/attest` from 4.1.1 to 4.2.0 - [Release notes](https://github.com/actions/attest/releases) - [Changelog](https://github.com/actions/attest/blob/main/RELEASE.md) - [Commits](https://github.com/actions/attest/compare/a1948c3f048ba23858d222213b7c278aabede763...f7c74d28b9d84cb8768d0b8ca14a4bac6ef463e6) Updates `zizmorcore/zizmor-action` from 0.5.7 to 0.6.0 - [Release notes](https://github.com/zizmorcore/zizmor-action/releases) - [Commits](https://github.com/zizmorcore/zizmor-action/compare/192e21d79ab29983730a13d1382995c2307fbcaa...6599ee8b7a49aef6a770f63d261d214911a7ce02) --- updated-dependencies: - dependency-name: actions/labeler dependency-version: 7.0.0 dependency-type: direct:production update-type: version-update:semver-major dependency-group: actions - dependency-name: actions/attest dependency-version: 4.2.0 dependency-type: direct:production update-type: version-update:semver-minor dependency-group: actions - dependency-name: zizmorcore/zizmor-action dependency-version: 0.6.0 dependency-type: direct:production update-type: version-update:semver-minor dependency-group: actions ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/needs_release_notes.yml | 2 +- .github/workflows/releases.yml | 2 +- .github/workflows/zarr-metadata-release.yml | 4 ++-- .github/workflows/zizmor.yml | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/needs_release_notes.yml b/.github/workflows/needs_release_notes.yml index fa555d1478..e001e8cd43 100644 --- a/.github/workflows/needs_release_notes.yml +++ b/.github/workflows/needs_release_notes.yml @@ -21,7 +21,7 @@ jobs: pull-requests: write # Required to add labels to PRs runs-on: ubuntu-latest steps: - - uses: actions/labeler@b8dd2d9be0f68b860e7dae5dae7d772984eacd6d # v6.2.0 + - uses: actions/labeler@bf12e9b00b37c5c0ca2b87b79b2daf7891dbda13 # v7.0.0 with: repo-token: ${{ secrets.GITHUB_TOKEN }} sync-labels: true diff --git a/.github/workflows/releases.yml b/.github/workflows/releases.yml index 4d460f4a56..fe0d09f300 100644 --- a/.github/workflows/releases.yml +++ b/.github/workflows/releases.yml @@ -81,7 +81,7 @@ jobs: name: releases path: dist - name: Generate artifact attestation - uses: actions/attest@a1948c3f048ba23858d222213b7c278aabede763 # v4.1.1 + uses: actions/attest@f7c74d28b9d84cb8768d0b8ca14a4bac6ef463e6 # v4.2.0 with: subject-path: dist/* - name: Publish package to PyPI diff --git a/.github/workflows/zarr-metadata-release.yml b/.github/workflows/zarr-metadata-release.yml index 5021f79d2e..bc9ecf9871 100644 --- a/.github/workflows/zarr-metadata-release.yml +++ b/.github/workflows/zarr-metadata-release.yml @@ -82,7 +82,7 @@ jobs: path: dist - name: Generate artifact attestation - uses: actions/attest@a1948c3f048ba23858d222213b7c278aabede763 # v4.1.1 + uses: actions/attest@f7c74d28b9d84cb8768d0b8ca14a4bac6ef463e6 # v4.2.0 with: subject-path: dist/* @@ -107,7 +107,7 @@ jobs: path: dist - name: Generate artifact attestation - uses: actions/attest@a1948c3f048ba23858d222213b7c278aabede763 # v4.1.1 + uses: actions/attest@f7c74d28b9d84cb8768d0b8ca14a4bac6ef463e6 # v4.2.0 with: subject-path: dist/* diff --git a/.github/workflows/zizmor.yml b/.github/workflows/zizmor.yml index 6250426bae..1567bea713 100644 --- a/.github/workflows/zizmor.yml +++ b/.github/workflows/zizmor.yml @@ -32,4 +32,4 @@ jobs: persist-credentials: false - name: Run zizmor - uses: zizmorcore/zizmor-action@192e21d79ab29983730a13d1382995c2307fbcaa # v0.5.7 + uses: zizmorcore/zizmor-action@6599ee8b7a49aef6a770f63d261d214911a7ce02 # v0.6.0 From 0b727571268c8962717ab5c23965c4642fe6bd36 Mon Sep 17 00:00:00 2001 From: Joe Hamman Date: Mon, 27 Jul 2026 19:31:14 -0700 Subject: [PATCH 2/3] feat: ZipStore accepts open binary file-like objects (#4187) Allows constructing a ZipStore from any seekable binary reader, enabling zip archives on remote storage: - io objects (BytesIO, fsspec file objects) are used directly - minimal readers that are not io.IOBase instances and whose read() may return buffer-protocol objects rather than bytes (e.g. obstore.ReadableFile) are adapted via a small io.RawIOBase wrapper when opened for reading clear()/move() raise NotImplementedError for file-object-backed stores. Co-authored-by: Claude Fable 5 --- changes/4187.feature.md | 4 + docs/user-guide/storage.md | 13 +++ src/zarr/storage/_zip.py | 114 +++++++++++++++++++++++-- tests/test_store/test_zip.py | 159 +++++++++++++++++++++++++++++++++++ 4 files changed, 281 insertions(+), 9 deletions(-) create mode 100644 changes/4187.feature.md diff --git a/changes/4187.feature.md b/changes/4187.feature.md new file mode 100644 index 0000000000..87133e2034 --- /dev/null +++ b/changes/4187.feature.md @@ -0,0 +1,4 @@ +`ZipStore` now accepts an open binary file-like object in place of a path, enabling +zip archives on remote storage (e.g. a file opened with `fsspec` or an +`obstore.ReadableFile`). Operations that require a filesystem location +(`clear`, `move`) raise `NotImplementedError` for file-object-backed stores. diff --git a/docs/user-guide/storage.md b/docs/user-guide/storage.md index 7e0154b2a0..0ba6202c76 100644 --- a/docs/user-guide/storage.md +++ b/docs/user-guide/storage.md @@ -124,6 +124,19 @@ array = zarr.create_array(store=store, shape=(2,), dtype='float64') print(array) ``` +In place of a path, `ZipStore` also accepts an open binary file object (for +example a file opened with `fsspec`, or an `obstore` reader), enabling zip +archives on remote storage. The file must stay open for as long as the store +is in use: + +```python exec="true" session="storage" source="above" result="ansi" +store.close() +f = open('data.zip', mode='rb') # must stay open while the store is used +array = zarr.open_array(store=zarr.storage.ZipStore(f), mode='r') +print(array[:]) +f.close() +``` + ### Remote Store The [`zarr.storage.FsspecStore`][] stores the contents of a Zarr hierarchy following the same diff --git a/src/zarr/storage/_zip.py b/src/zarr/storage/_zip.py index 430b0c3e2a..69ae18bc2c 100644 --- a/src/zarr/storage/_zip.py +++ b/src/zarr/storage/_zip.py @@ -1,12 +1,13 @@ from __future__ import annotations +import io import os import shutil import threading import time import zipfile from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal +from typing import IO, TYPE_CHECKING, Any, Literal from zarr.abc.store import ( ByteRequest, @@ -23,14 +24,67 @@ ZipStoreAccessModeLiteral = Literal["r", "w", "a"] +class _RawReaderAdapter(io.RawIOBase): + """ + Adapt a minimal seekable reader to the `io` interface `zipfile` needs. + + Some file-like objects (e.g. `obstore.ReadableFile`) implement + `read`/`seek`/`tell` but are not `io.IOBase` instances, and their + `read` may return a buffer-protocol object rather than `bytes`. + Wrapping in this adapter plus `io.BufferedReader` yields real `bytes`. + + Reads are clamped to the bytes remaining before EOF: some readers + (obstore < 0.6) raise on short reads rather than returning fewer bytes. + The size is cached, which is safe because the adapter is only used for + read-only access. + """ + + def __init__(self, fileobj: IO[bytes]) -> None: + self._fileobj = fileobj + self._size: int | None = None + + def _get_size(self) -> int: + if self._size is None: + pos = self._fileobj.tell() + self._size = self._fileobj.seek(0, os.SEEK_END) + self._fileobj.seek(pos) + return self._size + + def readable(self) -> bool: + return True + + def seekable(self) -> bool: + return True + + def seek(self, pos: int, whence: int = 0) -> int: + return self._fileobj.seek(pos, whence) + + def tell(self) -> int: + return self._fileobj.tell() + + def readinto(self, b: Any) -> int: + n_requested = min(len(b), self._get_size() - self._fileobj.tell()) + if n_requested <= 0: + return 0 + data = self._fileobj.read(n_requested) + n = len(data) + b[:n] = memoryview(data) + return n + + class ZipStore(Store): """ Store using a ZIP file. Parameters ---------- - path : str - Location of file. + path : str, Path, or IO[bytes] + Location of file, or an open binary file object. A file object must + support `read`, `seek`, and `tell`; objects that are not `io.IOBase` + instances (e.g. an `obstore` reader) are adapted automatically but + can only be used for reading (`mode="r"`). The file object must stay + open for the lifetime of the store, and operations that require a + filesystem location (`clear`, `move`, pickling) are not supported. mode : str, optional One of 'r' to read an existing file, 'w' to truncate and write a new file, 'a' to append to an existing file, or 'x' to exclusively create @@ -58,16 +112,17 @@ class ZipStore(Store): supports_deletes: bool = False supports_listing: bool = True - path: Path + path: Path | None compression: int allowZip64: bool _zf: zipfile.ZipFile _lock: threading.RLock + _fileobj: IO[bytes] | None def __init__( self, - path: Path | str, + path: Path | str | IO[bytes], *, mode: ZipStoreAccessModeLiteral = "r", read_only: bool | None = None, @@ -81,8 +136,28 @@ def __init__( if isinstance(path, str): path = Path(path) - assert isinstance(path, Path) - self.path = path # root? + if isinstance(path, Path): + self.path = path # root? + self._fileobj = None + else: + self.path = None + if not isinstance(path, io.IOBase): + if not all( + callable(getattr(path, attr, None)) for attr in ("read", "seek", "tell") + ): + raise TypeError( + f"expected a path or an open binary file object supporting " + f"read/seek/tell, got {type(path).__name__}" + ) + if mode != "r": + raise TypeError( + f"a file object that is not an io.IOBase instance can only be " + f"opened for reading (mode='r', got mode={mode!r})" + ) + # e.g. an obstore ReadableFile: readable and seekable, but + # not an io object and reads may not return bytes + path = io.BufferedReader(_RawReaderAdapter(path)) + self._fileobj = path self._zmode = mode self.compression = compression @@ -95,7 +170,7 @@ def _sync_open(self) -> None: self._lock = threading.RLock() self._zf = zipfile.ZipFile( - self.path, + self.path if self.path is not None else self._fileobj, # type: ignore[arg-type] mode=self._zmode, compression=self.compression, allowZip64=self.allowZip64, @@ -107,6 +182,13 @@ async def _open(self) -> None: self._sync_open() def __getstate__(self) -> dict[str, Any]: + if self.path is None: + # A path-backed store pickles its path and reopens the file on + # unpickling; an open file object cannot be serialized that way. + raise TypeError( + "cannot pickle a ZipStore backed by a file-like object; " + "construct the store from a path instead" + ) # We need a copy to not modify the state of the original store state = self.__dict__.copy() for attr in ["_zf", "_lock"]: @@ -130,6 +212,10 @@ async def clear(self) -> None: # docstring inherited with self._lock: self._check_writable() + if self.path is None: + raise NotImplementedError( + "clear() is not supported for a ZipStore backed by a file-like object" + ) self._zf.close() os.remove(self.path) self._zf = zipfile.ZipFile( @@ -137,13 +223,19 @@ async def clear(self) -> None: ) def __str__(self) -> str: + if self.path is None: + return f"zip://{self._fileobj!r}" return f"zip://{self.path}" def __repr__(self) -> str: return f"ZipStore('{self}')" def __eq__(self, other: object) -> bool: - return isinstance(other, type(self)) and self.path == other.path + return ( + isinstance(other, type(self)) + and self.path == other.path + and self._fileobj is other._fileobj + ) def _get( self, @@ -297,6 +389,10 @@ async def move(self, path: Path | str) -> None: """ Move the store to another path. """ + if self.path is None: + raise NotImplementedError( + "move() is not supported for a ZipStore backed by a file-like object" + ) if isinstance(path, str): path = Path(path) self.close() diff --git a/tests/test_store/test_zip.py b/tests/test_store/test_zip.py index ed69114b51..0d8dadd18a 100644 --- a/tests/test_store/test_zip.py +++ b/tests/test_store/test_zip.py @@ -1,6 +1,8 @@ from __future__ import annotations +import io import os +import pickle import shutil import tempfile import zipfile @@ -188,6 +190,163 @@ async def test_move(self, tmp_path: Path) -> None: assert np.array_equal(array[...], np.arange(10)) +class TestZipStoreFileObj: + """ZipStore backed by an open binary file-like object instead of a path.""" + + @pytest.fixture + def zip_bytes(self, tmp_path: Path) -> bytes: + path = tmp_path / "data.zip" + store = ZipStore(path, mode="w") + zarr.create_array(store, data=np.arange(10), chunks=(5,)) + store.close() + return path.read_bytes() + + def test_read_from_fileobj(self, zip_bytes: bytes) -> None: + # an existing archive can be read through any seekable binary reader + store = ZipStore(io.BytesIO(zip_bytes), mode="r") + array = zarr.open_array(store, mode="r") + assert np.array_equal(array[...], np.arange(10)) + assert store.path is None + + def test_write_to_fileobj(self) -> None: + # a writable file object receives the archive; the bytes it holds + # after close() are a complete, reopenable zip + buffer = io.BytesIO() + store = ZipStore(buffer, mode="w", read_only=False) + zarr.create_array(store, data=np.arange(4)) + store.close() + + roundtrip = ZipStore(io.BytesIO(buffer.getvalue()), mode="r") + array = zarr.open_array(roundtrip, mode="r") + assert np.array_equal(array[...], np.arange(4)) + + async def test_clear_unsupported(self, zip_bytes: bytes) -> None: + # clear() requires a filesystem location, so it raises a clear error + # for file-object-backed stores + store = ZipStore(io.BytesIO(zip_bytes), mode="a", read_only=False) + store._sync_open() + with pytest.raises(NotImplementedError, match="clear.*file-like"): + await store.clear() + + async def test_move_unsupported(self, zip_bytes: bytes) -> None: + # move() requires a filesystem location, so it raises a clear error + # for file-object-backed stores + store = ZipStore(io.BytesIO(zip_bytes), mode="a", read_only=False) + store._sync_open() + with pytest.raises(NotImplementedError, match="move.*file-like"): + await store.move("elsewhere.zip") + + def test_invalid_file_object_rejected(self) -> None: + # objects without read/seek/tell are rejected at construction, not + # deep inside zipfile + with pytest.raises(TypeError, match="read/seek/tell"): + ZipStore(42, mode="r") # type: ignore[arg-type] + + @pytest.mark.parametrize("mode", ["w", "a", "x"]) + def test_non_iobase_reader_write_modes_rejected(self, zip_bytes: bytes, mode: str) -> None: + # readers that are not io.IOBase instances are adapted for reading + # only; write modes are rejected at construction with a clear error + class MinimalReader: + def __init__(self, data: bytes) -> None: + self._buffer = io.BytesIO(data) + + def read(self, size: int, /) -> bytes: + return self._buffer.read(size) + + def seek(self, pos: int, whence: int = 0, /) -> int: + return self._buffer.seek(pos, whence) + + def tell(self) -> int: + return self._buffer.tell() + + with pytest.raises(TypeError, match="opened for reading"): + ZipStore(MinimalReader(zip_bytes), mode=mode, read_only=False) # type: ignore[arg-type] + + def test_fsspec_file(self, tmp_path: Path, zip_bytes: bytes) -> None: + # a file opened through fsspec (already an io.IOBase) is used directly; + # fsspec's local filesystem stands in for a remote one + fsspec = pytest.importorskip("fsspec") + + path = tmp_path / "fsspec.zip" + path.write_bytes(zip_bytes) + with fsspec.open(f"local://{path}", "rb") as fileobj: + store = ZipStore(fileobj, mode="r") + array = zarr.open_array(store, mode="r") + assert np.array_equal(array[...], np.arange(10)) + assert store.path is None + + def test_obstore_reader(self, tmp_path: Path, zip_bytes: bytes) -> None: + # obstore's ReadableFile is not an io.IOBase and its read() returns a + # buffer-protocol object; ZipStore adapts it via _RawReaderAdapter + obstore = pytest.importorskip("obstore") + from obstore.store import LocalStore as ObstoreLocalStore + + (tmp_path / "obstore.zip").write_bytes(zip_bytes) + reader = obstore.open_reader(ObstoreLocalStore(str(tmp_path)), "obstore.zip") + store = ZipStore(reader, mode="r") + array = zarr.open_array(store, mode="r") + assert np.array_equal(array[...], np.arange(10)) + + def test_raw_reader_adapter_eof(self) -> None: + from zarr.storage._zip import _RawReaderAdapter + + class MinimalReader: + """Non-io.IOBase reader exposing only read/seek/tell, like obstore.""" + + def __init__(self, data: bytes) -> None: + self._buffer = io.BytesIO(data) + + def read(self, size: int, /) -> bytes: + return self._buffer.read(size) + + def seek(self, pos: int, whence: int = 0, /) -> int: + return self._buffer.seek(pos, whence) + + def tell(self) -> int: + return self._buffer.tell() + + # the adapter must clamp reads to EOF: some readers (obstore < 0.6) + # raise on short reads instead of returning fewer bytes + data = b"0123456789" + adapter = _RawReaderAdapter(MinimalReader(data)) # type: ignore[arg-type] + + # A read straddling EOF returns only the remaining bytes. + adapter.seek(len(data) - 3) + buf = bytearray(8) + assert adapter.readinto(buf) == 3 + assert bytes(buf[:3]) == data[-3:] + + # A read at EOF returns 0. + assert adapter.tell() == len(data) + assert adapter.readinto(bytearray(8)) == 0 + + def test_pickle_fileobj_raises(self, zip_bytes: bytes) -> None: + # an open file object cannot be reliably serialized, so pickling a + # file-object-backed store raises with a pointer at the alternative + store = ZipStore(io.BytesIO(zip_bytes), mode="r") + with pytest.raises(TypeError, match="cannot pickle a ZipStore backed by a file-like"): + pickle.dumps(store) + + def test_pickle_path_backed_roundtrip(self, tmp_path: Path, zip_bytes: bytes) -> None: + # path-backed stores remain picklable: the path is serialized and the + # archive is reopened on unpickling + path = tmp_path / "pickled.zip" + path.write_bytes(zip_bytes) + store = ZipStore(path, mode="r") + unpickled = pickle.loads(pickle.dumps(store)) + array = zarr.open_array(unpickled, mode="r") + assert np.array_equal(array[...], np.arange(10)) + + def test_str_and_eq(self, zip_bytes: bytes) -> None: + # file-object-backed stores stringify with the object repr and + # compare equal only when backed by the very same file object + fileobj = io.BytesIO(zip_bytes) + store = ZipStore(fileobj, mode="r") + assert str(store).startswith("zip://<") + assert store == ZipStore(fileobj, mode="r") + assert store != ZipStore(io.BytesIO(zip_bytes), mode="r") + + class ZipStoreLifecycleMachine(RuleBasedStateMachine): """Drive a ZipStore through construct / open / write / close transitions. From aa766b41bbd4135e5ae701c615503ef16b64647a Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 28 Jul 2026 07:13:12 +0200 Subject: [PATCH 3/3] fix(codec_pipeline): keep FusedCodecPipeline compute off the event-loop thread MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit FusedCodecPipeline.read/write ran their synchronous fast path inline on the coroutine servicing the request — i.e. on the global zarr_io event loop thread. Single-chunk batches decoded inline on the loop and multi-chunk batches blocked the loop in pool.map, so every sync-API call from every user thread serialized behind each other's codec compute. The blocked window scales with codec cost, which is why users reported the fused pipeline as "slower for zstd-compressed data" under multi-threaded (dask-style, one chunk per call) access: at 8 reader threads on 4 MiB zstd chunks it was 3.4x slower than BatchedCodecPipeline, and throughput did not scale with threads at all (336 -> 439 ms from 1 to 8 threads, versus 669 -> 121 ms for batched). Offload the synchronous batch to a worker thread with asyncio.to_thread: one hop per batch, not per chunk, preserving the fused pipeline's win over per-chunk async scheduling while keeping the loop free. After the fix the same workload scales 625 -> 109 ms from 1 to 8 threads, beating batched at every thread count; single-threaded performance is unchanged (the hop costs ~75 us per batch). The regression test asserts deterministically (no timing) that codec compute never runs on a thread with a running event loop, covering single- and multi-chunk reads and writes through the sync API. A new benchmark covers the many-threads/one-chunk-per-call access pattern. Assisted-by: ClaudeCode:claude-fable-5 --- changes/247.bugfix.md | 10 +++++++ src/zarr/core/codec_pipeline.py | 16 ++++++++-- tests/benchmarks/test_e2e.py | 48 +++++++++++++++++++++++++++++ tests/test_fused_pipeline.py | 53 +++++++++++++++++++++++++++++++++ 4 files changed, 125 insertions(+), 2 deletions(-) create mode 100644 changes/247.bugfix.md diff --git a/changes/247.bugfix.md b/changes/247.bugfix.md new file mode 100644 index 0000000000..21a4924664 --- /dev/null +++ b/changes/247.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..70eaee0398 100644 --- a/tests/test_fused_pipeline.py +++ b/tests/test_fused_pipeline.py @@ -34,6 +34,59 @@ 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} + 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: + 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: + 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 + 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 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