From a49296e1d721fef265e0f8173ca244a77f178081 Mon Sep 17 00:00:00 2001 From: TheGreatAlgo <37487508+TheGreatAlgo@users.noreply.github.com> Date: Tue, 28 Jul 2026 09:44:10 -0400 Subject: [PATCH 1/2] fix: auto mode --- py_hamt/sharded_zarr_store.py | 140 ++++++++++++++-- tests/test_benchmark_stores.py | 118 ++++++++++++-- tests/test_sharded_zarr_store_v2.py | 245 +++++++++++++++++++++++++++- 3 files changed, 472 insertions(+), 31 deletions(-) diff --git a/py_hamt/sharded_zarr_store.py b/py_hamt/sharded_zarr_store.py index 7d0d072..b8ded42 100644 --- a/py_hamt/sharded_zarr_store.py +++ b/py_hamt/sharded_zarr_store.py @@ -39,7 +39,7 @@ _V1_INFERENCE_CONCURRENCY = 8 ShardCacheKey = int | tuple[str, int] -ShardReadMode = Literal["full", "sparse"] +ShardReadMode = Literal["full", "sparse", "auto"] def _read_cbor_argument( @@ -473,6 +473,22 @@ class ShardedZarrStore(zarr.abc.store.Store): "sharded_zarr_v2 writes require an explicit Zarr group. Write the " "dataset with ds.to_zarr(..., group='0') or another group name." ) + # Sparse reads of a *single* shard before "auto" latches the whole store to + # full decodes. Measured crossover, one CAS round trip per sparse entry: + # + # lookups full sparse + # 1 775ms 3.6ms + # 32 775ms 82.2ms + # 128 830ms 340.0ms + # 256 861ms 706.3ms <- sparse still ahead + # 512 832ms 1334.9ms <- full ahead + # + # Sparse costs ~2.6ms/lookup against a ~800ms flat full decode, so the + # break-even is ~300; 256 trips just before it. Unlike the jaxray reference + # (threshold 32, local blockstore reads), this is a *scan-detection* + # threshold rather than a per-shard promotion point: it is paid once for the + # whole store, not once per shard. + _SPARSE_PROMOTE_THRESHOLD: ClassVar[int] = 256 def __init__( self, @@ -481,14 +497,14 @@ def __init__( root_cid: Optional[str] = None, *, max_cache_memory_bytes: int = 100 * 1024 * 1024, # 100MB default - shard_read_mode: ShardReadMode = "sparse", + shard_read_mode: ShardReadMode = "auto", ): """Use the async `open()` classmethod to instantiate this class.""" super().__init__(read_only=read_only) - if shard_read_mode not in {"full", "sparse"}: + if shard_read_mode not in {"full", "sparse", "auto"}: raise ValueError( f"Unsupported shard_read_mode: {shard_read_mode!r}. " - "Expected 'full' or 'sparse'." + "Expected 'full', 'sparse', or 'auto'." ) self.cas = cas self._root_cid = root_cid @@ -506,6 +522,14 @@ def __init__( self._shard_data_cache = MemoryBoundedLRUCache(max_cache_memory_bytes) self._pending_shard_loads: Dict[ShardCacheKey, asyncio.Event] = {} + # Per-shard sparse-read counts, used only to detect the scan pattern in + # "auto" mode. Detection is per-shard because 256 reads spread across + # 256 distinct shards is a point-read workload, not a scan. + self._sparse_read_counts: Dict[ShardCacheKey, int] = {} + # Latched once any single shard crosses the threshold: the caller is + # scanning, so every shard gets the full path from here on. Store-wide + # because the access pattern belongs to the caller, not the shard. + self._full_mode_latched: bool = False self._metadata_read_cache: Dict[str, bytes] = {} self.array_indices: Dict[str, ArrayIndex] = {} @@ -610,7 +634,7 @@ async def open( max_cache_memory_bytes: int = 100 * 1024 * 1024, # 100MB default manifest_version: Optional[str] = None, primary_array_path: str = "", - shard_read_mode: ShardReadMode = "sparse", + shard_read_mode: ShardReadMode = "auto", ) -> "ShardedZarrStore": """ Asynchronously opens an existing ShardedZarrStore or initializes a new one. @@ -618,6 +642,31 @@ async def open( Shape-based creation remains the v1 compatibility path. To create a new path-aware v2 store, pass ``manifest_version="sharded_zarr_v2"`` or omit ``array_shape``/``chunk_shape`` and provide ``chunks_per_shard``. + + ``shard_read_mode`` controls how a **read-only** cache miss resolves a + chunk pointer. It has no effect on writes: a writable store always goes + through the shard cache so pending writes stay visible, so writes behave + as ``"full"`` does regardless of this setting. + + - ``"auto"`` (the default) starts sparse, then latches the **entire + store** to full decodes once any *single* shard has been read + ``_SPARSE_PROMOTE_THRESHOLD`` times. The latch is store-wide because + the access pattern belongs to the caller rather than the shard: a + caller reading one shard that heavily is scanning and will scan the + rest too, so making every other shard re-learn that independently + would re-pay the detection cost on each one. It is permanent for the + store's lifetime and unaffected by cache eviction. Below the + threshold it is byte-for-byte the ``"sparse"`` path, so point reads + pay nothing for the safety net. + - ``"sparse"`` fetches only the requested entry and caches nothing. Far + cheaper for point reads, but degrades without bound on a scan, + eventually costing more than ``"full"``. Pin this when you know the + workload is point reads and want to rule out the latch entirely -- + for instance a long-lived reader that hammers one hot shard without + ever scanning, which ``"auto"`` would latch on. + - ``"full"`` decodes and caches the whole shard. Flat cost regardless of + how many chunks are then read from it, so it suits known scans and + skips ``"auto"``'s detection cost. """ store = cls( cas, @@ -1493,6 +1542,55 @@ async def _fetch_and_cache_full_shard( f"Failed to fetch shard {shard_idx} after {max_retries} attempts: {e}" ) from e + def _sparse_read_is_eligible( + self, + array_index: ArrayIndex, + shard_idx: int, + cached_shard: Optional[List[Optional[CID]]], + byte_range: Optional[zarr.abc.store.ByteRequest], + ) -> bool: + """ + Whether a single-entry shard decode is structurally legal here. + + Independent of ``shard_read_mode``: a writable store must go through the + cache so pending writes stay visible, a cache hit is already cheaper + than any fetch, a byte range needs the full CID resolution path, and an + absent or out-of-range shard CID has nothing to sparsely decode. + """ + return ( + self.read_only + and cached_shard is None + and byte_range is None + and 0 <= shard_idx < array_index.num_shards + and array_index.shard_cids[shard_idx] is not None + ) + + def _auto_mode_wants_sparse(self, cache_key: ShardCacheKey) -> bool: + """ + Record one ``auto``-mode sparse read and report whether to stay sparse. + + Called only under ``self._shard_locks[cache_key]``, so each shard's + read-modify-write of the counter is serialized. Once any single shard + crosses the threshold the whole store latches to full mode: a caller + reading one shard that heavily is scanning, and will scan the rest too. + + ``_full_mode_latched`` is written under one shard's lock and read under + others, so two shards can latch concurrently. That race is benign and + deliberate — the write is idempotent (never ``True`` back to ``False``) + and the worst outcome is one extra sparse read on a shard that was about + to latch anyway, so it does not warrant a second lock. + """ + if self._full_mode_latched: + return False + count = self._sparse_read_counts.get(cache_key, 0) + 1 + if count >= self._SPARSE_PROMOTE_THRESHOLD: + self._full_mode_latched = True + # Never read again once latched; drop whatever it accumulated. + self._sparse_read_counts.clear() + return False + self._sparse_read_counts[cache_key] = count + return True + async def _load_sparse_shard_entry( self, cache_key: ShardCacheKey, @@ -1877,6 +1975,14 @@ def with_read_only(self, read_only: bool = False) -> "ShardedZarrStore": clone._shard_data_cache = self._shard_data_cache clone._pending_shard_loads = self._pending_shard_loads + # Shared by reference like the cache above: a clone that copied would + # restart counting from zero while reading through the *same* cache. + clone._sparse_read_counts = self._sparse_read_counts + # Copied by value, deliberately. The clone inherits the decision made so + # far but latches independently afterward — it has the opposite + # read/write posture, and a writable clone can never take the sparse + # path at all, so there is nothing for a shared latch to coordinate. + clone._full_mode_latched = self._full_mode_latched clone._metadata_read_cache = self._metadata_read_cache clone.array_indices = self.array_indices @@ -2098,14 +2204,15 @@ async def get( shard_lock = self._shard_locks[cache_key] async with shard_lock: cached_shard = await self._shard_data_cache.get(cache_key) - if ( - self.read_only - and self.shard_read_mode == "sparse" - and cached_shard is None - and byte_range is None - and 0 <= shard_idx < array_index.num_shards - and array_index.shard_cids[shard_idx] is not None - ): + use_sparse = ( + self.shard_read_mode != "full" + and self._sparse_read_is_eligible( + array_index, shard_idx, cached_shard, byte_range + ) + ) + if use_sparse and self.shard_read_mode == "auto": + use_sparse = self._auto_mode_wants_sparse(cache_key) + if use_sparse: chunk_cid_obj = await self._load_sparse_shard_entry( cache_key, shard_idx, @@ -2386,6 +2493,9 @@ async def _clear_v2_unlocked(self) -> None: pending_load.set() self._pending_shard_loads.clear() await self._shard_data_cache.clear() + # A cleared store has a new access pattern to learn. + self._sparse_read_counts.clear() + self._full_mode_latched = False self._root_obj["metadata"] = {} self._root_obj["arrays"] = {} self.array_indices.clear() @@ -2739,6 +2849,10 @@ async def _migrate_v1_to_v2_unlocked(self, primary_array_path: str) -> str: await self._flush_unlocked() await self._shard_data_cache.clear() + # Cache keys change shape from int to tuple[str, int] across the + # migration, so stale counter entries would be unreachable garbage. + self._sparse_read_counts.clear() + self._full_mode_latched = False source_array_path = self._infer_v1_migration_source_array_path(normalized_path) old_metadata = dict(self._root_obj.get("metadata", {})) diff --git a/tests/test_benchmark_stores.py b/tests/test_benchmark_stores.py index dd1865f..9b5a7fb 100644 --- a/tests/test_benchmark_stores.py +++ b/tests/test_benchmark_stores.py @@ -335,7 +335,9 @@ async def test_benchmark_sharded_store( with _Phase("ShardedZarrStore read") as phase: # Explicit: this phase is a whole-array scan, which is what "full" - # is for. The default is "sparse", tuned for point reads instead. + # is for. The default is "auto", which would reach the same place + # via the latch but pay a detection cost first -- pinned here so + # this measures steady-state scan cost, not the detection ramp. read_store = await ShardedZarrStore.open( cas=cas, read_only=True, root_cid=root_cid, shard_read_mode="full" ) @@ -438,23 +440,34 @@ async def test_sharded_full_mode_read_batches_chunks_into_shard_fetches( async def test_sparse_mode_trades_scan_cost_for_point_read_latency( - create_ipfs: tuple[str, str], report: Any + create_ipfs: tuple[str, str], report: Any, monkeypatch: pytest.MonkeyPatch ) -> None: - """Characterize the default read mode, so its cost is not mistaken for a bug. - - ``shard_read_mode`` defaults to ``"sparse"`` (PR #87): a read-only cache miss - with no byte range fetches just the requested entry rather than decoding the - whole shard. That is a large win for point reads over long time ranges -- the - AEGIS workload the mode was built for -- and a deliberate loss on a full - scan, where every entry is wanted anyway and per-chunk fetches add up. - - Asserted here so the trade-off is visible and intentional: a full scan in - sparse mode legitimately costs *more* than in full mode. Anyone benchmarking - a whole-array read should set ``shard_read_mode="full"``. + """Characterize each read mode on a scan, so costs are not mistaken for bugs. + + ``"sparse"`` (PR #87) resolves a read-only cache miss by fetching just the + requested entry rather than decoding the whole shard. That is a large win + for point reads over long time ranges -- the AEGIS workload the mode was + built for -- and a deliberate loss on a full scan, where every entry is + wanted anyway and per-chunk fetches add up. + + Asserted here so the trade-off stays visible: a full scan in sparse mode + legitimately costs *more* than in full mode. + + The ``"auto"`` arm is why that penalty is no longer the default: it starts + sparse, detects the scan, and latches to full decodes, so it lands between + the two rather than paying the sparse penalty all the way through. It does + not match ``"full"`` outright, and should not -- the reads before the latch + trips are genuinely sparse. That gap is the bounded detection cost. """ rpc, gateway = create_ipfs ds = _make_dataset("temp", periods=100, time_chunk=5) + # This dataset is far smaller than a real one: a whole-array scan touches + # ~20 chunks per shard, so the production threshold of 256 would never trip + # and the "auto" arm below would silently measure plain sparse mode. Scale + # the detector to the fixture so the arm exercises the latch it claims to. + monkeypatch.setattr(ShardedZarrStore, "_SPARSE_PROMOTE_THRESHOLD", 8) + async with KuboCAS(rpc_base_url=rpc, gateway_base_url=gateway) as cas: store = await ShardedZarrStore.open( cas=cas, @@ -467,7 +480,7 @@ async def test_sparse_mode_trades_scan_cost_for_point_read_latency( root_cid = await store.flush() results: dict[str, BenchmarkResult] = {} - for mode in ("sparse", "full"): + for mode in ("sparse", "full", "auto"): with _Phase(f"full scan @ shard_read_mode={mode}") as phase: read_store = await ShardedZarrStore.open( cas=cas, read_only=True, root_cid=root_cid, shard_read_mode=mode @@ -476,18 +489,89 @@ async def test_sparse_mode_trades_scan_cost_for_point_read_latency( scanned.load() results[mode] = report(phase.result) xr.testing.assert_identical(ds, scanned) + if mode == "auto": + assert read_store._full_mode_latched, ( + "a whole-array scan should have latched auto mode to full " + "decodes; if it did not, the threshold is now above the " + "per-shard chunk count and auto has silently become sparse" + ) assert results["sparse"].cas_loads > results["full"].cas_loads, ( "expected sparse mode to cost more on a full scan; if this now holds " "the other way, sparse decoding has changed and the default may want " "revisiting" ) + # The whole point of auto: pay a bounded detection cost, then stop paying + # the sparse scan penalty. Asserted against sparse rather than pinned to + # full because auto is *expected* to cost a little more than full -- the + # reads before the latch trips are genuinely sparse. + assert results["auto"].cas_loads < results["sparse"].cas_loads, ( + f"auto scan cost {results['auto'].cas_loads:.0f} CAS loads vs sparse's " + f"{results['sparse'].cas_loads:.0f}; auto should detect the scan and " + "latch to full decodes rather than staying sparse throughout" + ) -async def test_sharded_store_defaults_to_sparse_read_mode( +async def test_auto_mode_matches_sparse_on_point_reads( + create_ipfs: tuple[str, str], report: Any +) -> None: + """The other half of the auto claim: no scan, no latch, no added cost. + + ``test_sparse_mode_trades_scan_cost_for_point_read_latency`` shows auto + escaping the sparse penalty on a scan. This shows it does not *pay* anything + for that safety net on the workload sparse exists to serve -- a handful of + scattered point reads must cost exactly what plain sparse costs, and must + leave the store unlatched. + """ + rpc, gateway = create_ipfs + ds = _make_dataset("temp", periods=100, time_chunk=5) + + async with KuboCAS(rpc_base_url=rpc, gateway_base_url=gateway) as cas: + store = await ShardedZarrStore.open( + cas=cas, + read_only=False, + array_shape=tuple(ds.sizes.values()), + chunk_shape=_chunk_shape(ds), + chunks_per_shard=50, + ) + ds.to_zarr(store=store, mode="w") + root_cid = await store.flush() + + results: dict[str, BenchmarkResult] = {} + for mode in ("sparse", "auto"): + with _Phase(f"point reads @ shard_read_mode={mode}") as phase: + read_store = await ShardedZarrStore.open( + cas=cas, read_only=True, root_cid=root_cid, shard_read_mode=mode + ) + opened = xr.open_zarr(store=read_store) + for step in (0, 25, 50, 75, 99): + opened["temp"].isel(time=step).load() + results[mode] = report(phase.result) + + assert not read_store._full_mode_latched, ( + "scattered point reads are not a scan and must not latch auto mode " + "to full decodes -- that would hand this workload the whole-shard " + "cost sparse mode exists to avoid" + ) + + assert results["auto"].cas_loads == results["sparse"].cas_loads, ( + f"auto cost {results['auto'].cas_loads:.0f} CAS loads on point reads vs " + f"sparse's {results['sparse'].cas_loads:.0f}; below the threshold auto " + "must be byte-for-byte the sparse path" + ) + + +async def test_sharded_store_defaults_to_auto_read_mode( create_ipfs: tuple[str, str], ) -> None: - """Pin the default, since it determines which trade-off users get.""" + """Pin the default, since it determines which trade-off users get. + + ``"auto"`` rather than ``"sparse"``: no caller in this repo passes + ``shard_read_mode`` explicitly, so the default is what essentially everyone + runs, and it should not hand an unbounded scan penalty to callers who never + learned the knob exists. ``"auto"`` matches ``"sparse"`` below the threshold + and escapes it above -- see the two benchmarks above. + """ rpc, gateway = create_ipfs async with KuboCAS(rpc_base_url=rpc, gateway_base_url=gateway) as cas: @@ -499,7 +583,7 @@ async def test_sharded_store_defaults_to_sparse_read_mode( chunks_per_shard=4, ) - assert store.shard_read_mode == "sparse" + assert store.shard_read_mode == "auto" async def test_read_cache_prevents_duplicate_cid_fetches( diff --git a/tests/test_sharded_zarr_store_v2.py b/tests/test_sharded_zarr_store_v2.py index 0b4ecc2..b293c4f 100644 --- a/tests/test_sharded_zarr_store_v2.py +++ b/tests/test_sharded_zarr_store_v2.py @@ -269,12 +269,14 @@ async def test_read_only_get_defaults_to_sparse_shard_decode() -> None: root_cid = await store.flush() read_store = await ShardedZarrStore.open(cas=cas, read_only=True, root_cid=root_cid) - assert read_store.shard_read_mode == "sparse" + assert read_store.shard_read_mode == "auto" buffer = await read_store.get("a/c/0", proto) assert buffer is not None assert buffer.to_bytes() == b"value" + # Below the threshold "auto" is the sparse path, so nothing is cached. assert await read_store._shard_data_cache.get(("a", 0)) is None + assert read_store._full_mode_latched is False @pytest.mark.asyncio @@ -287,6 +289,247 @@ async def test_shard_read_mode_rejects_unknown_value() -> None: ) +async def _seed_auto_mode_store( + cas: LocalCIDCAS, *, num_chunks: int, chunks_per_shard: int +) -> str: + """Write a single v2 array of ``num_chunks`` 1-element chunks, return its CID.""" + proto = zarr.core.buffer.default_buffer_prototype() + store = await ShardedZarrStore.open( + cas=cas, + read_only=False, + chunks_per_shard=chunks_per_shard, + manifest_version=SHARDED_ZARR_V2, + ) + metadata = { + "zarr_format": 3, + "node_type": "array", + "shape": [num_chunks], + "data_type": "uint8", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [1]}}, + } + await store.set( + "a/zarr.json", proto.buffer.from_bytes(json.dumps(metadata).encode()) + ) + for idx in range(num_chunks): + await store.set(f"a/c/{idx}", proto.buffer.from_bytes(f"v{idx}".encode())) + return str(await store.flush()) + + +@pytest.mark.asyncio +async def test_auto_mode_stays_sparse_below_threshold( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Under the threshold, "auto" must behave exactly like "sparse".""" + monkeypatch.setattr(ShardedZarrStore, "_SPARSE_PROMOTE_THRESHOLD", 3) + cas = LocalCIDCAS() + proto = zarr.core.buffer.default_buffer_prototype() + root_cid = await _seed_auto_mode_store(cas, num_chunks=4, chunks_per_shard=4) + + store = await ShardedZarrStore.open( + cas=cas, read_only=True, root_cid=root_cid, shard_read_mode="auto" + ) + + async def fail_full_decode(*args: object, **kwargs: object) -> None: + raise AssertionError("full shard decode should not run below the threshold") + + monkeypatch.setattr(store, "_fetch_and_cache_full_shard", fail_full_decode) + + for _ in range(ShardedZarrStore._SPARSE_PROMOTE_THRESHOLD - 1): + buffer = await store.get("a/c/2", proto) + assert buffer is not None + assert buffer.to_bytes() == b"v2" + + assert store._full_mode_latched is False + assert await store._shard_data_cache.get(("a", 0)) is None + assert store._sparse_read_counts[("a", 0)] == 2 + + +@pytest.mark.asyncio +async def test_auto_mode_latches_to_full_after_threshold( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The Nth sparse read of one shard latches and clears the counter.""" + monkeypatch.setattr(ShardedZarrStore, "_SPARSE_PROMOTE_THRESHOLD", 3) + cas = LocalCIDCAS() + proto = zarr.core.buffer.default_buffer_prototype() + root_cid = await _seed_auto_mode_store(cas, num_chunks=4, chunks_per_shard=4) + + store = await ShardedZarrStore.open( + cas=cas, read_only=True, root_cid=root_cid, shard_read_mode="auto" + ) + + for _ in range(ShardedZarrStore._SPARSE_PROMOTE_THRESHOLD): + buffer = await store.get("a/c/2", proto) + assert buffer is not None + assert buffer.to_bytes() == b"v2" + + assert store._full_mode_latched is True + # The latching read full-decodes, so the shard is now cached. + assert await store._shard_data_cache.get(("a", 0)) is not None + # Never consulted again once latched, so it is dropped. + assert store._sparse_read_counts == {} + + +@pytest.mark.asyncio +async def test_auto_mode_latch_is_store_wide( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A shard that never sparse-read at all inherits the latch immediately. + + This is what distinguishes the store-wide latch from jaxray's per-shard + promotion: shard 1 must not restart its own count from zero. + """ + monkeypatch.setattr(ShardedZarrStore, "_SPARSE_PROMOTE_THRESHOLD", 3) + cas = LocalCIDCAS() + proto = zarr.core.buffer.default_buffer_prototype() + root_cid = await _seed_auto_mode_store(cas, num_chunks=4, chunks_per_shard=1) + + store = await ShardedZarrStore.open( + cas=cas, read_only=True, root_cid=root_cid, shard_read_mode="auto" + ) + for _ in range(ShardedZarrStore._SPARSE_PROMOTE_THRESHOLD): + assert await store.get("a/c/0", proto) is not None + assert store._full_mode_latched is True + + async def fail_sparse_decode(*args: object, **kwargs: object) -> None: + raise AssertionError("sparse decode should not run once latched") + + monkeypatch.setattr(store, "_load_sparse_shard_entry", fail_sparse_decode) + + # A different, never-before-read shard: full path on its very first read. + buffer = await store.get("a/c/1", proto) + assert buffer is not None + assert buffer.to_bytes() == b"v1" + assert await store._shard_data_cache.get(("a", 1)) is not None + + +@pytest.mark.asyncio +async def test_auto_mode_latch_survives_cache_eviction( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Eviction cannot un-latch: the store stays in full mode.""" + monkeypatch.setattr(ShardedZarrStore, "_SPARSE_PROMOTE_THRESHOLD", 3) + cas = LocalCIDCAS() + proto = zarr.core.buffer.default_buffer_prototype() + root_cid = await _seed_auto_mode_store(cas, num_chunks=4, chunks_per_shard=4) + + store = await ShardedZarrStore.open( + cas=cas, read_only=True, root_cid=root_cid, shard_read_mode="auto" + ) + for _ in range(ShardedZarrStore._SPARSE_PROMOTE_THRESHOLD): + assert await store.get("a/c/2", proto) is not None + assert store._full_mode_latched is True + + await store._shard_data_cache.discard(("a", 0)) + assert await store._shard_data_cache.get(("a", 0)) is None + + async def fail_sparse_decode(*args: object, **kwargs: object) -> None: + raise AssertionError("sparse decode should not run once latched") + + monkeypatch.setattr(store, "_load_sparse_shard_entry", fail_sparse_decode) + + buffer = await store.get("a/c/2", proto) + assert buffer is not None + assert buffer.to_bytes() == b"v2" + assert await store._shard_data_cache.get(("a", 0)) is not None + + +@pytest.mark.asyncio +async def test_auto_mode_does_not_latch_on_distinct_shards( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reads spread across distinct shards are point reads, not a scan. + + Detection is per-shard precisely so this workload keeps the sparse path; + a global sparse-read counter would wrongly latch it to full mode. + """ + monkeypatch.setattr(ShardedZarrStore, "_SPARSE_PROMOTE_THRESHOLD", 3) + cas = LocalCIDCAS() + proto = zarr.core.buffer.default_buffer_prototype() + root_cid = await _seed_auto_mode_store(cas, num_chunks=6, chunks_per_shard=1) + + store = await ShardedZarrStore.open( + cas=cas, read_only=True, root_cid=root_cid, shard_read_mode="auto" + ) + + async def fail_full_decode(*args: object, **kwargs: object) -> None: + raise AssertionError("distinct-shard point reads should not latch") + + monkeypatch.setattr(store, "_fetch_and_cache_full_shard", fail_full_decode) + + for idx in range(6): + buffer = await store.get(f"a/c/{idx}", proto) + assert buffer is not None + assert buffer.to_bytes() == f"v{idx}".encode() + + assert store._full_mode_latched is False + assert all(count == 1 for count in store._sparse_read_counts.values()) + + +@pytest.mark.asyncio +async def test_auto_mode_ineligible_reads_do_not_advance_counter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Byte-ranged reads take the full path without counting toward the latch.""" + monkeypatch.setattr(ShardedZarrStore, "_SPARSE_PROMOTE_THRESHOLD", 3) + cas = LocalCIDCAS() + proto = zarr.core.buffer.default_buffer_prototype() + root_cid = await _seed_auto_mode_store(cas, num_chunks=4, chunks_per_shard=4) + + store = await ShardedZarrStore.open( + cas=cas, read_only=True, root_cid=root_cid, shard_read_mode="auto" + ) + buffer = await store.get("a/c/2", proto, byte_range=RangeByteRequest(0, 1)) + assert buffer is not None + assert buffer.to_bytes() == b"v" + assert store._sparse_read_counts == {} + assert store._full_mode_latched is False + + +@pytest.mark.asyncio +async def test_full_mode_never_touches_auto_mode_state() -> None: + """ "full" short-circuits before the counter is ever consulted.""" + cas = LocalCIDCAS() + proto = zarr.core.buffer.default_buffer_prototype() + root_cid = await _seed_auto_mode_store(cas, num_chunks=4, chunks_per_shard=4) + + store = await ShardedZarrStore.open( + cas=cas, read_only=True, root_cid=root_cid, shard_read_mode="full" + ) + for _ in range(4): + assert await store.get("a/c/2", proto) is not None + + assert store._sparse_read_counts == {} + assert store._full_mode_latched is False + + +@pytest.mark.asyncio +async def test_auto_mode_state_threads_through_with_read_only( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The clone shares the counter dict and inherits the latch by value.""" + monkeypatch.setattr(ShardedZarrStore, "_SPARSE_PROMOTE_THRESHOLD", 3) + cas = LocalCIDCAS() + proto = zarr.core.buffer.default_buffer_prototype() + root_cid = await _seed_auto_mode_store(cas, num_chunks=4, chunks_per_shard=4) + + store = await ShardedZarrStore.open( + cas=cas, read_only=True, root_cid=root_cid, shard_read_mode="auto" + ) + assert await store.get("a/c/2", proto) is not None + + clone = store.with_read_only(False) + assert clone._sparse_read_counts is store._sparse_read_counts + assert clone._full_mode_latched is False + + for _ in range(ShardedZarrStore._SPARSE_PROMOTE_THRESHOLD - 1): + assert await store.get("a/c/2", proto) is not None + assert store._full_mode_latched is True + + latched_clone = store.with_read_only(False) + assert latched_clone._full_mode_latched is True + + def _pyramid_level(data: np.ndarray, *, coord_offset: int = 0) -> xr.Dataset: return xr.Dataset( {"FPAR": (("time", "y", "x"), data)}, From f481bb8ef1285e698b649257fc75dbefa3ae4e94 Mon Sep 17 00:00:00 2001 From: TheGreatAlgo <37487508+TheGreatAlgo@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:05:53 -0400 Subject: [PATCH 2/2] fix: improve latching --- py_hamt/sharded_zarr_store.py | 36 +++++++++++++++----- tests/test_sharded_zarr_store_v2.py | 53 ++++++++++++++++++++++++++++- 2 files changed, 80 insertions(+), 9 deletions(-) diff --git a/py_hamt/sharded_zarr_store.py b/py_hamt/sharded_zarr_store.py index b8ded42..68e088d 100644 --- a/py_hamt/sharded_zarr_store.py +++ b/py_hamt/sharded_zarr_store.py @@ -529,7 +529,14 @@ def __init__( # Latched once any single shard crosses the threshold: the caller is # scanning, so every shard gets the full path from here on. Store-wide # because the access pattern belongs to the caller, not the shard. - self._full_mode_latched: bool = False + # + # Held in a one-element list so with_read_only clones share the *cell* + # rather than a copied bool. They already share the counters and the + # cache, and a clone that latched would otherwise clear those shared + # counters while leaving its siblings believing they were still sparse + # -- so the next clone would resume sparse reads with no counter left + # to re-earn promotion. See _full_mode_latched. + self._full_mode_latched_cell: List[bool] = [False] self._metadata_read_cache: Dict[str, bytes] = {} self.array_indices: Dict[str, ArrayIndex] = {} @@ -1542,6 +1549,19 @@ async def _fetch_and_cache_full_shard( f"Failed to fetch shard {shard_idx} after {max_retries} attempts: {e}" ) from e + @property + def _full_mode_latched(self) -> bool: + """Whether ``auto`` mode has committed this store to full decodes. + + Backed by a cell shared with every ``with_read_only`` clone, so a latch + earned by one clone is immediately visible to all of them. + """ + return self._full_mode_latched_cell[0] + + @_full_mode_latched.setter + def _full_mode_latched(self, value: bool) -> None: + self._full_mode_latched_cell[0] = value + def _sparse_read_is_eligible( self, array_index: ArrayIndex, @@ -1975,14 +1995,14 @@ def with_read_only(self, read_only: bool = False) -> "ShardedZarrStore": clone._shard_data_cache = self._shard_data_cache clone._pending_shard_loads = self._pending_shard_loads - # Shared by reference like the cache above: a clone that copied would - # restart counting from zero while reading through the *same* cache. + # Both shared by reference, like the cache above. The counters must be + # shared so a clone does not restart counting while reading through the + # *same* cache; the latch cell must be shared for the same reason in + # reverse -- latching clears the shared counters, so a sibling holding + # a copied False would resume sparse reads with nothing left to re-earn + # promotion from. clone._sparse_read_counts = self._sparse_read_counts - # Copied by value, deliberately. The clone inherits the decision made so - # far but latches independently afterward — it has the opposite - # read/write posture, and a writable clone can never take the sparse - # path at all, so there is nothing for a shared latch to coordinate. - clone._full_mode_latched = self._full_mode_latched + clone._full_mode_latched_cell = self._full_mode_latched_cell clone._metadata_read_cache = self._metadata_read_cache clone.array_indices = self.array_indices diff --git a/tests/test_sharded_zarr_store_v2.py b/tests/test_sharded_zarr_store_v2.py index b293c4f..8bcfbce 100644 --- a/tests/test_sharded_zarr_store_v2.py +++ b/tests/test_sharded_zarr_store_v2.py @@ -507,7 +507,7 @@ async def test_full_mode_never_touches_auto_mode_state() -> None: async def test_auto_mode_state_threads_through_with_read_only( monkeypatch: pytest.MonkeyPatch, ) -> None: - """The clone shares the counter dict and inherits the latch by value.""" + """The clone shares both the counter dict and the latch cell.""" monkeypatch.setattr(ShardedZarrStore, "_SPARSE_PROMOTE_THRESHOLD", 3) cas = LocalCIDCAS() proto = zarr.core.buffer.default_buffer_prototype() @@ -520,16 +520,67 @@ async def test_auto_mode_state_threads_through_with_read_only( clone = store.with_read_only(False) assert clone._sparse_read_counts is store._sparse_read_counts + assert clone._full_mode_latched_cell is store._full_mode_latched_cell assert clone._full_mode_latched is False for _ in range(ShardedZarrStore._SPARSE_PROMOTE_THRESHOLD - 1): assert await store.get("a/c/2", proto) is not None assert store._full_mode_latched is True + # Shared cell: the pre-existing clone sees the latch, not just later ones. + assert clone._full_mode_latched is True latched_clone = store.with_read_only(False) assert latched_clone._full_mode_latched is True +@pytest.mark.asyncio +async def test_auto_mode_latch_is_visible_to_sibling_clones( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A latch earned by one clone must reach its siblings. + + Regression: the latch was originally copied by value while the counters and + cache were shared by reference. Latching in one read-only clone cleared the + *shared* counters but left a sibling clone of the same writable parent + holding a stale ``False`` -- so the sibling resumed sparse reads on evicted + shards with no counter left to re-earn promotion from, permanently. This is + reachable from stock zarr, which calls ``with_read_only(True)`` on writable + stores (zarr/storage/_common.py). + """ + monkeypatch.setattr(ShardedZarrStore, "_SPARSE_PROMOTE_THRESHOLD", 3) + cas = LocalCIDCAS() + proto = zarr.core.buffer.default_buffer_prototype() + root_cid = await _seed_auto_mode_store(cas, num_chunks=4, chunks_per_shard=4) + + parent = await ShardedZarrStore.open( + cas=cas, read_only=False, root_cid=root_cid, shard_read_mode="auto" + ) + first = parent.with_read_only(True) + for _ in range(ShardedZarrStore._SPARSE_PROMOTE_THRESHOLD): + assert await first.get("a/c/2", proto) is not None + + assert first._full_mode_latched is True + assert parent._full_mode_latched is True, "the writable parent must observe it" + assert first._sparse_read_counts == {} + + second = parent.with_read_only(True) + assert second._full_mode_latched is True, ( + "a sibling clone inherited a stale unlatched state; it would resume " + "sparse reads with the shared counters already cleared" + ) + + # Evict, then confirm the sibling takes the full path rather than sparse. + await second._shard_data_cache.discard(("a", 0)) + + async def fail_sparse_decode(*args: object, **kwargs: object) -> None: + raise AssertionError("a latched sibling must not sparse-decode") + + monkeypatch.setattr(second, "_load_sparse_shard_entry", fail_sparse_decode) + buffer = await second.get("a/c/2", proto) + assert buffer is not None + assert buffer.to_bytes() == b"v2" + + def _pyramid_level(data: np.ndarray, *, coord_offset: int = 0) -> xr.Dataset: return xr.Dataset( {"FPAR": (("time", "y", "x"), data)},