From 4a9c8e405a6b295520c3f9276083a741532aa780 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 29 Jul 2026 13:29:48 +0200 Subject: [PATCH] fix: gate fused sync fast paths on full store sync capability; wrappers forward it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fused write gate checked only SupportsSetSync, but write_sync also needs get_sync (partial-chunk read-modify-write) and delete_sync (all-fill chunk cleanup): a set-sync-only store passed the gate, wrote some chunks, then died mid-batch with TypeError. And WrapperStore forwarded no *_sync method, so every wrapped store (e.g. LatencyStore) silently lost the sync fast path — latency benchmarks measured the async fallback while claiming to measure the fused sync path. Both gates now consult _store_supports_sync_io: structural membership in SupportsSyncStore (the full get/set/delete sync surface) combined with a per-instance _supports_sync_io opt-out (absent means capable). This is a private, interim convention pending a formal sync/async store architecture — the store-side twin of the codec-side _sync_capable convention from #4179 — deliberately not new public API. WrapperStore delegates the three sync methods and forwards the wrapped store's capability, so wrapping a sync store keeps the fast path and wrapping an async-only store falls back cleanly; LoggingStore logs the delegated sync calls. LatencyStore fixes: sync reads/writes now sleep the configured latency on the worker thread; get_ranges/get_partial_values route through the latency-injecting get instead of bypassing the wrapper; _with_store passes the raw (loc, scale) latency config instead of a single sampled float, so derived stores keep the distribution. Assisted-by: ClaudeCode:claude-fable-5 --- changes/4206.bugfix.md | 1 + src/zarr/abc/store.py | 47 +++++- src/zarr/codecs/sharding.py | 4 +- src/zarr/core/codec_pipeline.py | 32 ++-- src/zarr/experimental/cache_store.py | 10 ++ src/zarr/storage/_logging.py | 21 +++ src/zarr/storage/_wrapper.py | 42 ++++- src/zarr/testing/store.py | 80 ++++++++- tests/test_codec_pipeline_suite.py | 18 ++- tests/test_experimental/test_cache_store.py | 37 +++++ tests/test_fused_pipeline.py | 170 +++++++++++++++++++- tests/test_store/test_get_ranges.py | 8 +- tests/test_store/test_latency.py | 108 +++++++++++++ tests/test_store/test_wrapper.py | 50 +++++- 14 files changed, 596 insertions(+), 32 deletions(-) create mode 100644 changes/4206.bugfix.md diff --git a/changes/4206.bugfix.md b/changes/4206.bugfix.md new file mode 100644 index 0000000000..a01c969449 --- /dev/null +++ b/changes/4206.bugfix.md @@ -0,0 +1 @@ +Fixed `FusedCodecPipeline`'s gating of its synchronous fast paths: stores exposing only part of the sync surface (e.g. `set_sync` without `get_sync`) now fall back cleanly to the async path instead of failing mid-write, and `WrapperStore` now forwards `get_sync`/`set_sync`/`delete_sync` to the wrapped store so wrapped sync-capable stores keep the fast path. The capability decision uses a private, interim convention (`zarr.abc.store._store_supports_sync_io`) rather than new public API, pending a formal sync/async store architecture. Also fixed `LatencyStore`: synchronous reads and writes now pay the configured latency, `get_ranges`/`get_partial_values` no longer bypass latency injection, and derived stores (e.g. from `with_read_only`) keep a stochastic `(loc, scale)` latency configuration instead of freezing a single sample. diff --git a/src/zarr/abc/store.py b/src/zarr/abc/store.py index c60d2468c5..af528ec533 100644 --- a/src/zarr/abc/store.py +++ b/src/zarr/abc/store.py @@ -662,6 +662,15 @@ def delete_sync(self) -> None: ... @runtime_checkable class SupportsGetSync(Protocol): + """Store protocol for synchronous reads (`get_sync`). + + The store sync surface is all-or-nothing: a store implementing any of the + `*_sync` methods must implement all of them (`SupportsSyncStore`), because + consumers mix sync reads, writes, and deletes within one operation. + Capability-gated callers consult `_store_supports_sync_io` rather than the + individual protocols. + """ + def get_sync( self, key: str, @@ -673,16 +682,52 @@ def get_sync( @runtime_checkable class SupportsSetSync(Protocol): + """Store protocol for synchronous writes (`set_sync`). + + See `SupportsGetSync` for the all-or-nothing contract on the store sync + surface. + """ + def set_sync(self, key: str, value: Buffer) -> None: ... @runtime_checkable class SupportsDeleteSync(Protocol): + """Store protocol for synchronous deletes (`delete_sync`). + + See `SupportsGetSync` for the all-or-nothing contract on the store sync + surface. + """ + def delete_sync(self, key: str) -> None: ... @runtime_checkable -class SupportsSyncStore(SupportsGetSync, SupportsSetSync, SupportsDeleteSync, Protocol): ... +class SupportsSyncStore(SupportsGetSync, SupportsSetSync, SupportsDeleteSync, Protocol): + """The full store sync surface: `get_sync`, `set_sync`, and `delete_sync`.""" + + +def _store_supports_sync_io(store: object) -> bool: + """Whether `store` can serve the full synchronous IO surface right now. + + Structural membership in `SupportsSyncStore` is necessary but not always + sufficient: a store can present the `*_sync` methods while its ability to + run them depends on runtime state the type system cannot see. Wrapper + stores are the canonical case — `WrapperStore` delegates the sync methods + to the store it wraps, so they only work when the wrapped store is itself + sync-capable. Such stores opt out dynamically via a `_supports_sync_io` + attribute/property (absent means capable). + + This is an interim, private convention pending a formal sync/async store + architecture — the store-side twin of the codec-side `_sync_capable` + convention consulted by `zarr.abc.codec._codec_supports_sync`. + + Synchronous IO is all-or-nothing: consumers such as the fused codec + pipeline mix synchronous reads, writes, and deletes within one batch + (e.g. a partial-chunk write reads existing bytes and an all-fill chunk is + deleted), so a partial sync surface never satisfies this predicate. + """ + return isinstance(store, SupportsSyncStore) and getattr(store, "_supports_sync_io", True) async def set_or_delete(byte_setter: ByteSetter, value: Buffer | None) -> None: diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index cdfdae6c89..d8ca8bdf62 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -23,7 +23,7 @@ RangeByteRequest, Store, SuffixByteRequest, - SupportsGetSync, + _store_supports_sync_io, ) from zarr.codecs._deprecated_enum import _coerce_enum_input, _DeprecatedStrEnumMeta from zarr.codecs.bytes import BytesCodec @@ -1721,7 +1721,7 @@ def _load_partial_shard_maybe_sync( shard_dict: ShardMutableMapping = {} store = byte_getter.store if hasattr(byte_getter, "store") else None - if isinstance(store, Store) and isinstance(store, SupportsGetSync): + if isinstance(store, Store) and _store_supports_sync_io(store): # External store: coalesce via get_ranges_sync (mirrors get_ranges). byte_ranges = [byte_range for _, byte_range in chunk_coord_byte_ranges] try: diff --git a/src/zarr/core/codec_pipeline.py b/src/zarr/core/codec_pipeline.py index 92fd0970fe..597f338c42 100644 --- a/src/zarr/core/codec_pipeline.py +++ b/src/zarr/core/codec_pipeline.py @@ -331,8 +331,8 @@ async def _async_read_fallback( then scatters each decoded chunk into `out` at its `out_selection`. Used by both `BatchedCodecPipeline.read_batch` (non-partial-decode - branch) and `FusedCodecPipeline.read` (when the store is not a - `SupportsGetSync` / sync transform is unavailable). + branch) and `FusedCodecPipeline.read` (when the store does not advertise + sync IO / sync transform is unavailable). """ chunk_array_batch: list[NDBuffer | None] @@ -393,8 +393,8 @@ async def _async_write_fallback( if encoding produced `None` or the chunk dropped). Used by both `BatchedCodecPipeline.write_batch` (non-partial-encode - branch) and `FusedCodecPipeline.write` (when the store is not a - `SupportsSetSync` / sync transform is unavailable). + branch) and `FusedCodecPipeline.write` (when the store does not advertise + sync IO / sync transform is unavailable). """ if use_sync := ( @@ -1265,16 +1265,17 @@ async def read( return () # Fast path: sync transform plus synchronous IO. For StorePath the gate - # is on the STORE's sync support (StorePath always has a get_sync - # method, but it only works when its store does); for other byte - # getters (e.g. the sharding codec's in-memory _ShardingByteGetter) the - # SyncByteGetter protocol is the gate. - from zarr.abc.store import SupportsGetSync, SyncByteGetter + # is the STORE's sync-IO capability (`_store_supports_sync_io`) (StorePath always has a + # get_sync method, but it only works when its store implements the full + # sync surface); for other byte getters (e.g. the sharding codec's + # in-memory _ShardingByteGetter) the SyncByteGetter protocol is the + # gate. + from zarr.abc.store import SyncByteGetter, _store_supports_sync_io from zarr.storage._common import StorePath first_bg = batch[0][0] if self.sync_transform is not None and ( - (isinstance(first_bg, StorePath) and isinstance(first_bg.store, SupportsGetSync)) + (isinstance(first_bg, StorePath) and _store_supports_sync_io(first_bg.store)) or (not isinstance(first_bg, StorePath) and isinstance(first_bg, SyncByteGetter)) ): # One thread hop for the WHOLE batch — not per chunk, so the fused @@ -1328,14 +1329,17 @@ async def write( return # Fast path: sync transform plus synchronous IO. Mirrors `read`: gate - # StorePath on the store's sync support, other byte setters (e.g. the - # sharding codec's in-memory _ShardingByteSetter) on SyncByteSetter. - from zarr.abc.store import SupportsSetSync, SyncByteSetter + # StorePath on the store's sync-IO capability (`_store_supports_sync_io`) — write_sync + # needs the FULL sync surface (get_sync for partial-chunk + # read-modify-write, delete_sync for all-fill chunks), not just + # set_sync — and other byte setters (e.g. the sharding codec's + # in-memory _ShardingByteSetter) on SyncByteSetter. + from zarr.abc.store import SyncByteSetter, _store_supports_sync_io from zarr.storage._common import StorePath first_bs = batch[0][0] if self.sync_transform is not None and ( - (isinstance(first_bs, StorePath) and isinstance(first_bs.store, SupportsSetSync)) + (isinstance(first_bs, StorePath) and _store_supports_sync_io(first_bs.store)) or (not isinstance(first_bs, StorePath) and isinstance(first_bs, SyncByteSetter)) ): # One thread hop for the whole batch; see the matching comment in diff --git a/src/zarr/experimental/cache_store.py b/src/zarr/experimental/cache_store.py index dd50693ad9..20cb4d4c0f 100644 --- a/src/zarr/experimental/cache_store.py +++ b/src/zarr/experimental/cache_store.py @@ -329,6 +329,16 @@ async def _get_no_cache( await self._cache_miss(key, byte_range, result) return result + @property + def _supports_sync_io(self) -> bool: + # The caching logic lives only in the async get/set/delete overrides; + # the sync methods inherited from `WrapperStore` delegate straight to + # the source store, so a sync-capable consumer (the fused codec + # pipeline) would write and delete around the cache, leaving stale + # entries that later async reads serve as current data. Opt out of + # sync IO until the sync surface is cache-aware. + return False + async def get( self, key: str, diff --git a/src/zarr/storage/_logging.py b/src/zarr/storage/_logging.py index c6f58ccd61..cdf0731430 100644 --- a/src/zarr/storage/_logging.py +++ b/src/zarr/storage/_logging.py @@ -204,6 +204,27 @@ async def delete(self, key: str) -> None: with self.log(key): return await self._store.delete(key=key) + def get_sync( + self, + key: str, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # docstring inherited + with self.log(key): + return super().get_sync(key, prototype=prototype, byte_range=byte_range) + + def set_sync(self, key: str, value: Buffer) -> None: + # docstring inherited + with self.log(key): + return super().set_sync(key, value) + + def delete_sync(self, key: str) -> None: + # docstring inherited + with self.log(key): + return super().delete_sync(key) + async def list(self) -> AsyncGenerator[str, None]: # docstring inherited with self.log(): diff --git a/src/zarr/storage/_wrapper.py b/src/zarr/storage/_wrapper.py index 37aeb8166f..6f498a655d 100644 --- a/src/zarr/storage/_wrapper.py +++ b/src/zarr/storage/_wrapper.py @@ -11,7 +11,13 @@ from zarr.abc.store import ByteRequest from zarr.core.buffer import BufferPrototype -from zarr.abc.store import Store +from zarr.abc.store import ( + Store, + SupportsDeleteSync, + SupportsGetSync, + SupportsSetSync, + _store_supports_sync_io, +) class WrapperStore[T_Store: Store](Store): @@ -149,6 +155,40 @@ def supports_writes(self) -> bool: def supports_deletes(self) -> bool: return self._store.supports_deletes + @property + def _supports_sync_io(self) -> bool: + # The delegating `*_sync` methods below make every wrapper structurally + # satisfy `SupportsSyncStore`; whether they can actually run depends on + # the wrapped store, so forward its capability (see + # `zarr.abc.store._store_supports_sync_io`). + return _store_supports_sync_io(self._store) + + def get_sync( + self, + key: str, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + """Forward `get_sync` to the wrapped store.""" + if not isinstance(self._store, SupportsGetSync): + raise TypeError(f"Store {type(self._store).__name__} does not support synchronous get.") + return self._store.get_sync(key, prototype=prototype, byte_range=byte_range) # type: ignore[unreachable] + + def set_sync(self, key: str, value: Buffer) -> None: + """Forward `set_sync` to the wrapped store.""" + if not isinstance(self._store, SupportsSetSync): + raise TypeError(f"Store {type(self._store).__name__} does not support synchronous set.") + self._store.set_sync(key, value) # type: ignore[unreachable] + + def delete_sync(self, key: str) -> None: + """Forward `delete_sync` to the wrapped store.""" + if not isinstance(self._store, SupportsDeleteSync): + raise TypeError( + f"Store {type(self._store).__name__} does not support synchronous delete." + ) + self._store.delete_sync(key) # type: ignore[unreachable] + async def delete(self, key: str) -> None: await self._store.delete(key) diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index d7011440e0..f64d8e9364 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -2,6 +2,7 @@ import asyncio import pickle +import time from abc import abstractmethod from typing import TYPE_CHECKING, Self @@ -10,6 +11,7 @@ from zarr.storage import WrapperStore if TYPE_CHECKING: + from collections.abc import AsyncIterator, Iterable, Sequence from typing import Any from zarr.core.buffer.core import BufferPrototype @@ -718,7 +720,10 @@ def set_latency(self) -> float: return max(0.0, np.random.normal(loc=self._set_latency[0], scale=self._set_latency[1])) def _with_store(self, store: Store) -> Self: - return type(self)(store, get_latency=self.get_latency, set_latency=self.set_latency) + # Pass the raw latency config, not the sampled `get_latency`/`set_latency` + # properties — sampling would freeze a `(loc, scale)` distribution into + # one fixed float on derived stores (e.g. via `with_read_only`). + return type(self)(store, get_latency=self._get_latency, set_latency=self._set_latency) async def set(self, key: str, value: Buffer) -> None: """ @@ -763,3 +768,76 @@ async def get( """ await asyncio.sleep(self.get_latency) return await self._store.get(key, prototype=prototype, byte_range=byte_range) + + def get_sync( + self, + key: str, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + """Add latency to `get_sync`. + + Sleeps `self.get_latency` on the calling thread (the sync path runs on + worker threads, not the event loop) before delegating to the wrapped + store. + """ + time.sleep(self.get_latency) + return super().get_sync(key, prototype=prototype, byte_range=byte_range) + + def set_sync(self, key: str, value: Buffer) -> None: + """Add latency to `set_sync`. + + Sleeps `self.set_latency` on the calling thread (the sync path runs on + worker threads, not the event loop) before delegating to the wrapped + store. + """ + time.sleep(self.set_latency) + super().set_sync(key, value) + + async def get_ranges( + self, + key: str, + byte_ranges: Sequence[ByteRequest | None], + *, + prototype: BufferPrototype, + max_concurrency: int | None = None, + max_gap_bytes: int | None = None, + max_coalesced_bytes: int | None = None, + ) -> AsyncIterator[Sequence[tuple[int, Buffer | None]]]: + """Byte-range reads built on `self.get`, so each fetch pays latency. + + Routes through the coalescing `Store.get_ranges` default instead of the + `WrapperStore` delegation, which would bypass this wrapper's `get` and + therefore the synthetic latency. `None` for a coalescing kwarg means + "use the `Store` default". + """ + kwargs: dict[str, int] = {} + if max_concurrency is not None: + kwargs["max_concurrency"] = max_concurrency + if max_gap_bytes is not None: + kwargs["max_gap_bytes"] = max_gap_bytes + if max_coalesced_bytes is not None: + kwargs["max_coalesced_bytes"] = max_coalesced_bytes + async for group in Store.get_ranges(self, key, byte_ranges, prototype=prototype, **kwargs): + yield group + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + """Partial-value reads built on `self.get`, so each fetch pays latency. + + Issues one `self.get` per `(key, byte_range)` pair instead of the + `WrapperStore` delegation, which would bypass this wrapper's `get` and + therefore the synthetic latency. + """ + return list( + await asyncio.gather( + *( + self.get(key, prototype=prototype, byte_range=byte_range) + for key, byte_range in key_ranges + ) + ) + ) diff --git a/tests/test_codec_pipeline_suite.py b/tests/test_codec_pipeline_suite.py index 07e1aa2ec4..f0376d185a 100644 --- a/tests/test_codec_pipeline_suite.py +++ b/tests/test_codec_pipeline_suite.py @@ -9,8 +9,8 @@ Each test also runs over a *store axis* that exercises both code paths the synchronous pipelines branch on: -* ``sync`` -> ``MemoryStore`` (supports ``get_sync``/``set_sync``: fast path) -* ``async`` -> ``LatencyStore(MemoryStore())`` (NOT sync-capable: async fallback) +* ``sync`` -> ``MemoryStore`` (full sync surface: fast path) +* ``async`` -> ``_NoSyncIOStore(MemoryStore())`` (NOT sync-capable: async fallback) The async axis is deliberate: a regression that only affects the async fallback of the default pipeline (e.g. a codec-spec-evolution bug that surfaces only on @@ -50,14 +50,22 @@ STORE_KINDS = ["sync", "async"] +class _NoSyncIOStore(LatencyStore): + """An in-memory store that advertises no sync IO capability, so a + synchronous pipeline must fall back to its async path. (A plain wrapper + won't do: `WrapperStore` forwards the wrapped store's sync capability.)""" + + @property + def _supports_sync_io(self) -> bool: + return False + + def _make_store(kind: str) -> Store: if kind == "sync": # MemoryStore supports get_sync/set_sync -> synchronous fast path. return MemoryStore() if kind == "async": - # LatencyStore is NOT SupportsGetSync/SupportsSetSync, so a synchronous - # pipeline must fall back to its async path. Zero latency keeps it fast. - return LatencyStore(MemoryStore(), get_latency=0.0, set_latency=0.0) + return _NoSyncIOStore(MemoryStore(), get_latency=0.0, set_latency=0.0) raise AssertionError(kind) diff --git a/tests/test_experimental/test_cache_store.py b/tests/test_experimental/test_cache_store.py index 5ad56a4335..f688a6ca02 100644 --- a/tests/test_experimental/test_cache_store.py +++ b/tests/test_experimental/test_cache_store.py @@ -1036,3 +1036,40 @@ async def test_delete_invalidates_cached_byte_ranges(self) -> None: # Key is gone from source result = await cached_store.get("key", proto) assert result is None + + +def test_cache_store_opts_out_of_sync_io() -> None: + """`CacheStore` must not advertise sync IO capability. + + Its caching logic lives only in the async `get`/`set`/`delete` overrides, + while the inherited `WrapperStore` sync methods delegate straight to the + source store. If the fused codec pipeline took the sync fast path, writes + and deletes would bypass the cache and later async reads would serve stale + entries. The opt-out forces sync-capable consumers onto the async path, + which keeps the cache coherent. + """ + from zarr.abc.store import _store_supports_sync_io + from zarr.storage import MemoryStore + + cached = CacheStore(MemoryStore(), cache_store=MemoryStore()) + assert _store_supports_sync_io(cached) is False + + +async def test_cache_coherent_after_fused_pipeline_write() -> None: + """Writing through the fused pipeline must not leave stale cache entries.""" + import numpy as np + + import zarr + from zarr.core.config import config as zarr_config + from zarr.storage import MemoryStore + + source = MemoryStore() + cached = CacheStore(source, cache_store=MemoryStore()) + with zarr_config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}): + arr = zarr.create_array(cached, shape=(8,), chunks=(8,), dtype="int32", fill_value=0) + arr[:] = np.arange(8, dtype="int32") + np.testing.assert_array_equal(arr[:], np.arange(8)) + # Overwrite, then read back through the same cached handle: the read + # must observe the overwrite, not a cached copy of the first write. + arr[:] = np.arange(100, 108, dtype="int32") + np.testing.assert_array_equal(arr[:], np.arange(100, 108)) diff --git a/tests/test_fused_pipeline.py b/tests/test_fused_pipeline.py index 7fa3ef2277..5c712fa97a 100644 --- a/tests/test_fused_pipeline.py +++ b/tests/test_fused_pipeline.py @@ -16,6 +16,7 @@ ArrayBytesCodecPartialEncodeMixin, BytesBytesCodec, ) +from zarr.abc.store import Store, _store_supports_sync_io from zarr.codecs.bytes import BytesCodec from zarr.codecs.gzip import GzipCodec from zarr.codecs.transpose import TransposeCodec @@ -23,9 +24,13 @@ from zarr.core.codec_pipeline import FusedCodecPipeline from zarr.core.config import config as zarr_config from zarr.registry import register_codec -from zarr.storage import MemoryStore, StorePath +from zarr.storage import MemoryStore, StorePath, WrapperStore +from zarr.storage._utils import _normalize_byte_range_index +from zarr.testing.store import LatencyStore if TYPE_CHECKING: + from collections.abc import AsyncIterator, Callable, Iterable + from zarr.abc.store import ByteRequest from zarr.core.array_spec import ArraySpec from zarr.core.buffer import Buffer, BufferPrototype, NDBuffer @@ -1065,3 +1070,166 @@ def test_partial_mixin_codec_async_partial_only_round_trip(dtype: str) -> None: with zarr_config.set(_BATCHED): np.testing.assert_array_equal(zarr.open_array(store, mode="r")[:], expected) + + +# Sync-IO capability gating (`zarr.abc.store._store_supports_sync_io`) +# +# The fused read/write fast paths must engage iff the store advertises the +# FULL synchronous IO surface (get_sync + set_sync + delete_sync): write_sync +# needs get_sync for partial-chunk read-modify-write and delete_sync for +# all-fill chunk cleanup, so gating on any single protocol can crash +# mid-batch. Wrappers must forward the capability of the wrapped store. +# --------------------------------------------------------------------------- + + +class AsyncOnlyStore(Store): + """Dict-backed store implementing only the async `Store` surface (no `*_sync`).""" + + def __init__(self) -> None: + super().__init__(read_only=False) + self._data: dict[str, Buffer] = {} + + def __eq__(self, other: object) -> bool: + return other is self + + @property + def supports_writes(self) -> bool: + return True + + @property + def supports_deletes(self) -> bool: + return True + + @property + def supports_listing(self) -> bool: + return True + + async def get( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + try: + value = self._data[key] + except KeyError: + return None + start, stop = _normalize_byte_range_index(value, byte_range) + return prototype.buffer.from_buffer(value[start:stop]) + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + return [await self.get(key, prototype, byte_range) for key, byte_range in key_ranges] + + async def exists(self, key: str) -> bool: + return key in self._data + + async def set(self, key: str, value: Buffer) -> None: + self._check_writable() + self._data[key] = value + + async def delete(self, key: str) -> None: + self._check_writable() + self._data.pop(key, None) + + async def list(self) -> AsyncIterator[str]: + for key in list(self._data): + yield key + + async def list_prefix(self, prefix: str) -> AsyncIterator[str]: + for key in list(self._data): + if key.startswith(prefix): + yield key + + async def list_dir(self, prefix: str) -> AsyncIterator[str]: + if prefix and not prefix.endswith("/"): + prefix += "/" + seen: set[str] = set() + for key in list(self._data): + if key.startswith(prefix): + head = key.removeprefix(prefix).split("/")[0] + if head not in seen: + seen.add(head) + yield head + + +class SetOnlySyncStore(AsyncOnlyStore): + """Implements `set_sync` but not `get_sync`/`delete_sync` (partial sync surface).""" + + def set_sync(self, key: str, value: Buffer) -> None: + self._check_writable() + self._data[key] = value + + +@pytest.mark.parametrize( + ("store_factory", "expect_sync_path"), + [ + (MemoryStore, True), + (lambda: WrapperStore(MemoryStore()), True), + (lambda: LatencyStore(MemoryStore()), True), + (SetOnlySyncStore, False), + (lambda: WrapperStore(AsyncOnlyStore()), False), + ], + ids=[ + "full-sync", + "wrapper-of-sync", + "latency-wrapper-of-sync", + "set-sync-only", + "wrapper-of-async-only", + ], +) +def test_sync_io_capability_gates_fused_paths( + store_factory: Callable[[], Store], expect_sync_path: bool +) -> None: + """The fused pipeline takes the sync fast path iff the store satisfies + `_store_supports_sync_io`, + and every store round-trips correctly through full writes, partial + (read-modify-write) writes, and all-fill (delete) writes — a store with a + partial sync surface must get a clean async fallback, never a mid-batch + error.""" + from unittest.mock import patch + + store = store_factory() + assert _store_supports_sync_io(store) is expect_sync_path + + calls = {"read_sync": 0, "write_sync": 0} + orig_read_sync = FusedCodecPipeline.read_sync + orig_write_sync = FusedCodecPipeline.write_sync + + def spy_read_sync(self: FusedCodecPipeline, *args: Any, **kwargs: Any) -> Any: + calls["read_sync"] += 1 + return orig_read_sync(self, *args, **kwargs) + + def spy_write_sync(self: FusedCodecPipeline, *args: Any, **kwargs: Any) -> Any: + calls["write_sync"] += 1 + return orig_write_sync(self, *args, **kwargs) + + with zarr_config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}): + arr = zarr.create_array( + store=store, + shape=(8,), + chunks=(4,), + dtype="uint8", + compressors=None, + fill_value=0, + ) + assert isinstance(arr._async_array.codec_pipeline, FusedCodecPipeline) + with ( + patch.object(FusedCodecPipeline, "read_sync", spy_read_sync), + patch.object(FusedCodecPipeline, "write_sync", spy_write_sync), + ): + data = np.arange(8, dtype="uint8") + arr[:] = data # complete-chunk writes + arr[:3] = 7 # partial write -> read-modify-write needs get + data[:3] = 7 + np.testing.assert_array_equal(arr[:], data) + arr[4:8] = 0 # all-fill chunk -> delete needed + data[4:8] = 0 + np.testing.assert_array_equal(arr[:], data) + + if expect_sync_path: + assert calls["write_sync"] > 0, "sync-capable store did not take the sync write path" + assert calls["read_sync"] > 0, "sync-capable store did not take the sync read path" + else: + assert calls["write_sync"] == 0, "non-sync store took the sync write path" + assert calls["read_sync"] == 0, "non-sync store took the sync read path" diff --git a/tests/test_store/test_get_ranges.py b/tests/test_store/test_get_ranges.py index f04251adf4..522d6565aa 100644 --- a/tests/test_store/test_get_ranges.py +++ b/tests/test_store/test_get_ranges.py @@ -16,12 +16,12 @@ from zarr.abc.store import RangeByteRequest from zarr.core.buffer import default_buffer_prototype -from zarr.storage import MemoryStore +from zarr.storage import MemoryStore, ZipStore from zarr.storage._wrapper import WrapperStore -from zarr.testing.store import LatencyStore if TYPE_CHECKING: from collections.abc import AsyncIterator, Sequence + from pathlib import Path from zarr.abc.store import ByteRequest from zarr.core.buffer import Buffer, BufferPrototype @@ -89,11 +89,11 @@ def test_get_ranges_sync_missing_key_raises() -> None: store.get_ranges_sync("does-not-exist", [RangeByteRequest(0, 10)], prototype=proto) -def test_get_ranges_sync_on_non_sync_store_raises_type_error() -> None: +def test_get_ranges_sync_on_non_sync_store_raises_type_error(tmp_path: Path) -> None: """`get_ranges_sync` requires the store to support synchronous reads (`SupportsGetSync`); a non-sync store raises TypeError rather than silently falling back.""" - store = LatencyStore(MemoryStore(), get_latency=0.0, set_latency=0.0) + store = ZipStore(tmp_path / "store.zip", mode="w") proto = default_buffer_prototype() with pytest.raises(TypeError, match="does not support synchronous reads"): store.get_ranges_sync("k", [RangeByteRequest(0, 10)], prototype=proto) diff --git a/tests/test_store/test_latency.py b/tests/test_store/test_latency.py index 38ffb17dd6..9cb71fd6a0 100644 --- a/tests/test_store/test_latency.py +++ b/tests/test_store/test_latency.py @@ -1,8 +1,16 @@ from __future__ import annotations +import time +from unittest.mock import patch + +import numpy as np import pytest +import zarr +from zarr.abc.store import RangeByteRequest from zarr.core.buffer import default_buffer_prototype +from zarr.core.codec_pipeline import FusedCodecPipeline +from zarr.core.config import config as zarr_config from zarr.storage import MemoryStore from zarr.testing.store import LatencyStore @@ -55,3 +63,103 @@ async def test_latency_store_with_read_only_round_trip() -> None: # The original read-only wrapper remains read-only assert latency_ro.read_only + + +@pytest.mark.parametrize( + ("get_latency", "set_latency"), + [ + (0.01, 0.02), + ((0.1, 0.05), (0.2, 0.01)), + ], + ids=["scalar", "distribution"], +) +def test_with_store_preserves_latency_config( + get_latency: float | tuple[float, float], set_latency: float | tuple[float, float] +) -> None: + """Derived stores (e.g. via `with_read_only`) keep the raw latency config — + a `(loc, scale)` distribution must not collapse to one sampled float.""" + store = LatencyStore(MemoryStore(), get_latency=get_latency, set_latency=set_latency) + derived = store.with_read_only(True) + assert derived._get_latency == store._get_latency + assert derived._set_latency == store._set_latency + + +def test_sync_methods_inject_latency(monkeypatch: pytest.MonkeyPatch) -> None: + """`get_sync`/`set_sync` sleep the configured latency on the calling thread + before delegating to the wrapped store.""" + sleeps: list[float] = [] + monkeypatch.setattr(time, "sleep", sleeps.append) + + store = LatencyStore(MemoryStore(), get_latency=0.123, set_latency=0.456) + buf = default_buffer_prototype().buffer.from_bytes(b"abcd") + store.set_sync("key", buf) + assert sleeps == [pytest.approx(0.456)] + out = store.get_sync("key", prototype=default_buffer_prototype()) + assert out is not None + assert out.to_bytes() == b"abcd" + assert sleeps == [pytest.approx(0.456), pytest.approx(0.123)] + + +async def test_get_ranges_pays_latency_per_fetch() -> None: + """`get_ranges` routes through the coalescing default built on `self.get`, + so each merged fetch pays the configured latency instead of bypassing it + via WrapperStore delegation. Two ranges further apart than `max_gap_bytes` + cannot coalesce -> exactly two `get` calls.""" + proto = default_buffer_prototype() + inner = MemoryStore() + await inner.set("blob", proto.buffer.from_bytes(bytes(4 << 20))) + store = LatencyStore(inner, get_latency=0.0) + + requests = [RangeByteRequest(0, 10), RangeByteRequest(2 << 20, (2 << 20) + 10)] + results: list[tuple[int, object]] = [] + with patch.object(store, "get", wraps=store.get) as get_spy: + async for group in store.get_ranges("blob", requests, prototype=proto): + results.extend(group) + assert get_spy.await_count == 2 + assert sorted(idx for idx, _ in results) == [0, 1] + for _, buf in results: + assert buf is not None + assert len(buf) == 10 # type: ignore[arg-type] + + +async def test_get_partial_values_routes_through_get() -> None: + """`get_partial_values` issues one `self.get` per key-range so each fetch + pays the configured latency instead of bypassing it via WrapperStore + delegation.""" + proto = default_buffer_prototype() + inner = MemoryStore() + await inner.set("blob", proto.buffer.from_bytes(b"0123456789")) + store = LatencyStore(inner, get_latency=0.0) + + with patch.object(store, "get", wraps=store.get) as get_spy: + results = await store.get_partial_values( + proto, [("blob", RangeByteRequest(0, 4)), ("blob", None)] + ) + assert get_spy.await_count == 2 + assert results[0] is not None + assert results[0].to_bytes() == b"0123" + assert results[1] is not None + assert results[1].to_bytes() == b"0123456789" + + +def test_latency_store_engages_fused_sync_path() -> None: + """A LatencyStore wrapping a sync-capable store must take the fused sync + fast path: reads go through the inner store's `get_sync`, not the async + fallback.""" + inner = MemoryStore() + store = LatencyStore(inner, get_latency=0.0, set_latency=0.0) + with zarr_config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}): + arr = zarr.create_array( + store=store, + shape=(8,), + chunks=(4,), + dtype="uint8", + compressors=None, + fill_value=0, + ) + assert isinstance(arr._async_array.codec_pipeline, FusedCodecPipeline) + data = np.arange(8, dtype="uint8") + arr[:] = data + with patch.object(inner, "get_sync", wraps=inner.get_sync) as get_sync_spy: + np.testing.assert_array_equal(arr[:], data) + assert get_sync_spy.call_count == 2 # one per chunk diff --git a/tests/test_store/test_wrapper.py b/tests/test_store/test_wrapper.py index b34a63d5d0..e556a108c5 100644 --- a/tests/test_store/test_wrapper.py +++ b/tests/test_store/test_wrapper.py @@ -4,12 +4,12 @@ import pytest -from zarr.abc.store import ByteRequest, Store +from zarr.abc.store import ByteRequest, Store, _store_supports_sync_io from zarr.core.buffer import Buffer from zarr.core.buffer.cpu import Buffer as CPUBuffer from zarr.core.buffer.cpu import buffer_prototype -from zarr.storage import LocalStore, WrapperStore -from zarr.testing.store import StoreTests +from zarr.storage import LocalStore, MemoryStore, WrapperStore, ZipStore +from zarr.testing.store import LatencyStore, StoreTests if TYPE_CHECKING: from pathlib import Path @@ -123,3 +123,47 @@ async def get( await store_wrapped.get(key, buffer_prototype) captured = capsys.readouterr() assert f"getting {key}" in captured.out + + +@pytest.mark.parametrize( + ("store_factory", "expected"), + [ + (lambda tmp: MemoryStore(), True), + (lambda tmp: LocalStore(str(tmp)), True), + (lambda tmp: WrapperStore(MemoryStore()), True), + (lambda tmp: LatencyStore(MemoryStore()), True), + (lambda tmp: ZipStore(tmp / "store.zip", mode="w"), False), + (lambda tmp: WrapperStore(ZipStore(tmp / "store.zip", mode="w")), False), + ], + ids=[ + "memory", + "local", + "wrapper-of-memory", + "latency-wrapper-of-memory", + "zip", + "wrapper-of-zip", + ], +) +def test_supports_sync_io(store_factory: Any, expected: bool, tmp_path: Path | Any) -> None: + """`_store_supports_sync_io` is True only for stores implementing the full + sync surface (get_sync + set_sync + delete_sync); wrappers forward the + wrapped store's capability via `_supports_sync_io`.""" + assert _store_supports_sync_io(store_factory(tmp_path)) is expected + + +def test_wrapper_get_sync_without_inner_sync_raises(tmp_path: Any) -> None: + store = WrapperStore(ZipStore(tmp_path / "store.zip", mode="w")) + with pytest.raises(TypeError, match="does not support synchronous get"): + store.get_sync("key") + + +def test_wrapper_set_sync_without_inner_sync_raises(tmp_path: Any) -> None: + store = WrapperStore(ZipStore(tmp_path / "store.zip", mode="w")) + with pytest.raises(TypeError, match="does not support synchronous set"): + store.set_sync("key", CPUBuffer.from_bytes(b"data")) + + +def test_wrapper_delete_sync_without_inner_sync_raises(tmp_path: Any) -> None: + store = WrapperStore(ZipStore(tmp_path / "store.zip", mode="w")) + with pytest.raises(TypeError, match="does not support synchronous delete"): + store.delete_sync("key")