From 1a20c4c50a424dc3a35ae55efc7c4fe83d1ceb92 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 29 Jul 2026 13:51:13 +0200 Subject: [PATCH] fix: make ManagedMemoryStore/GpuMemoryStore sync methods parity-safe ManagedMemoryStore inherited get_sync/set_sync/delete_sync from MemoryStore, which use the raw key, while every async method prefixed keys with self.path. Any code taking the sync fast path (e.g. FusedCodecPipeline) wrote/read chunks outside the store's path prefix, so a fresh handle re-reading through the prefix silently got fill values. Override the three sync methods to prefix like their async counterparts. GpuMemoryStore.set_sync gets the same treatment: it now converts its value to a gpu.Buffer like set does, preserving the store's all-values-are-gpu invariant for the sync API. Also fix ManagedMemoryStore.get_partial_values, which applied self.path twice whenever path was non-empty (it pre-prefixed keys, then delegated to MemoryStore.get_partial_values, which itself dispatches through the already-overridden self.get) -- this made it return None for every key. Discovered via the strengthened test fixture below. Add sync/async parity laws to the shared StoreTests suite so every store subclass exercises this invariant: set through one API and read through the other (including byte_range variants), and confirm delete_sync is visible to async get. These are the tests that would have caught the ManagedMemoryStore bug. TestManagedMemoryStore's raw set/get test helpers now respect self.path, and store_kwargs uses a non-empty path, so prefix handling is actually exercised instead of passing vacuously. Add an end-to-end regression with FusedCodecPipeline writing to a ManagedMemoryStore(path=...) sharing a dict with a fresh handle. Add a LocalStore.delete_sync directory-branch test. Assisted-by: ClaudeCode:claude-sonnet-5 --- changes/4204.bugfix.md | 16 +++++++ src/zarr/storage/_memory.py | 44 +++++++++++++++---- src/zarr/testing/store.py | 65 ++++++++++++++++++++++++++++ tests/test_store/test_local.py | 14 ++++++ tests/test_store/test_memory.py | 76 ++++++++++++++++++++++++++++++--- 5 files changed, 202 insertions(+), 13 deletions(-) create mode 100644 changes/4204.bugfix.md diff --git a/changes/4204.bugfix.md b/changes/4204.bugfix.md new file mode 100644 index 0000000000..90101d1059 --- /dev/null +++ b/changes/4204.bugfix.md @@ -0,0 +1,16 @@ +`ManagedMemoryStore.get_sync`/`set_sync`/`delete_sync` now apply the store's +`path` prefix, matching the async `get`/`set`/`delete` methods. Previously the +sync methods were inherited unchanged from `MemoryStore` and used the raw key, +so code that takes the sync fast path (e.g. `FusedCodecPipeline`) would read +and write chunks outside the store's `path` prefix, silently returning fill +values when the data was re-read through a fresh handle. `GpuMemoryStore.set_sync` +now converts its value to a `gpu.Buffer`, matching `set`, so writes through the +sync API keep the store's all-values-are-GPU invariant. Also fixed +`ManagedMemoryStore.get_partial_values` applying its `path` prefix twice +whenever `path` is non-empty, which made it always return `None` for every +requested key. + +The shared store test suite (`zarr.testing.store.StoreTests`) gained +sync/async parity checks — comparing sync and async observations of the same +key on the same store instance, including with a `byte_range` — so every +store subclass now exercises this invariant. diff --git a/src/zarr/storage/_memory.py b/src/zarr/storage/_memory.py index 97dd355515..f42c38df69 100644 --- a/src/zarr/storage/_memory.py +++ b/src/zarr/storage/_memory.py @@ -314,6 +314,19 @@ async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None gpu_value = value if isinstance(value, gpu.Buffer) else gpu.Buffer.from_buffer(value) await super().set(key, gpu_value, byte_range=byte_range) + def set_sync(self, key: str, value: Buffer) -> None: + # docstring inherited + self._check_writable() + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"GpuMemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." + ) + # Convert to gpu.Buffer, mirroring `set` above: every value in this store's + # backing dict must be a gpu.Buffer, regardless of which API wrote it. + gpu_value = value if isinstance(value, gpu.Buffer) else gpu.Buffer.from_buffer(value) + super().set_sync(key, gpu_value) + # ----------------------------------------------------------------------------- # ManagedMemoryStore and its registry @@ -572,25 +585,40 @@ def from_url(cls, url: str, *, read_only: bool = False) -> ManagedMemoryStore: # Override MemoryStore methods to use path prefix and check process - async def get( + def get_sync( self, key: str, + *, prototype: BufferPrototype | None = None, byte_range: ByteRequest | None = None, ) -> Buffer | None: # docstring inherited - return await super().get( + return super().get_sync( _join_paths([self.path, key]), prototype=prototype, byte_range=byte_range ) - async def get_partial_values( + def set_sync(self, key: str, value: Buffer) -> None: + # docstring inherited + super().set_sync(_join_paths([self.path, key]), value) + + def delete_sync(self, key: str) -> None: + # docstring inherited + super().delete_sync(_join_paths([self.path, key])) + + async def get( self, - prototype: BufferPrototype, - key_ranges: Iterable[tuple[str, ByteRequest | None]], - ) -> list[Buffer | None]: + key: str, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: # docstring inherited - key_ranges = [(_join_paths([self.path, key]), byte_range) for key, byte_range in key_ranges] - return await super().get_partial_values(prototype, key_ranges) + return await super().get( + _join_paths([self.path, key]), prototype=prototype, byte_range=byte_range + ) + + # get_partial_values is intentionally NOT overridden here: MemoryStore.get_partial_values + # dispatches per-key through `self.get`, which already resolves to the override above. + # Re-prefixing the keys here as well would apply `self.path` twice. async def exists(self, key: str) -> bool: # docstring inherited diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index 46287ccffb..d7011440e0 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -618,6 +618,71 @@ def test_delete_sync_missing(self, store: S) -> None: # should not raise deleter.delete_sync("nonexistent_sync") + # ------------------------------------------------------------------- + # Sync/async parity laws + # ------------------------------------------------------------------- + # A store's sync and async methods must observe the same key the same + # way. This is stronger than the individual test_get_sync/test_set_sync/ + # test_delete_sync tests above: those write and read back through the + # *same* API (sync-only or, via `self.set`/`self.get`, bypassing the + # store entirely), so a sync method that skips logic the async method + # applies (e.g. a path prefix) can still pass them. These laws write + # through one API and observe through the other. + + @pytest.mark.parametrize("direction", ["set_async_get_sync", "set_sync_get_async"]) + async def test_sync_async_set_get_parity(self, store: S, direction: str) -> None: + setter = self._require_set_sync(store) + getter = self._require_get_sync(store) + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "parity_set_get" + if direction == "set_async_get_sync": + await store.set(key, data_buf) + result = getter.get_sync(key) + else: + setter.set_sync(key, data_buf) + result = await store.get(key, prototype=default_buffer_prototype()) + assert result is not None + assert_bytes_equal(result, data_buf) + + async def test_delete_sync_visible_to_async_get(self, store: S) -> None: + deleter = self._require_delete_sync(store) + if not store.supports_deletes: + pytest.skip("store does not support deletes") + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "parity_delete" + await store.set(key, data_buf) + deleter.delete_sync(key) + result = await store.get(key, prototype=default_buffer_prototype()) + assert result is None + + @pytest.mark.parametrize( + "byte_range", + [ + None, + RangeByteRequest(1, 4), + OffsetByteRequest(1), + SuffixByteRequest(1), + RangeByteRequest(10, 20), + ], + ids=["none", "range", "offset", "suffix", "range-past-eof"], + ) + async def test_get_sync_byte_range_parity( + self, store: S, byte_range: ByteRequest | None + ) -> None: + getter = self._require_get_sync(store) + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "parity_byte_range" + await store.set(key, data_buf) + sync_result = getter.get_sync(key, byte_range=byte_range) + async_result = await store.get( + key, prototype=default_buffer_prototype(), byte_range=byte_range + ) + if async_result is None: + assert sync_result is None + else: + assert sync_result is not None + assert_bytes_equal(sync_result, async_result) + class LatencyStore(WrapperStore[Store]): """ diff --git a/tests/test_store/test_local.py b/tests/test_store/test_local.py index f65f618d65..61e48a269f 100644 --- a/tests/test_store/test_local.py +++ b/tests/test_store/test_local.py @@ -46,6 +46,20 @@ async def test_empty_with_empty_subdir(self, store: LocalStore) -> None: (store.root / "foo/bar").mkdir(parents=True) assert await store.is_empty("") + def test_delete_sync_directory(self, store: LocalStore) -> None: + """`delete_sync` on a key that is a directory must remove the whole tree. + + Mirrors the async `delete_dir` behavior: deleting `"foo"` where + `"foo"` is a directory containing further nested paths should remove + everything under it, not just fail or delete a single file. + """ + (store.root / "foo" / "bar").mkdir(parents=True) + (store.root / "foo" / "bar" / "baz").write_bytes(b"data") + + store.delete_sync("foo") + + assert not (store.root / "foo").exists() + def test_creates_new_directory(self, tmp_path: pathlib.Path) -> None: target = tmp_path.joinpath("a", "b", "c") assert not target.exists() diff --git a/tests/test_store/test_memory.py b/tests/test_store/test_memory.py index 36265423e6..a976f3738e 100644 --- a/tests/test_store/test_memory.py +++ b/tests/test_store/test_memory.py @@ -11,6 +11,7 @@ from zarr.core.buffer import Buffer, cpu, default_buffer_prototype, gpu from zarr.errors import ZarrUserWarning from zarr.storage import GpuMemoryStore, ManagedMemoryStore, MemoryStore +from zarr.storage._utils import _join_paths from zarr.testing.store import StoreTests from zarr.testing.utils import gpu_test @@ -233,31 +234,48 @@ def test_from_dict(self) -> None: for v in result._store_dict.values(): assert type(v) is gpu.Buffer + def test_set_sync_converts_to_gpu_buffer(self, store: GpuMemoryStore) -> None: + """`set_sync` must convert its value to a `gpu.Buffer`, mirroring `set`. + + `GpuMemoryStore`'s invariant is that every stored value is a + `gpu.Buffer`. Without this override, the inherited `MemoryStore.set_sync` + would store the CPU buffer it was given as-is, breaking that invariant + for whichever code path (e.g. the fused pipeline) uses the sync API. + """ + cpu_value = cpu.Buffer.from_bytes(b"aaaa") + msg = "Creating a zarr.buffer.gpu.Buffer with an array that does not support the __cuda_array_interface__ for zero-copy transfers, falling back to slow copy based path" + with pytest.warns(ZarrUserWarning, match=msg): + store.set_sync("k", cpu_value) + assert type(store._store_dict["k"]) is gpu.Buffer + class TestManagedMemoryStore(StoreTests[ManagedMemoryStore, cpu.Buffer]): store_cls = ManagedMemoryStore buffer_cls = cpu.Buffer async def set(self, store: ManagedMemoryStore, key: str, value: Buffer) -> None: - store._store_dict[key] = value + store._store_dict[_join_paths([store.path, key])] = value async def get(self, store: ManagedMemoryStore, key: str) -> Buffer: - return store._store_dict[key] + return store._store_dict[_join_paths([store.path, key])] @pytest.fixture def store_kwargs(self, request: pytest.FixtureRequest) -> dict[str, Any]: # Use a unique name per test to avoid sharing state between tests # but ensure the name is deterministic for equality tests # Replace '/' with '-' since store names cannot contain '/' + # A non-empty path exercises prefix handling; a store with an + # unprefixed key in its backing dict would pass these tests + # vacuously with path="". sanitized_name = request.node.name.replace("/", "-") - return {"name": f"test-{sanitized_name}"} + return {"name": f"test-{sanitized_name}", "path": "prefix"} @pytest.fixture async def store(self, store_kwargs: dict[str, Any]) -> ManagedMemoryStore: return self.store_cls(**store_kwargs) def test_store_repr(self, store: ManagedMemoryStore) -> None: - assert str(store) == f"memory://{store.name}" + assert str(store) == _join_paths([f"memory://{store.name}", store.path]) async def test_serializable_store(self, store: ManagedMemoryStore) -> None: """ @@ -383,7 +401,10 @@ def test_from_url(self, store: ManagedMemoryStore) -> None: def test_from_url_with_path(self, store: ManagedMemoryStore) -> None: """Test that from_url extracts path component from URL.""" - url = f"{store}/some/path" + # Reconnect to the fixture's dict via its name, but with an empty + # path, so appending "/some/path" below yields exactly that path. + base = ManagedMemoryStore(name=store.name) + url = f"{base}/some/path" store2 = ManagedMemoryStore.from_url(url) assert store2._store_dict is store._store_dict assert store2.path == "some/path" @@ -512,3 +533,48 @@ def test_garbage_collection(self) -> None: # URL should no longer resolve with pytest.raises(ValueError, match="garbage collected"): ManagedMemoryStore.from_url(url) + + def test_sync_methods_respect_path_prefix(self) -> None: + """`get_sync`/`set_sync`/`delete_sync` must prefix keys with `self.path`, + exactly like the async `get`/`set`/`delete` methods. + + `ManagedMemoryStore` used to inherit these from `MemoryStore`, which + writes/reads the raw key. Two stores sharing a dict with different + `path` values would then cross-talk through the sync API. + """ + store = ManagedMemoryStore(name="sync-prefix-test", path="subdir") + data_buf = self.buffer_cls.from_bytes(b"value") + + store.set_sync("key", data_buf) + assert "subdir/key" in store._store_dict + assert "key" not in store._store_dict + + result = store.get_sync("key") + assert result is not None + assert result.to_bytes() == b"value" + + store.delete_sync("key") + assert "subdir/key" not in store._store_dict + + def test_fused_pipeline_respects_path_prefix(self) -> None: + """End-to-end regression: the fused pipeline's sync store fast path must + write chunks under the store's path prefix. + + `FusedCodecPipeline` uses `set_sync`/`get_sync` when a store implements + the sync protocols. If those methods skip the prefix that the async + methods apply, chunk data lands outside `self.path` and a fresh handle + re-reading through the prefix silently sees fill values instead. + """ + with zarr.config.set( + {"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"} + ): + store = ManagedMemoryStore(name="fused-prefix-test", path="subdir") + arr = zarr.create_array(store, shape=(4,), chunks=(4,), dtype="uint8", zarr_format=3) + arr[:] = np.arange(4, dtype="uint8") + + bad_keys = [k for k in store._store_dict if not k.startswith("subdir/")] + assert bad_keys == [], f"keys written outside the store's path prefix: {bad_keys}" + + store2 = ManagedMemoryStore.from_url("memory://fused-prefix-test/subdir") + arr2 = zarr.open_array(store2, mode="r") + np.testing.assert_array_equal(arr2[:], np.arange(4, dtype="uint8"))