Skip to content

Commit c3dfafc

Browse files
committed
fix: gate fused sync fast paths on full store sync capability; wrappers forward it
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
1 parent 123268a commit c3dfafc

12 files changed

Lines changed: 556 additions & 34 deletions

File tree

changes/4206.bugfix.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
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.

src/zarr/abc/store.py

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -662,6 +662,15 @@ def delete_sync(self) -> None: ...
662662

663663
@runtime_checkable
664664
class SupportsGetSync(Protocol):
665+
"""Store protocol for synchronous reads (`get_sync`).
666+
667+
The store sync surface is all-or-nothing: a store implementing any of the
668+
`*_sync` methods must implement all of them (`SupportsSyncStore`), because
669+
consumers mix sync reads, writes, and deletes within one operation.
670+
Capability-gated callers consult `_store_supports_sync_io` rather than the
671+
individual protocols.
672+
"""
673+
665674
def get_sync(
666675
self,
667676
key: str,
@@ -673,16 +682,52 @@ def get_sync(
673682

674683
@runtime_checkable
675684
class SupportsSetSync(Protocol):
685+
"""Store protocol for synchronous writes (`set_sync`).
686+
687+
See `SupportsGetSync` for the all-or-nothing contract on the store sync
688+
surface.
689+
"""
690+
676691
def set_sync(self, key: str, value: Buffer) -> None: ...
677692

678693

679694
@runtime_checkable
680695
class SupportsDeleteSync(Protocol):
696+
"""Store protocol for synchronous deletes (`delete_sync`).
697+
698+
See `SupportsGetSync` for the all-or-nothing contract on the store sync
699+
surface.
700+
"""
701+
681702
def delete_sync(self, key: str) -> None: ...
682703

683704

684705
@runtime_checkable
685-
class SupportsSyncStore(SupportsGetSync, SupportsSetSync, SupportsDeleteSync, Protocol): ...
706+
class SupportsSyncStore(SupportsGetSync, SupportsSetSync, SupportsDeleteSync, Protocol):
707+
"""The full store sync surface: `get_sync`, `set_sync`, and `delete_sync`."""
708+
709+
710+
def _store_supports_sync_io(store: object) -> bool:
711+
"""Whether `store` can serve the full synchronous IO surface right now.
712+
713+
Structural membership in `SupportsSyncStore` is necessary but not always
714+
sufficient: a store can present the `*_sync` methods while its ability to
715+
run them depends on runtime state the type system cannot see. Wrapper
716+
stores are the canonical case — `WrapperStore` delegates the sync methods
717+
to the store it wraps, so they only work when the wrapped store is itself
718+
sync-capable. Such stores opt out dynamically via a `_supports_sync_io`
719+
attribute/property (absent means capable).
720+
721+
This is an interim, private convention pending a formal sync/async store
722+
architecture — the store-side twin of the codec-side `_sync_capable`
723+
convention consulted by `zarr.abc.codec._codec_supports_sync`.
724+
725+
Synchronous IO is all-or-nothing: consumers such as the fused codec
726+
pipeline mix synchronous reads, writes, and deletes within one batch
727+
(e.g. a partial-chunk write reads existing bytes and an all-fill chunk is
728+
deleted), so a partial sync surface never satisfies this predicate.
729+
"""
730+
return isinstance(store, SupportsSyncStore) and getattr(store, "_supports_sync_io", True)
686731

687732

688733
async def set_or_delete(byte_setter: ByteSetter, value: Buffer | None) -> None:

src/zarr/codecs/sharding.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
RangeByteRequest,
2424
Store,
2525
SuffixByteRequest,
26-
SupportsGetSync,
26+
_store_supports_sync_io,
2727
)
2828
from zarr.codecs._deprecated_enum import _coerce_enum_input, _DeprecatedStrEnumMeta
2929
from zarr.codecs.bytes import BytesCodec
@@ -1679,7 +1679,7 @@ def _load_partial_shard_maybe_sync(
16791679

16801680
shard_dict: ShardMutableMapping = {}
16811681
store = byte_getter.store if hasattr(byte_getter, "store") else None
1682-
if isinstance(store, Store) and isinstance(store, SupportsGetSync):
1682+
if isinstance(store, Store) and _store_supports_sync_io(store):
16831683
# External store: coalesce via get_ranges_sync (mirrors get_ranges).
16841684
byte_ranges = [byte_range for _, byte_range in chunk_coord_byte_ranges]
16851685
try:

src/zarr/core/codec_pipeline.py

Lines changed: 18 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -283,8 +283,8 @@ async def _async_read_fallback(
283283
then scatters each decoded chunk into `out` at its `out_selection`.
284284
285285
Used by both `BatchedCodecPipeline.read_batch` (non-partial-decode
286-
branch) and `FusedCodecPipeline.read` (when the store is not a
287-
`SupportsGetSync` / sync transform is unavailable).
286+
branch) and `FusedCodecPipeline.read` (when the store does not advertise
287+
sync IO / sync transform is unavailable).
288288
"""
289289

290290
chunk_array_batch: list[NDBuffer | None]
@@ -345,8 +345,8 @@ async def _async_write_fallback(
345345
if encoding produced `None` or the chunk dropped).
346346
347347
Used by both `BatchedCodecPipeline.write_batch` (non-partial-encode
348-
branch) and `FusedCodecPipeline.write` (when the store is not a
349-
`SupportsSetSync` / sync transform is unavailable).
348+
branch) and `FusedCodecPipeline.write` (when the store does not advertise
349+
sync IO / sync transform is unavailable).
350350
"""
351351

352352
if use_sync := (
@@ -1170,16 +1170,17 @@ async def read(
11701170
return ()
11711171

11721172
# Fast path: sync transform plus synchronous IO. For StorePath the gate
1173-
# is on the STORE's sync support (StorePath always has a get_sync
1174-
# method, but it only works when its store does); for other byte
1175-
# getters (e.g. the sharding codec's in-memory _ShardingByteGetter) the
1176-
# SyncByteGetter protocol is the gate.
1177-
from zarr.abc.store import SupportsGetSync, SyncByteGetter
1173+
# is the STORE's sync-IO capability (`_store_supports_sync_io`) (StorePath always has a
1174+
# get_sync method, but it only works when its store implements the full
1175+
# sync surface); for other byte getters (e.g. the sharding codec's
1176+
# in-memory _ShardingByteGetter) the SyncByteGetter protocol is the
1177+
# gate.
1178+
from zarr.abc.store import SyncByteGetter, _store_supports_sync_io
11781179
from zarr.storage._common import StorePath
11791180

11801181
first_bg = batch[0][0]
11811182
if self.sync_transform is not None and (
1182-
(isinstance(first_bg, StorePath) and isinstance(first_bg.store, SupportsGetSync))
1183+
(isinstance(first_bg, StorePath) and _store_supports_sync_io(first_bg.store))
11831184
or (not isinstance(first_bg, StorePath) and isinstance(first_bg, SyncByteGetter))
11841185
):
11851186
# One thread hop for the WHOLE batch — not per chunk, so the fused
@@ -1233,14 +1234,17 @@ async def write(
12331234
return
12341235

12351236
# Fast path: sync transform plus synchronous IO. Mirrors `read`: gate
1236-
# StorePath on the store's sync support, other byte setters (e.g. the
1237-
# sharding codec's in-memory _ShardingByteSetter) on SyncByteSetter.
1238-
from zarr.abc.store import SupportsSetSync, SyncByteSetter
1237+
# StorePath on the store's sync-IO capability (`_store_supports_sync_io`) — write_sync
1238+
# needs the FULL sync surface (get_sync for partial-chunk
1239+
# read-modify-write, delete_sync for all-fill chunks), not just
1240+
# set_sync — and other byte setters (e.g. the sharding codec's
1241+
# in-memory _ShardingByteSetter) on SyncByteSetter.
1242+
from zarr.abc.store import SyncByteSetter, _store_supports_sync_io
12391243
from zarr.storage._common import StorePath
12401244

12411245
first_bs = batch[0][0]
12421246
if self.sync_transform is not None and (
1243-
(isinstance(first_bs, StorePath) and isinstance(first_bs.store, SupportsSetSync))
1247+
(isinstance(first_bs, StorePath) and _store_supports_sync_io(first_bs.store))
12441248
or (not isinstance(first_bs, StorePath) and isinstance(first_bs, SyncByteSetter))
12451249
):
12461250
# One thread hop for the whole batch; see the matching comment in

src/zarr/storage/_logging.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,27 @@ async def delete(self, key: str) -> None:
204204
with self.log(key):
205205
return await self._store.delete(key=key)
206206

207+
def get_sync(
208+
self,
209+
key: str,
210+
*,
211+
prototype: BufferPrototype | None = None,
212+
byte_range: ByteRequest | None = None,
213+
) -> Buffer | None:
214+
# docstring inherited
215+
with self.log(key):
216+
return super().get_sync(key, prototype=prototype, byte_range=byte_range)
217+
218+
def set_sync(self, key: str, value: Buffer) -> None:
219+
# docstring inherited
220+
with self.log(key):
221+
return super().set_sync(key, value)
222+
223+
def delete_sync(self, key: str) -> None:
224+
# docstring inherited
225+
with self.log(key):
226+
return super().delete_sync(key)
227+
207228
async def list(self) -> AsyncGenerator[str, None]:
208229
# docstring inherited
209230
with self.log():

src/zarr/storage/_wrapper.py

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,13 @@
1111
from zarr.abc.store import ByteRequest
1212
from zarr.core.buffer import BufferPrototype
1313

14-
from zarr.abc.store import Store
14+
from zarr.abc.store import (
15+
Store,
16+
SupportsDeleteSync,
17+
SupportsGetSync,
18+
SupportsSetSync,
19+
_store_supports_sync_io,
20+
)
1521

1622

1723
class WrapperStore[T_Store: Store](Store):
@@ -149,6 +155,40 @@ def supports_writes(self) -> bool:
149155
def supports_deletes(self) -> bool:
150156
return self._store.supports_deletes
151157

158+
@property
159+
def _supports_sync_io(self) -> bool:
160+
# The delegating `*_sync` methods below make every wrapper structurally
161+
# satisfy `SupportsSyncStore`; whether they can actually run depends on
162+
# the wrapped store, so forward its capability (see
163+
# `zarr.abc.store._store_supports_sync_io`).
164+
return _store_supports_sync_io(self._store)
165+
166+
def get_sync(
167+
self,
168+
key: str,
169+
*,
170+
prototype: BufferPrototype | None = None,
171+
byte_range: ByteRequest | None = None,
172+
) -> Buffer | None:
173+
"""Forward `get_sync` to the wrapped store."""
174+
if not isinstance(self._store, SupportsGetSync):
175+
raise TypeError(f"Store {type(self._store).__name__} does not support synchronous get.")
176+
return self._store.get_sync(key, prototype=prototype, byte_range=byte_range) # type: ignore[unreachable]
177+
178+
def set_sync(self, key: str, value: Buffer) -> None:
179+
"""Forward `set_sync` to the wrapped store."""
180+
if not isinstance(self._store, SupportsSetSync):
181+
raise TypeError(f"Store {type(self._store).__name__} does not support synchronous set.")
182+
self._store.set_sync(key, value) # type: ignore[unreachable]
183+
184+
def delete_sync(self, key: str) -> None:
185+
"""Forward `delete_sync` to the wrapped store."""
186+
if not isinstance(self._store, SupportsDeleteSync):
187+
raise TypeError(
188+
f"Store {type(self._store).__name__} does not support synchronous delete."
189+
)
190+
self._store.delete_sync(key) # type: ignore[unreachable]
191+
152192
async def delete(self, key: str) -> None:
153193
await self._store.delete(key)
154194

src/zarr/testing/store.py

Lines changed: 79 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import asyncio
44
import pickle
5+
import time
56
from abc import abstractmethod
67
from typing import TYPE_CHECKING, Self
78

@@ -10,6 +11,7 @@
1011
from zarr.storage import WrapperStore
1112

1213
if TYPE_CHECKING:
14+
from collections.abc import AsyncIterator, Iterable, Sequence
1315
from typing import Any
1416

1517
from zarr.core.buffer.core import BufferPrototype
@@ -653,7 +655,10 @@ def set_latency(self) -> float:
653655
return max(0.0, np.random.normal(loc=self._set_latency[0], scale=self._set_latency[1]))
654656

655657
def _with_store(self, store: Store) -> Self:
656-
return type(self)(store, get_latency=self.get_latency, set_latency=self.set_latency)
658+
# Pass the raw latency config, not the sampled `get_latency`/`set_latency`
659+
# properties — sampling would freeze a `(loc, scale)` distribution into
660+
# one fixed float on derived stores (e.g. via `with_read_only`).
661+
return type(self)(store, get_latency=self._get_latency, set_latency=self._set_latency)
657662

658663
async def set(self, key: str, value: Buffer) -> None:
659664
"""
@@ -698,3 +703,76 @@ async def get(
698703
"""
699704
await asyncio.sleep(self.get_latency)
700705
return await self._store.get(key, prototype=prototype, byte_range=byte_range)
706+
707+
def get_sync(
708+
self,
709+
key: str,
710+
*,
711+
prototype: BufferPrototype | None = None,
712+
byte_range: ByteRequest | None = None,
713+
) -> Buffer | None:
714+
"""Add latency to `get_sync`.
715+
716+
Sleeps `self.get_latency` on the calling thread (the sync path runs on
717+
worker threads, not the event loop) before delegating to the wrapped
718+
store.
719+
"""
720+
time.sleep(self.get_latency)
721+
return super().get_sync(key, prototype=prototype, byte_range=byte_range)
722+
723+
def set_sync(self, key: str, value: Buffer) -> None:
724+
"""Add latency to `set_sync`.
725+
726+
Sleeps `self.set_latency` on the calling thread (the sync path runs on
727+
worker threads, not the event loop) before delegating to the wrapped
728+
store.
729+
"""
730+
time.sleep(self.set_latency)
731+
super().set_sync(key, value)
732+
733+
async def get_ranges(
734+
self,
735+
key: str,
736+
byte_ranges: Sequence[ByteRequest | None],
737+
*,
738+
prototype: BufferPrototype,
739+
max_concurrency: int | None = None,
740+
max_gap_bytes: int | None = None,
741+
max_coalesced_bytes: int | None = None,
742+
) -> AsyncIterator[Sequence[tuple[int, Buffer | None]]]:
743+
"""Byte-range reads built on `self.get`, so each fetch pays latency.
744+
745+
Routes through the coalescing `Store.get_ranges` default instead of the
746+
`WrapperStore` delegation, which would bypass this wrapper's `get` and
747+
therefore the synthetic latency. `None` for a coalescing kwarg means
748+
"use the `Store` default".
749+
"""
750+
kwargs: dict[str, int] = {}
751+
if max_concurrency is not None:
752+
kwargs["max_concurrency"] = max_concurrency
753+
if max_gap_bytes is not None:
754+
kwargs["max_gap_bytes"] = max_gap_bytes
755+
if max_coalesced_bytes is not None:
756+
kwargs["max_coalesced_bytes"] = max_coalesced_bytes
757+
async for group in Store.get_ranges(self, key, byte_ranges, prototype=prototype, **kwargs):
758+
yield group
759+
760+
async def get_partial_values(
761+
self,
762+
prototype: BufferPrototype,
763+
key_ranges: Iterable[tuple[str, ByteRequest | None]],
764+
) -> list[Buffer | None]:
765+
"""Partial-value reads built on `self.get`, so each fetch pays latency.
766+
767+
Issues one `self.get` per `(key, byte_range)` pair instead of the
768+
`WrapperStore` delegation, which would bypass this wrapper's `get` and
769+
therefore the synthetic latency.
770+
"""
771+
return list(
772+
await asyncio.gather(
773+
*(
774+
self.get(key, prototype=prototype, byte_range=byte_range)
775+
for key, byte_range in key_ranges
776+
)
777+
)
778+
)

tests/test_codec_pipeline_suite.py

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@
99
Each test also runs over a *store axis* that exercises both code paths the
1010
synchronous pipelines branch on:
1111
12-
* ``sync`` -> ``MemoryStore`` (supports ``get_sync``/``set_sync``: fast path)
13-
* ``async`` -> ``LatencyStore(MemoryStore())`` (NOT sync-capable: async fallback)
12+
* ``sync`` -> ``MemoryStore`` (full sync surface: fast path)
13+
* ``async`` -> ``_NoSyncIOStore(MemoryStore())`` (NOT sync-capable: async fallback)
1414
1515
The async axis is deliberate: a regression that only affects the async fallback
1616
of the default pipeline (e.g. a codec-spec-evolution bug that surfaces only on
@@ -50,14 +50,22 @@
5050
STORE_KINDS = ["sync", "async"]
5151

5252

53+
class _NoSyncIOStore(LatencyStore):
54+
"""An in-memory store that advertises no sync IO capability, so a
55+
synchronous pipeline must fall back to its async path. (A plain wrapper
56+
won't do: `WrapperStore` forwards the wrapped store's sync capability.)"""
57+
58+
@property
59+
def _supports_sync_io(self) -> bool:
60+
return False
61+
62+
5363
def _make_store(kind: str) -> Store:
5464
if kind == "sync":
5565
# MemoryStore supports get_sync/set_sync -> synchronous fast path.
5666
return MemoryStore()
5767
if kind == "async":
58-
# LatencyStore is NOT SupportsGetSync/SupportsSetSync, so a synchronous
59-
# pipeline must fall back to its async path. Zero latency keeps it fast.
60-
return LatencyStore(MemoryStore(), get_latency=0.0, set_latency=0.0)
68+
return _NoSyncIOStore(MemoryStore(), get_latency=0.0, set_latency=0.0)
6169
raise AssertionError(kind)
6270

6371

0 commit comments

Comments
 (0)