From 1d8233946e524a66d792f2403154917a53f100b9 Mon Sep 17 00:00:00 2001 From: Tom Nicholas Date: Fri, 8 Aug 2025 15:39:00 +0100 Subject: [PATCH 1/7] move definition of LatencyStore --- src/zarr/testing/store.py | 63 --------------------------------------- 1 file changed, 63 deletions(-) diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index d2946705f0..66eae91552 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -1,17 +1,14 @@ from __future__ import annotations -import asyncio import pickle from abc import abstractmethod from typing import TYPE_CHECKING, Generic, TypeVar -from zarr.storage import WrapperStore if TYPE_CHECKING: from typing import Any from zarr.abc.store import ByteRequest - from zarr.core.buffer.core import BufferPrototype import pytest @@ -525,63 +522,3 @@ async def test_set_if_not_exists(self, store: S) -> None: result = await store.get("k2", default_buffer_prototype()) assert result == new - - -class LatencyStore(WrapperStore[Store]): - """ - A wrapper class that takes any store class in its constructor and - adds latency to the `set` and `get` methods. This can be used for - performance testing. - """ - - get_latency: float - set_latency: float - - def __init__(self, cls: Store, *, get_latency: float = 0, set_latency: float = 0) -> None: - self.get_latency = float(get_latency) - self.set_latency = float(set_latency) - self._store = cls - - async def set(self, key: str, value: Buffer) -> None: - """ - Add latency to the ``set`` method. - - Calls ``asyncio.sleep(self.set_latency)`` before invoking the wrapped ``set`` method. - - Parameters - ---------- - key : str - The key to set - value : Buffer - The value to set - - Returns - ------- - None - """ - await asyncio.sleep(self.set_latency) - await self._store.set(key, value) - - async def get( - self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None - ) -> Buffer | None: - """ - Add latency to the ``get`` method. - - Calls ``asyncio.sleep(self.get_latency)`` before invoking the wrapped ``get`` method. - - Parameters - ---------- - key : str - The key to get - prototype : BufferPrototype - The BufferPrototype to use. - byte_range : ByteRequest, optional - An optional byte range. - - Returns - ------- - buffer : Buffer or None - """ - await asyncio.sleep(self.get_latency) - return await self._store.get(key, prototype=prototype, byte_range=byte_range) From c8812689ae53b6acce771b73e9c3eb7b860c7222 Mon Sep 17 00:00:00 2001 From: Tom Nicholas Date: Fri, 8 Aug 2025 15:39:39 +0100 Subject: [PATCH 2/7] expose publicly --- src/zarr/storage/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/zarr/storage/__init__.py b/src/zarr/storage/__init__.py index 00df50214f..4b09a44cbb 100644 --- a/src/zarr/storage/__init__.py +++ b/src/zarr/storage/__init__.py @@ -6,6 +6,7 @@ from zarr.errors import ZarrDeprecationWarning from zarr.storage._common import StoreLike, StorePath from zarr.storage._fsspec import FsspecStore +from zarr.storage._latency import LatencyStore from zarr.storage._local import LocalStore from zarr.storage._logging import LoggingStore from zarr.storage._memory import GpuMemoryStore, MemoryStore @@ -16,6 +17,7 @@ __all__ = [ "FsspecStore", "GpuMemoryStore", + "LatencyStore", "LocalStore", "LoggingStore", "MemoryStore", From 0195167bdaefc008dff81e6beee59a703badc0d6 Mon Sep 17 00:00:00 2001 From: Tom Nicholas Date: Fri, 8 Aug 2025 15:39:49 +0100 Subject: [PATCH 3/7] fix import in tests --- tests/test_group.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_group.py b/tests/test_group.py index e5cfe82daa..3bcc4a8637 100644 --- a/tests/test_group.py +++ b/tests/test_group.py @@ -50,7 +50,7 @@ from zarr.storage import LocalStore, MemoryStore, StorePath, ZipStore from zarr.storage._common import make_store_path from zarr.storage._utils import _join_paths, normalize_path -from zarr.testing.store import LatencyStore +from zarr.storage import LatencyStore from .conftest import meta_from_array, parse_store From 60f1f5c9732910b3a211ee16a6cabc3644283fba Mon Sep 17 00:00:00 2001 From: Tom Nicholas Date: Fri, 8 Aug 2025 15:41:35 +0100 Subject: [PATCH 4/7] actually add file with implementation in --- src/zarr/storage/_latency.py | 83 ++++++++++++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 src/zarr/storage/_latency.py diff --git a/src/zarr/storage/_latency.py b/src/zarr/storage/_latency.py new file mode 100644 index 0000000000..0c5fa86d7c --- /dev/null +++ b/src/zarr/storage/_latency.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import asyncio + +from zarr.storage._wrapper import WrapperStore +from zarr.abc.store import ByteRequest, Store +from zarr.core.buffer import Buffer, BufferPrototype + + +class LatencyStore(WrapperStore[Store]): + """ + A wrapper class that takes any store class in its constructor and + adds latency to the `set` and `get` methods. This can be used for + performance testing. + + Particularly useful for testing downstream applications which will + interact with a high-latency zarr store implementation, + such as one which read from or writes to remote object storage. + For example, by using this class to wrap a ``MemoryStore`` instance, + you can (crudely) simulate the latency of reading and writing from S3 + without having to actually use the network, or a mock like MinIO. + + Parameters + ---------- + store : Store + Store to wrap + get_latency : float + Amount of latency to add to each get call, in seconds. Default is 0. + set_latency : float + Amount of latency to add to each set call, in seconds. Default is 0. + """ + + get_latency: float + set_latency: float + + def __init__(self, cls: Store, *, get_latency: float = 0, set_latency: float = 0) -> None: + self.get_latency = float(get_latency) + self.set_latency = float(set_latency) + self._store = cls + + async def set(self, key: str, value: Buffer) -> None: + """ + Add latency to the ``set`` method. + + Calls ``asyncio.sleep(self.set_latency)`` before invoking the wrapped ``set`` method. + + Parameters + ---------- + key : str + The key to set + value : Buffer + The value to set + + Returns + ------- + None + """ + await asyncio.sleep(self.set_latency) + await self._store.set(key, value) + + async def get( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + """ + Add latency to the ``get`` method. + + Calls ``asyncio.sleep(self.get_latency)`` before invoking the wrapped ``get`` method. + + Parameters + ---------- + key : str + The key to get + prototype : BufferPrototype + The BufferPrototype to use. + byte_range : ByteRequest, optional + An optional byte range. + + Returns + ------- + buffer : Buffer or None + """ + await asyncio.sleep(self.get_latency) + return await self._store.get(key, prototype=prototype, byte_range=byte_range) From 85027c3fd14db5e45011e93e94dd5fa0c9a29f6a Mon Sep 17 00:00:00 2001 From: Tom Nicholas Date: Fri, 8 Aug 2025 15:46:28 +0100 Subject: [PATCH 5/7] changelog --- changes/3359.feature.rst | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/3359.feature.rst diff --git a/changes/3359.feature.rst b/changes/3359.feature.rst new file mode 100644 index 0000000000..97a9a6e801 --- /dev/null +++ b/changes/3359.feature.rst @@ -0,0 +1 @@ +Expose `LatencyStore` as public API. \ No newline at end of file From 825cdc4653c498a54b49ee49db8d15a7461d28cc Mon Sep 17 00:00:00 2001 From: Tom Nicholas Date: Fri, 8 Aug 2025 15:48:19 +0100 Subject: [PATCH 6/7] some pre-commit fixes --- src/zarr/storage/_latency.py | 12 ++++++------ src/zarr/testing/store.py | 1 - tests/test_group.py | 3 +-- 3 files changed, 7 insertions(+), 9 deletions(-) diff --git a/src/zarr/storage/_latency.py b/src/zarr/storage/_latency.py index 0c5fa86d7c..4af96281de 100644 --- a/src/zarr/storage/_latency.py +++ b/src/zarr/storage/_latency.py @@ -2,9 +2,9 @@ import asyncio -from zarr.storage._wrapper import WrapperStore from zarr.abc.store import ByteRequest, Store from zarr.core.buffer import Buffer, BufferPrototype +from zarr.storage._wrapper import WrapperStore class LatencyStore(WrapperStore[Store]): @@ -13,11 +13,11 @@ class LatencyStore(WrapperStore[Store]): adds latency to the `set` and `get` methods. This can be used for performance testing. - Particularly useful for testing downstream applications which will - interact with a high-latency zarr store implementation, - such as one which read from or writes to remote object storage. - For example, by using this class to wrap a ``MemoryStore`` instance, - you can (crudely) simulate the latency of reading and writing from S3 + Particularly useful for testing downstream applications which will + interact with a high-latency zarr store implementation, + such as one which read from or writes to remote object storage. + For example, by using this class to wrap a ``MemoryStore`` instance, + you can (crudely) simulate the latency of reading and writing from S3 without having to actually use the network, or a mock like MinIO. Parameters diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index 66eae91552..2cca4118f4 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -4,7 +4,6 @@ from abc import abstractmethod from typing import TYPE_CHECKING, Generic, TypeVar - if TYPE_CHECKING: from typing import Any diff --git a/tests/test_group.py b/tests/test_group.py index 3bcc4a8637..31ad14d62f 100644 --- a/tests/test_group.py +++ b/tests/test_group.py @@ -47,10 +47,9 @@ ZarrDeprecationWarning, ZarrUserWarning, ) -from zarr.storage import LocalStore, MemoryStore, StorePath, ZipStore +from zarr.storage import LatencyStore, LocalStore, MemoryStore, StorePath, ZipStore from zarr.storage._common import make_store_path from zarr.storage._utils import _join_paths, normalize_path -from zarr.storage import LatencyStore from .conftest import meta_from_array, parse_store From eb4aa285ed9f89bd42ff088a32b6091f8a29afba Mon Sep 17 00:00:00 2001 From: Tom Nicholas Date: Fri, 8 Aug 2025 15:52:33 +0100 Subject: [PATCH 7/7] more pre-commit --- src/zarr/storage/_latency.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/zarr/storage/_latency.py b/src/zarr/storage/_latency.py index 4af96281de..12ad9a85a5 100644 --- a/src/zarr/storage/_latency.py +++ b/src/zarr/storage/_latency.py @@ -1,11 +1,14 @@ from __future__ import annotations import asyncio +from typing import TYPE_CHECKING from zarr.abc.store import ByteRequest, Store -from zarr.core.buffer import Buffer, BufferPrototype from zarr.storage._wrapper import WrapperStore +if TYPE_CHECKING: + from zarr.core.buffer import Buffer, BufferPrototype + class LatencyStore(WrapperStore[Store]): """