|
4 | 4 | import threading |
5 | 5 | from concurrent.futures import ThreadPoolExecutor |
6 | 6 | from dataclasses import dataclass, field |
7 | | -from itertools import batched, pairwise |
| 7 | +from itertools import batched, chain, pairwise |
8 | 8 | from typing import TYPE_CHECKING, Any, cast |
9 | 9 | from warnings import warn |
10 | 10 |
|
@@ -54,10 +54,23 @@ def _resolve_max_workers() -> int: |
54 | 54 | """Helper for getting the maximum number of workers available to the `FusedCodecPipeline`""" |
55 | 55 | import os as _os |
56 | 56 |
|
| 57 | + default = _os.cpu_count() or 1 |
57 | 58 | cfg = config.get("codec_pipeline.max_workers", default=None) |
58 | 59 | 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 |
61 | 74 |
|
62 | 75 |
|
63 | 76 | def _get_pool(max_workers: int) -> ThreadPoolExecutor: |
@@ -170,6 +183,23 @@ def pipeline_supports_partial_encode( |
170 | 183 | return isinstance(array_bytes_codec, ArrayBytesCodecPartialEncodeMixin) |
171 | 184 |
|
172 | 185 |
|
| 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 | + |
173 | 203 | async def _fetch_and_decode_as_completed( |
174 | 204 | batch: Sequence[tuple[ByteGetter | None, ArraySpec]], |
175 | 205 | transform: ChunkTransform, |
@@ -202,20 +232,29 @@ def _decode(buffer: Buffer | None, chunk_spec: ArraySpec) -> NDBuffer | None: |
202 | 232 | _fetch, |
203 | 233 | config.get("async.concurrency"), |
204 | 234 | ) |
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)) |
217 | 248 |
|
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)) |
219 | 258 |
|
220 | 259 |
|
221 | 260 | async def _encode_and_write_as_completed( |
@@ -264,10 +303,20 @@ async def _write(idx: int, chunk_bytes: Buffer | None) -> None: |
264 | 303 | # Kick off each chunk's write the instant its encode lands, so writes of |
265 | 304 | # already-compressed chunks proceed while the rest are still encoding. |
266 | 305 | 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)) |
271 | 320 |
|
272 | 321 |
|
273 | 322 | async def _async_read_fallback( |
@@ -469,7 +518,11 @@ class AsyncChunkTransform: |
469 | 518 | _bb_codecs: tuple[BytesBytesCodec, ...] = field(init=False, repr=False, compare=False) |
470 | 519 |
|
471 | 520 | 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)) |
473 | 526 | self._aa_codecs = aa |
474 | 527 | self._ab_codec = ab |
475 | 528 | self._bb_codecs = bb |
@@ -533,7 +586,19 @@ class BatchedCodecPipeline(CodecPipeline): |
533 | 586 | batch_size: int |
534 | 587 |
|
535 | 588 | 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 | + ) |
537 | 602 |
|
538 | 603 | @classmethod |
539 | 604 | def from_codecs(cls, codecs: Iterable[Codec], *, batch_size: int | None = None) -> Self: |
@@ -795,21 +860,44 @@ async def write( |
795 | 860 | def codecs_from_list( |
796 | 861 | codecs: Iterable[Codec], |
797 | 862 | ) -> 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 | + """ |
798 | 873 | from zarr.codecs.sharding import ShardingCodec |
799 | 874 |
|
800 | 875 | codecs = tuple(codecs) # materialize to avoid generator consumption issues |
801 | 876 |
|
802 | | - array_array: tuple[ArrayArrayCodec, ...] = () |
803 | | - array_bytes_maybe: ArrayBytesCodec | None = None |
804 | | - bytes_bytes: tuple[BytesBytesCodec, ...] = () |
805 | | - |
806 | 877 | if any(isinstance(codec, ShardingCodec) for codec in codecs) and len(codecs) > 1: |
807 | 878 | warn( |
808 | 879 | "Combining a `sharding_indexed` codec disables partial reads and " |
809 | 880 | "writes, which may lead to inefficient performance.", |
810 | 881 | category=ZarrUserWarning, |
811 | 882 | stacklevel=3, |
812 | 883 | ) |
| 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, ...] = () |
813 | 901 |
|
814 | 902 | for prev_codec, cur_codec in pairwise((None, *codecs)): |
815 | 903 | if isinstance(cur_codec, ArrayArrayCodec): |
@@ -912,8 +1000,11 @@ def from_codecs(cls, codecs: Iterable[Codec], *, batch_size: int | None = None) |
912 | 1000 | ) |
913 | 1001 |
|
914 | 1002 | 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. |
915 | 1006 | 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) |
917 | 1008 |
|
918 | 1009 | try: |
919 | 1010 | sync_transform: ChunkTransform | None = ChunkTransform(codecs=evolved_codecs) |
|
0 commit comments