diff --git a/src/zarr/abc/codec.py b/src/zarr/abc/codec.py index eed2119aff..099974583d 100644 --- a/src/zarr/abc/codec.py +++ b/src/zarr/abc/codec.py @@ -9,7 +9,7 @@ from zarr.abc.metadata import Metadata from zarr.core.buffer import Buffer, NDBuffer from zarr.core.common import NamedConfig, concurrent_map -from zarr.core.config import config +from zarr.core.config import get_async_concurrency if TYPE_CHECKING: from collections.abc import Awaitable, Callable, Iterable @@ -249,7 +249,7 @@ async def decode_partial( return await concurrent_map( list(batch_info), self._decode_partial_single, - config.get("async.concurrency"), + get_async_concurrency(), ) @@ -286,7 +286,7 @@ async def encode_partial( await concurrent_map( list(batch_info), self._encode_partial_single, - config.get("async.concurrency"), + get_async_concurrency(), ) @@ -493,7 +493,7 @@ async def _batching_helper[CI: CodecInput, CO: CodecOutput]( return await concurrent_map( list(batch_info), _noop_for_none(func), - config.get("async.concurrency"), + get_async_concurrency(), ) diff --git a/src/zarr/core/codec_pipeline.py b/src/zarr/core/codec_pipeline.py index 5c26681d6b..472b564212 100644 --- a/src/zarr/core/codec_pipeline.py +++ b/src/zarr/core/codec_pipeline.py @@ -17,7 +17,7 @@ SupportsSyncCodec, ) from zarr.core.common import concurrent_map -from zarr.core.config import config +from zarr.core.config import async_concurrency_scope, config, get_async_concurrency from zarr.core.indexing import SelectorTuple, is_scalar from zarr.errors import ZarrUserWarning from zarr.registry import register_pipeline @@ -395,7 +395,7 @@ async def read_batch( for byte_getter, array_spec, *_ in batch_info_list ], lambda byte_getter, prototype: byte_getter.get(prototype), - config.get("async.concurrency"), + get_async_concurrency(), ) chunk_array_batch = await self.decode_batch( [ @@ -505,7 +505,7 @@ async def _read_key( for byte_setter, chunk_spec, chunk_selection, _, is_complete_chunk in batch_info ], _read_key, - config.get("async.concurrency"), + get_async_concurrency(), ) chunk_array_decoded = await self.decode_batch( [ @@ -571,7 +571,7 @@ async def _write_key(byte_setter: ByteSetter, chunk_bytes: Buffer | None) -> Non ) ], _write_key, - config.get("async.concurrency"), + get_async_concurrency(), ) async def decode( @@ -579,8 +579,9 @@ async def decode( chunk_bytes_and_specs: Iterable[tuple[Buffer | None, ArraySpec]], ) -> Iterable[NDBuffer | None]: output: list[NDBuffer | None] = [] - for batch_info in batched(chunk_bytes_and_specs, self.batch_size): - output.extend(await self.decode_batch(batch_info)) + with async_concurrency_scope(): + for batch_info in batched(chunk_bytes_and_specs, self.batch_size): + output.extend(await self.decode_batch(batch_info)) return output async def encode( @@ -588,8 +589,9 @@ async def encode( chunk_arrays_and_specs: Iterable[tuple[NDBuffer | None, ArraySpec]], ) -> Iterable[Buffer | None]: output: list[Buffer | None] = [] - for single_batch_info in batched(chunk_arrays_and_specs, self.batch_size): - output.extend(await self.encode_batch(single_batch_info)) + with async_concurrency_scope(): + for single_batch_info in batched(chunk_arrays_and_specs, self.batch_size): + output.extend(await self.encode_batch(single_batch_info)) return output async def read( @@ -598,14 +600,15 @@ async def read( out: NDBuffer, drop_axes: tuple[int, ...] = (), ) -> tuple[GetResult, ...]: - batch_results = await concurrent_map( - [ - (single_batch_info, out, drop_axes) - for single_batch_info in batched(batch_info, self.batch_size) - ], - self.read_batch, - config.get("async.concurrency"), - ) + with async_concurrency_scope(): + batch_results = await concurrent_map( + [ + (single_batch_info, out, drop_axes) + for single_batch_info in batched(batch_info, self.batch_size) + ], + self.read_batch, + get_async_concurrency(), + ) results: list[GetResult] = [] for batch in batch_results: results.extend(batch) @@ -617,14 +620,15 @@ async def write( value: NDBuffer, drop_axes: tuple[int, ...] = (), ) -> None: - await concurrent_map( - [ - (single_batch_info, value, drop_axes) - for single_batch_info in batched(batch_info, self.batch_size) - ], - self.write_batch, - config.get("async.concurrency"), - ) + with async_concurrency_scope(): + await concurrent_map( + [ + (single_batch_info, value, drop_axes) + for single_batch_info in batched(batch_info, self.batch_size) + ], + self.write_batch, + get_async_concurrency(), + ) def codecs_from_list( diff --git a/src/zarr/core/config.py b/src/zarr/core/config.py index 7dcbc78e31..5b3c6585b0 100644 --- a/src/zarr/core/config.py +++ b/src/zarr/core/config.py @@ -29,11 +29,15 @@ from __future__ import annotations +import contextlib +import contextvars from typing import TYPE_CHECKING, Any, Literal, cast from donfig import Config as DConfig if TYPE_CHECKING: + from collections.abc import Iterator + from donfig.config_obj import ConfigSet @@ -148,6 +152,47 @@ def enable_gpu(self) -> ConfigSet: ) +# Operation-scoped cache for ``async.concurrency``. +# +# The per-chunk codec read/write path looks up ``async.concurrency`` many times +# (once per chunk, plus once per codec per chunk via the batching helper). Each +# lookup walks Donfig's nested config. To avoid that, a single read/write/ +# encode/decode operation captures the value once via ``async_concurrency_scope`` +# and the hot paths read it through ``get_async_concurrency``. Outside of a +# scope we fall back to a live ``config.get``, so configuration changes always +# take effect between operations. +_async_concurrency_var: contextvars.ContextVar[int | None] = contextvars.ContextVar( + "zarr_async_concurrency" +) + + +def get_async_concurrency() -> int | None: + """Return the ``async.concurrency`` limit for the current operation. + + Within an :func:`async_concurrency_scope`, returns the value captured once at + the start of the operation. Outside of a scope, falls back to a live + ``config.get`` lookup. + """ + try: + return _async_concurrency_var.get() + except LookupError: + return cast("int | None", config.get("async.concurrency")) + + +@contextlib.contextmanager +def async_concurrency_scope() -> Iterator[None]: + """Capture ``async.concurrency`` once for the duration of an operation. + + Reads of :func:`get_async_concurrency` inside this context return the + captured value rather than performing a fresh ``config.get`` lookup. + """ + token = _async_concurrency_var.set(cast("int | None", config.get("async.concurrency"))) + try: + yield + finally: + _async_concurrency_var.reset(token) + + def parse_indexing_order(data: Any) -> Literal["C", "F"]: if data in ("C", "F"): return cast("Literal['C', 'F']", data) diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index 52eaa3e144..8aad48a2fb 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -3286,6 +3286,13 @@ async def _iter_members( raise ValueError(f"Unexpected type: {type(fetched_node)}") +class _Drained: + """Sentinel signalling that a subgroup iterator has been fully drained.""" + + +_DRAINED = _Drained() + + async def _iter_members_deep( group: AsyncGroup, *, @@ -3344,10 +3351,48 @@ async def _iter_members_deep( node, max_depth=new_depth, skip_keys=skip_keys, semaphore=semaphore ) - for prefix, subgroup_iter in to_recurse.items(): - async for name, node in subgroup_iter: - key = f"{prefix}/{name}".lstrip("/") - yield key, node + # Recurse into sibling subgroups concurrently rather than draining each + # subgroup's iterator in sequence, so that sibling subtrees overlap their + # latency. The shared ``semaphore`` (threaded into every recursive call) + # still bounds the total number of concurrent ``getitem`` operations. + # Members are not guaranteed to be ordered, so we yield them as they arrive. + queue: asyncio.Queue[tuple[str, AnyAsyncArray | AsyncGroup] | Exception | _Drained] = ( + asyncio.Queue() + ) + + async def _drain( + prefix: str, + subgroup_iter: AsyncGenerator[tuple[str, AnyAsyncArray | AsyncGroup], None], + ) -> None: + try: + async for name, node in subgroup_iter: + key = f"{prefix}/{name}".lstrip("/") + await queue.put((key, node)) + except Exception as exc: + # Forward the error to the consumer so it surfaces from this + # generator just as it would from sequential iteration. + await queue.put(exc) + finally: + await queue.put(_DRAINED) + + tasks = [ + asyncio.create_task(_drain(prefix, subgroup_iter)) + for prefix, subgroup_iter in to_recurse.items() + ] + remaining = len(tasks) + try: + while remaining: + item = await queue.get() + if isinstance(item, _Drained): + remaining -= 1 + elif isinstance(item, Exception): + raise item + else: + yield item + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) async def _read_metadata_v3(store: Store, path: str) -> ArrayV3Metadata | GroupMetadata: