Skip to content

Commit 6d30b53

Browse files
committed
fix: minor correctness and hygiene fixes from the sync-pipeline audit
- Codec construction warnings (e.g. sharding's "disables partial reads") fired twice per array open, and on every decode/encode through the fused pipeline's async fallback. Re-constructions of an already-validated codec chain now go through codecs_from_list_unchecked, which validates structure without repeating first-construction advisory warnings; each warning fires exactly once per open under both pipelines. - concurrent_iter returned a lazy generator while its docstring promised eagerly scheduled tasks; it now materializes the task list so awaiting one at a time cannot serialize the batch. - A garbage codec_pipeline.max_workers value (e.g. from the environment) raised ValueError mid-read; it now warns and falls back to the default, consistent with tolerant handling of config input. - The as-completed pipeline helpers abandoned in-flight tasks when one failed, leaving stray background writes and "Task exception was never retrieved" warnings; failures now cancel and drain outstanding tasks. - Benchmarks: seed the data generator for reproducibility; fix a copy-pasted docstring. - Remove dead commented-out test blocks referencing the removed set_range API. Assisted-by: ClaudeCode:claude-sonnet-5
1 parent 123268a commit 6d30b53

11 files changed

Lines changed: 329 additions & 149 deletions

File tree

changes/4205.bugfix.md

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
Fixed several small correctness issues from the codec-pipeline performance work: construction-time
2+
codec warnings (e.g. sharding's "disables partial reads" warning) no longer fire twice per array
3+
open — including for `FusedCodecPipeline`, which previously re-warned via its own codec-chain
4+
reconstruction and, on the async fallback path, on every decode/encode call; `concurrent_iter` now
5+
schedules its tasks eagerly, matching its documented contract; an invalid
6+
`codec_pipeline.max_workers` config/environment value now warns and falls back to the default
7+
instead of raising mid-read; and `FusedCodecPipeline`'s async fallback helpers now cancel
8+
already-spawned fetch/decode/write tasks instead of abandoning them in the background when one
9+
fails.

src/zarr/core/array.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,12 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None
228228
pass
229229

230230
if isinstance(metadata, ArrayV3Metadata):
231+
# The pipeline built here is a throwaway: `evolve_from_array_spec` below
232+
# reconstructs codecs against the evolved spec. `from_codecs` is the
233+
# chain's first construction, so its advisory warnings (e.g. sharding's
234+
# "disables partial reads" warning) fire here; `evolve_from_array_spec`
235+
# re-splits the same already-warned-about chain via
236+
# `codecs_from_list_unchecked`, so it does not re-emit them.
231237
pipeline = get_pipeline_class().from_codecs(metadata.codecs)
232238
from zarr.core.metadata.v3 import RegularChunkGridMetadata
233239

src/zarr/core/chunk_utils.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -238,7 +238,7 @@ class ChunkTransform:
238238
)
239239

240240
def __post_init__(self) -> None:
241-
from zarr.core.codec_pipeline import codecs_from_list
241+
from zarr.core.codec_pipeline import codecs_from_list_unchecked
242242

243243
# _codec_supports_sync, not a bare isinstance check: a codec can satisfy
244244
# the SupportsSyncCodec protocol structurally yet be unable to run
@@ -253,7 +253,11 @@ def __post_init__(self) -> None:
253253
f"All codecs must implement SupportsSyncCodec. The following do not: {names}"
254254
)
255255

256-
aa, ab, bb = codecs_from_list(list(self.codecs))
256+
# `ChunkTransform` is built from a codec chain that already went
257+
# through `codecs_from_list` when the owning pipeline was constructed
258+
# (see `FusedCodecPipeline.evolve_from_array_spec`), so re-splitting it
259+
# here must not re-emit that chain's advisory warnings.
260+
aa, ab, bb = codecs_from_list_unchecked(list(self.codecs))
257261
# SupportsSyncCodec was verified above; the cast is purely for mypy.
258262
self._aa_codecs = cast("tuple[SupportsSyncCodec[NDBuffer, NDBuffer], ...]", tuple(aa))
259263
self._ab_codec = cast("SupportsSyncCodec[NDBuffer, Buffer]", ab)

src/zarr/core/codec_pipeline.py

Lines changed: 118 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import threading
55
from concurrent.futures import ThreadPoolExecutor
66
from dataclasses import dataclass, field
7-
from itertools import batched, pairwise
7+
from itertools import batched, chain, pairwise
88
from typing import TYPE_CHECKING, Any, cast
99
from warnings import warn
1010

@@ -54,10 +54,23 @@ def _resolve_max_workers() -> int:
5454
"""Helper for getting the maximum number of workers available to the `FusedCodecPipeline`"""
5555
import os as _os
5656

57+
default = _os.cpu_count() or 1
5758
cfg = config.get("codec_pipeline.max_workers", default=None)
5859
if cfg is None:
59-
return _os.cpu_count() or 1
60-
return max(1, int(cfg))
60+
return default
61+
try:
62+
return max(1, int(cfg))
63+
except (TypeError, ValueError):
64+
# This value arrives via the config/env layer (e.g.
65+
# `ZARR_CODEC_PIPELINE__MAX_WORKERS`), so tolerate bad input here
66+
# instead of raising mid-read.
67+
warn(
68+
f"Ignoring invalid `codec_pipeline.max_workers` config value {cfg!r}; "
69+
f"falling back to {default}.",
70+
category=ZarrUserWarning,
71+
stacklevel=2,
72+
)
73+
return default
6174

6275

6376
def _get_pool(max_workers: int) -> ThreadPoolExecutor:
@@ -170,6 +183,23 @@ def pipeline_supports_partial_encode(
170183
return isinstance(array_bytes_codec, ArrayBytesCodecPartialEncodeMixin)
171184

172185

186+
async def _cancel_and_drain(futures: Iterable[asyncio.Future[Any]]) -> None:
187+
"""Cancel every not-yet-done future/task and await its outcome.
188+
189+
Used to clean up work spawned by a drain loop (`asyncio.as_completed` +
190+
`await`) when the loop exits early via exception. Without this, tasks
191+
already spawned keep running unattended after the caller has moved on,
192+
and an eventual failure surfaces as an unraisable "exception was never
193+
retrieved" warning instead of being observed here.
194+
"""
195+
pending = [f for f in futures if not f.done()]
196+
if len(pending) == 0:
197+
return
198+
for f in pending:
199+
f.cancel()
200+
await asyncio.gather(*pending, return_exceptions=True)
201+
202+
173203
async def _fetch_and_decode_as_completed(
174204
batch: Sequence[tuple[ByteGetter | None, ArraySpec]],
175205
transform: ChunkTransform,
@@ -202,20 +232,29 @@ def _decode(buffer: Buffer | None, chunk_spec: ArraySpec) -> NDBuffer | None:
202232
_fetch,
203233
config.get("async.concurrency"),
204234
)
205-
for fetch_coro in asyncio.as_completed(fetch_tasks):
206-
idx, buffer = await fetch_coro
207-
chunk_spec = batch[idx][1]
208-
# Bridge both paths to asyncio.Future so the final collection loop
209-
# can `await` uniformly without blocking the event loop. For the
210-
# pool path that means `wrap_future` (not `pool.submit(...).result()`,
211-
# which would block the loop thread for the duration of every decode
212-
# — freezing any unrelated coroutines sharing this loop).
213-
if pool is None:
214-
decode_futures[idx].set_result(_decode(buffer, chunk_spec))
215-
else:
216-
decode_futures[idx] = asyncio.wrap_future(pool.submit(_decode, buffer, chunk_spec))
235+
try:
236+
for fetch_coro in asyncio.as_completed(fetch_tasks):
237+
idx, buffer = await fetch_coro
238+
chunk_spec = batch[idx][1]
239+
# Bridge both paths to asyncio.Future so the final collection loop
240+
# can `await` uniformly without blocking the event loop. For the
241+
# pool path that means `wrap_future` (not `pool.submit(...).result()`,
242+
# which would block the loop thread for the duration of every decode
243+
# — freezing any unrelated coroutines sharing this loop).
244+
if pool is None:
245+
decode_futures[idx].set_result(_decode(buffer, chunk_spec))
246+
else:
247+
decode_futures[idx] = asyncio.wrap_future(pool.submit(_decode, buffer, chunk_spec))
217248

218-
return await asyncio.gather(*decode_futures)
249+
return await asyncio.gather(*decode_futures)
250+
finally:
251+
# On the happy path every future here is already done, so this is a
252+
# no-op; on failure it stops abandoned fetches/decodes from
253+
# continuing to run unattended after this function has raised. A
254+
# single call over both iterables (not two sequential calls) so that
255+
# outer-task cancellation during the first drain can't skip the
256+
# second, leaving its futures/tasks unobserved.
257+
await _cancel_and_drain(chain(fetch_tasks, decode_futures))
219258

220259

221260
async def _encode_and_write_as_completed(
@@ -264,10 +303,20 @@ async def _write(idx: int, chunk_bytes: Buffer | None) -> None:
264303
# Kick off each chunk's write the instant its encode lands, so writes of
265304
# already-compressed chunks proceed while the rest are still encoding.
266305
write_tasks: list[asyncio.Task[None]] = []
267-
for encode_coro in asyncio.as_completed(encode_futures):
268-
idx, chunk_bytes = await encode_coro
269-
write_tasks.append(asyncio.ensure_future(_write(idx, chunk_bytes)))
270-
await asyncio.gather(*write_tasks)
306+
try:
307+
for encode_coro in asyncio.as_completed(encode_futures):
308+
idx, chunk_bytes = await encode_coro
309+
write_tasks.append(asyncio.ensure_future(_write(idx, chunk_bytes)))
310+
await asyncio.gather(*write_tasks)
311+
finally:
312+
# On the happy path every future here is already done, so this is a
313+
# no-op; on failure (an encode or a write raising) it stops
314+
# already-spawned writes from continuing in the background after
315+
# this function has raised. A single call over both iterables (not
316+
# two sequential calls) so that outer-task cancellation during the
317+
# first drain can't skip the second, leaving its futures/tasks
318+
# unobserved.
319+
await _cancel_and_drain(chain(write_tasks, encode_futures))
271320

272321

273322
async def _async_read_fallback(
@@ -469,7 +518,11 @@ class AsyncChunkTransform:
469518
_bb_codecs: tuple[BytesBytesCodec, ...] = field(init=False, repr=False, compare=False)
470519

471520
def __post_init__(self) -> None:
472-
aa, ab, bb = codecs_from_list(list(self.codecs))
521+
# `AsyncChunkTransform` is (re)constructed per decode/encode call from a
522+
# codec chain that already went through `codecs_from_list` when the
523+
# pipeline itself was built, so re-splitting it here must not re-emit
524+
# that chain's advisory warnings on every call.
525+
aa, ab, bb = codecs_from_list_unchecked(list(self.codecs))
473526
self._aa_codecs = aa
474527
self._ab_codec = ab
475528
self._bb_codecs = bb
@@ -533,7 +586,19 @@ class BatchedCodecPipeline(CodecPipeline):
533586
batch_size: int
534587

535588
def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
536-
return type(self).from_codecs(evolve_codecs(self, array_spec))
589+
# Re-splits an already-`codecs_from_list`-validated (and warned-about)
590+
# chain against the evolved spec, so this uses the quiet variant rather
591+
# than routing through `from_codecs` (which would re-warn).
592+
evolved_codecs = evolve_codecs(self, array_spec)
593+
array_array_codecs, array_bytes_codec, bytes_bytes_codecs = codecs_from_list_unchecked(
594+
evolved_codecs
595+
)
596+
return type(self)(
597+
array_array_codecs=array_array_codecs,
598+
array_bytes_codec=array_bytes_codec,
599+
bytes_bytes_codecs=bytes_bytes_codecs,
600+
batch_size=self.batch_size,
601+
)
537602

538603
@classmethod
539604
def from_codecs(cls, codecs: Iterable[Codec], *, batch_size: int | None = None) -> Self:
@@ -795,21 +860,44 @@ async def write(
795860
def codecs_from_list(
796861
codecs: Iterable[Codec],
797862
) -> tuple[tuple[ArrayArrayCodec, ...], ArrayBytesCodec, tuple[BytesBytesCodec, ...]]:
863+
"""Split `codecs` into `(array_array, array_bytes, bytes_bytes)`, validating order.
864+
865+
Emits user-facing advisory warnings about the codec chain (e.g. sharding's
866+
"disables partial reads" warning). Use this for the FIRST construction of a
867+
codec chain from user-supplied codecs. Use `codecs_from_list_unchecked` when
868+
re-splitting a chain that was already validated and warned about by a prior
869+
`codecs_from_list` call (e.g. `evolve_from_array_spec` re-splitting the same
870+
codecs against an evolved spec) — re-warning there would fire the same
871+
advisory once per reconstruction instead of once per user-facing chain.
872+
"""
798873
from zarr.codecs.sharding import ShardingCodec
799874

800875
codecs = tuple(codecs) # materialize to avoid generator consumption issues
801876

802-
array_array: tuple[ArrayArrayCodec, ...] = ()
803-
array_bytes_maybe: ArrayBytesCodec | None = None
804-
bytes_bytes: tuple[BytesBytesCodec, ...] = ()
805-
806877
if any(isinstance(codec, ShardingCodec) for codec in codecs) and len(codecs) > 1:
807878
warn(
808879
"Combining a `sharding_indexed` codec disables partial reads and "
809880
"writes, which may lead to inefficient performance.",
810881
category=ZarrUserWarning,
811882
stacklevel=3,
812883
)
884+
return codecs_from_list_unchecked(codecs)
885+
886+
887+
def codecs_from_list_unchecked(
888+
codecs: Iterable[Codec],
889+
) -> tuple[tuple[ArrayArrayCodec, ...], ArrayBytesCodec, tuple[BytesBytesCodec, ...]]:
890+
"""Split `codecs` into `(array_array, array_bytes, bytes_bytes)`, validating order.
891+
892+
Same structural validation as `codecs_from_list` (raises on bad codec
893+
ordering or a missing/duplicate array-bytes codec) but does NOT emit
894+
user-facing advisory warnings. See `codecs_from_list` for when to use each.
895+
"""
896+
codecs = tuple(codecs) # materialize to avoid generator consumption issues
897+
898+
array_array: tuple[ArrayArrayCodec, ...] = ()
899+
array_bytes_maybe: ArrayBytesCodec | None = None
900+
bytes_bytes: tuple[BytesBytesCodec, ...] = ()
813901

814902
for prev_codec, cur_codec in pairwise((None, *codecs)):
815903
if isinstance(cur_codec, ArrayArrayCodec):
@@ -912,8 +1000,11 @@ def from_codecs(cls, codecs: Iterable[Codec], *, batch_size: int | None = None)
9121000
)
9131001

9141002
def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
1003+
# Re-splits an already-`codecs_from_list`-validated (and warned-about)
1004+
# chain against the evolved spec, so this uses the quiet variant to
1005+
# avoid re-emitting the same advisory warning on every array open.
9151006
evolved_codecs = evolve_codecs(self.codecs, array_spec)
916-
aa, ab, bb = codecs_from_list(evolved_codecs)
1007+
aa, ab, bb = codecs_from_list_unchecked(evolved_codecs)
9171008

9181009
try:
9191010
sync_transform: ChunkTransform | None = ChunkTransform(codecs=evolved_codecs)

src/zarr/core/common.py

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -93,34 +93,34 @@ def concurrent_iter[T: tuple[Any, ...], V](
9393
items: Iterable[T],
9494
func: Callable[..., Awaitable[V]],
9595
limit: int | None = None,
96-
) -> Iterator[asyncio.Task[V]]:
96+
) -> list[asyncio.Task[V]]:
9797
"""Launch `func(*item)` for each item concurrently, returning the tasks.
9898
9999
When `limit` is set, no more than `limit` calls are in flight at once.
100100
Tasks are returned in input order; callers that want completion order
101101
should wrap the result in `asyncio.as_completed`.
102102
103-
Note on `ensure_future`: when the result is passed to `asyncio.gather` or
104-
`asyncio.as_completed`, those already wrap awaitables into tasks, so the
105-
`ensure_future` here is redundant. It matters for callers that iterate and
106-
await tasks one at a time — without eager scheduling, each coroutine would
107-
only start when individually awaited, serializing the work and defeating
108-
the semaphore. It also makes the return type honest (real `Task`s support
109-
`.cancel()`, `.done()`, callbacks) rather than bare coroutines.
103+
Every task is scheduled (via `ensure_future`) before this function
104+
returns, not on first iteration of the result. That matters for callers
105+
that await the returned tasks one at a time — without eager scheduling,
106+
each coroutine would only start when individually awaited, serializing
107+
the work and defeating the semaphore. It also makes the return type
108+
honest (real `Task`s support `.cancel()`, `.done()`, callbacks) rather
109+
than bare coroutines.
110110
111111
See https://docs.python.org/3/library/asyncio-task.html#coroutines:
112112
"Note that simply calling a coroutine will not schedule it to be executed:"
113113
"""
114114
if limit is None:
115-
return (asyncio.ensure_future(func(*item)) for item in items)
115+
return [asyncio.ensure_future(func(*item)) for item in items]
116116

117117
sem = asyncio.Semaphore(limit)
118118

119119
async def run(item: T) -> V:
120120
async with sem:
121121
return await func(*item)
122122

123-
return (asyncio.ensure_future(run(item)) for item in items)
123+
return [asyncio.ensure_future(run(item)) for item in items]
124124

125125

126126
async def concurrent_map[T: tuple[Any, ...], V](

tests/benchmarks/test_e2e.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,8 @@ def _data(shape: tuple[int]) -> np.ndarray:
6363
noise_level = 1
6464
pattern = (np.sin(np.linspace(0, 2 * np.pi, period)) * 50 + 128).round().astype(np.uint8)
6565
data = np.tile(pattern, int(np.ceil(n / period)))[:n].astype(np.int16)
66-
data += np.random.randint(-noise_level, noise_level + 1, size=n, dtype=np.int16)
66+
rng = np.random.default_rng(0)
67+
data += rng.integers(-noise_level, noise_level + 1, size=n, dtype=np.int16)
6768
return np.clip(data, 0, 255).astype(np.uint8)
6869

6970

@@ -189,7 +190,7 @@ def test_read_array(
189190
get_data: Callable[[tuple[int]], np.ndarray | int],
190191
) -> None:
191192
"""
192-
Test the time required to fill an array with a single value
193+
Test the time required to read the entirety of an array
193194
"""
194195
arr = create_array(
195196
bench_store,

tests/test_codecs/test_codecs.py

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from __future__ import annotations
22

33
import json
4+
import warnings
45
from dataclasses import dataclass
56
from typing import TYPE_CHECKING, Any
67

@@ -22,7 +23,7 @@
2223
from zarr.core.metadata.v3 import ArrayV3Metadata
2324
from zarr.dtype import UInt8
2425
from zarr.errors import ZarrUserWarning
25-
from zarr.storage import StorePath
26+
from zarr.storage import MemoryStore, StorePath
2627

2728
if TYPE_CHECKING:
2829
from zarr.abc.codec import Codec
@@ -375,6 +376,49 @@ def test_invalid_metadata_create_array() -> None:
375376
)
376377

377378

379+
@pytest.mark.parametrize(
380+
"pipeline_path",
381+
[
382+
"zarr.core.codec_pipeline.BatchedCodecPipeline",
383+
"zarr.core.codec_pipeline.FusedCodecPipeline",
384+
],
385+
)
386+
def test_sharding_warning_fires_once_per_open(pipeline_path: str) -> None:
387+
"""Construction-time codec warnings (e.g. sharding's partial-reads warning)
388+
must fire exactly once per array open, not once per internal codec-chain
389+
reconstruction.
390+
391+
`create_codec_pipeline` builds a throwaway pipeline via `from_codecs` (which
392+
warns) and then calls `evolve_from_array_spec` on it, which re-splits the
393+
(already-warned-about) codec chain against the evolved spec. That re-split
394+
goes through `codecs_from_list_unchecked` rather than `codecs_from_list`, so
395+
it does not re-emit the warning. `FusedCodecPipeline` additionally builds a
396+
`ChunkTransform` (and, on the async fallback path, an `AsyncChunkTransform`
397+
per call) from the same evolved codec chain, which must use the same quiet
398+
variant.
399+
"""
400+
with config.set({"codec_pipeline.path": pipeline_path}):
401+
store = MemoryStore()
402+
with warnings.catch_warnings():
403+
warnings.simplefilter("ignore")
404+
zarr.create_array(
405+
store,
406+
shape=(16, 16),
407+
chunks=(16, 16),
408+
dtype=np.dtype("uint8"),
409+
fill_value=0,
410+
serializer=ShardingCodec(chunk_shape=(8, 8)),
411+
compressors=[GzipCodec()],
412+
)
413+
414+
with warnings.catch_warnings(record=True) as caught:
415+
warnings.simplefilter("always")
416+
zarr.open_array(store, mode="r")
417+
418+
matches = [w for w in caught if "disables partial reads" in str(w.message)]
419+
assert len(matches) == 1
420+
421+
378422
@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
379423
async def test_resize(store: Store) -> None:
380424
data = np.zeros((16, 18), dtype="uint16")

0 commit comments

Comments
 (0)