From e2928f9b8d3f822246879bbb671e36c2d2fdd9b4 Mon Sep 17 00:00:00 2001 From: Ilan Gold Date: Wed, 24 Jun 2026 15:26:23 +0000 Subject: [PATCH] perf: benchmarking non-repeated values --- src/zarr/testing/store.py | 4 +- tests/benchmarks/test_e2e.py | 86 ++++++++++++++++++++++++++---------- tests/conftest.py | 2 +- 3 files changed, 66 insertions(+), 26 deletions(-) diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index a0be9601ea..b3cd9985cb 100644 --- a/src/zarr/testing/store.py +++ b/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 diff --git a/tests/benchmarks/test_e2e.py b/tests/benchmarks/test_e2e.py index 377924135f..4ad9341c6f 100644 --- a/tests/benchmarks/test_e2e.py +++ b/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,12 +100,14 @@ 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. @@ -82,19 +115,20 @@ def pipeline(request: pytest.FixtureRequest) -> Iterator[str]: 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) diff --git a/tests/conftest.py b/tests/conftest.py index 031c4a4283..205371b78c 100644 --- a/tests/conftest.py +++ b/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