perf: benchmarking larger chunks + non-repeated values

src/zarr/testing/store.py

@@ -633,8 +633,8 @@ def __init__(
         self,
         store: Store,
         *,
-        get_latency: float | tuple[float, float] = 0,
-        set_latency: float | tuple[float, float] = 0,
+        get_latency: float | tuple[float, float] = 0.,
+        set_latency: float | tuple[float, float] = 0.,
     ) -> None:
         """
         Parameters

tests/benchmarks/test_e2e.py

@@ -4,10 +4,26 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
+from functools import lru_cache
+import os
+import numpy as np
 from typing import TYPE_CHECKING
 
 from tests.benchmarks.common import Layout
 
+
+import platform
+import subprocess
+
+def clear_cache():
+    if platform.system() == "Darwin":
+        subprocess.call(['sync', '&&', 'sudo', 'purge'])
+    elif platform.system() == "Linux":
+        subprocess.call(['sudo', 'sh', '-c', "sync; echo 3 > /proc/sys/vm/drop_caches"])
+    else:
+        raise Exception("Unsupported platform")
+
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
@@ -25,29 +41,44 @@
 from zarr.core.config import config as zarr_config
 from zarr.testing.store import LatencyStore
 
-CompressorName = Literal["gzip"] | None
+@lru_cache
+def _data(shape: tuple[int]) -> np.ndarray:
+    n = shape[0]
+    period = 256
+    noise_level = 1
+    pattern = (np.sin(np.linspace(0, 2 * np.pi, period)) * 50 + 128).round().astype(np.uint8)
+    data = np.tile(pattern, int(np.ceil(n / period)))[:n].astype(np.int16)
+    data += np.random.randint(-noise_level, noise_level + 1, size=n, dtype=np.int16)
+    return np.clip(data, 0, 255).astype(np.uint8)
+
+
+CompressorName = Literal["zstd"] | None
 
 compressors: dict[CompressorName, NamedConfig[Any, Any] | None] = {
     None: None,
-    "gzip": {"name": "gzip", "configuration": {"level": 1}},
+    # Default v3
+    "zstd": {"name": "zstd", "configuration": {"level": 0, "checksum": False}}
 }
 
 
 layouts: tuple[Layout, ...] = (
     # No shards, just 1000 chunks
-    Layout(shape=(1_000_000,), chunks=(1000,), shards=None),
+    Layout(shape=(100_000_000,), chunks=(100_000,), shards=None),
     # 1:1 chunk:shard shape, should measure overhead of sharding
-    Layout(shape=(1_000_000,), chunks=(1000,), shards=(1000,)),
-    # One shard with all the chunks, should measure overhead of handling inner shard chunks
-    Layout(shape=(1_000_000,), chunks=(100,), shards=(10000 * 100,)),
+    Layout(shape=(100_000_000,), chunks=(100_000,), shards=(100_000,)),
+    # One shard with all the chunks, should measure over/under-head of handling inner shard chunks
+    Layout(shape=(100_000_000,), chunks=(100_000,), shards=(100_000 * 1_000,)),
+    # Mixed layout balancing inner vs. outer concurrency (likely the most real-world case)
+    Layout(shape=(1_000_000_000,), chunks=(100_000,), shards=(100_000 * 100,)),
 )
 
-_PIPELINE_PATHS = {
-    "batched": "zarr.core.codec_pipeline.BatchedCodecPipeline",
-    "fused": "zarr.core.codec_pipeline.FusedCodecPipeline",
+_PIPELINE_SETTINGS = {
+    "batched": {"codec_pipeline.path": "zarr.core.codec_pipeline.BatchedCodecPipeline" },
+    "fused_full_threaded": {"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline","codec_pipeline.max_workers": None},
+    "fused_single_threaded": {"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline","codec_pipeline.max_workers": 1},
 }
 
-_LATENCY_VALUES = (0.0, 0.001, 0.05, 0.2)
+_LATENCY_VALUES = (0, 0.03)
 
 
 @pytest.fixture(params=_LATENCY_VALUES, ids=lambda v: f"latency={v}")
@@ -69,32 +100,35 @@ def bench_store(store: Store, latency: float, request: pytest.FixtureRequest) ->
         if store_kind == "local":
             pytest.skip("latency injection only applies to in-memory store")
         return LatencyStore(
-            store, get_latency=(latency, latency / 4), set_latency=(latency, latency / 4)
+            store, get_latency=(latency, latency * 1.2), set_latency=(latency, latency * 1.2)
         )
+    if store_kind == "memory":
+        pytest.skip("memory store doesn't offer much over local without latency")
     return store
 
 
-@pytest.fixture(params=["batched", "fused"])
+@pytest.fixture(params=["batched", "fused_full_threaded", "fused_single_threaded"])
 def pipeline(request: pytest.FixtureRequest) -> Iterator[str]:
     """Set ``codec_pipeline.path`` for the duration of the benchmark.
 
     Yields the pipeline name so each parametrize cell has a distinct
     benchmark id.
     """
     name = request.param
-    with zarr_config.set({"codec_pipeline.path": _PIPELINE_PATHS[name]}):
+    with zarr_config.set(_PIPELINE_SETTINGS[name]):
         yield name
 
-
-@pytest.mark.parametrize("compression_name", [None, "gzip"])
+@pytest.mark.parametrize("get_data", [lambda shape: 1, lambda shape: _data(shape)], ids=["repeated", "semi_random"])
+@pytest.mark.parametrize("compression_name", ["zstd"])
 @pytest.mark.parametrize("layout", layouts, ids=str)
-@pytest.mark.parametrize("store", ["memory", "local"], indirect=["store"])
+@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
 def test_write_array(
     bench_store: Store,
     layout: Layout,
     compression_name: CompressorName,
     pipeline: str,
     benchmark: BenchmarkFixture,
+    get_data: Callable[[tuple[int]], np.ndarray | int]
 ) -> None:
     """
     Test the time required to fill an array with a single value
@@ -108,19 +142,22 @@ def test_write_array(
         compressors=compressors[compression_name],  # type: ignore[arg-type]
         fill_value=0,
     )
+    def setup():
+        clear_cache()
+        return (arr, Ellipsis, get_data(layout.shape)), {}
+    benchmark.pedantic(setitem, setup=setup, rounds=3)
 
-    benchmark(setitem, arr, Ellipsis, 1)
-
-
-@pytest.mark.parametrize("compression_name", [None, "gzip"])
+@pytest.mark.parametrize("get_data", [lambda shape: 1, lambda shape: _data(shape)], ids=["repeated", "semi_random"])
+@pytest.mark.parametrize("compression_name", ["zstd"])
 @pytest.mark.parametrize("layout", layouts, ids=str)
-@pytest.mark.parametrize("store", ["memory", "local"], indirect=["store"])
+@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
 def test_read_array(
     bench_store: Store,
     layout: Layout,
     compression_name: CompressorName,
     pipeline: str,
     benchmark: BenchmarkFixture,
+    get_data: Callable[[tuple[int]], np.ndarray | int]
 ) -> None:
     """
     Test the time required to fill an array with a single value
@@ -134,5 +171,8 @@ def test_read_array(
         compressors=compressors[compression_name],  # type: ignore[arg-type]
         fill_value=0,
     )
-    arr[:] = 1
-    benchmark(getitem, arr, Ellipsis)
+    arr[:] = get_data(layout.shape)
+    def setup():
+        clear_cache()
+        return (arr, Ellipsis), {}
+    benchmark.pedantic(getitem, setup=setup, rounds=3)

tests/conftest.py

@@ -110,7 +110,7 @@ async def parse_store(
     if store == "zip":
         return await ZipStore.open(f"{path}/zarr.zip", mode="w")
     if store == "memory_get_latency":
-        return LatencyStore(MemoryStore(), get_latency=0.0001, set_latency=0)
+        return LatencyStore(MemoryStore(), get_latency=0.0001, set_latency=0.)
     raise AssertionError