diff --git a/docs/release-notes.md b/docs/release-notes.md index 3fd8a5f360..16f92ea571 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -30,6 +30,7 @@ emits), instead of falling back to the system's native byte order. ([#3417](https://github.com/zarr-developers/zarr-python/issues/3417)) - Fixed `save_array`, `Group.__setitem__`, and `load` for 0-dimensional arrays. ([#3469](https://github.com/zarr-developers/zarr-python/issues/3469)) - Fixed inner-codec spec evolution for sharded arrays. The sharding codec now threads the array spec through its inner codec chain when evolving codecs, so a codec that changes the dtype upstream of `BytesCodec` no longer leaves the inner chain evolved against the wrong spec (which previously failed at decode time). This runs on the default `BatchedCodecPipeline` as well. Standard inner chains (`[BytesCodec]`, `[BytesCodec, ZstdCodec]`, transpose + bytes) are byte-identical to before. Restores the behavior of #2179. ([#3885](https://github.com/zarr-developers/zarr-python/issues/3885)) +- Fixed the opt-in `FusedCodecPipeline` for sharded arrays whose inner or index codec chain contains a codec implementing only the async codec interface (no `SupportsSyncCodec`). Such arrays previously raised `TypeError: All codecs must implement SupportsSyncCodec` on both read and write; the pipeline now declines its synchronous fast path for them and falls back to the async path, matching the behavior of the default `BatchedCodecPipeline`. Fully sync-capable codec chains keep the fast path unchanged. ([#4179](https://github.com/zarr-developers/zarr-python/issues/4179)) - Make chunk normalization properly handle `-1` as a compact representation of the length of an entire axis. Reject several previously-accepted but ill-defined chunk specifications: `chunks=True` (previously silently produced size-1 chunks), diff --git a/src/zarr/abc/codec.py b/src/zarr/abc/codec.py index 61c5dc9948..34d349e6d1 100644 --- a/src/zarr/abc/codec.py +++ b/src/zarr/abc/codec.py @@ -82,6 +82,20 @@ def _decode_sync(self, chunk_data: CO, chunk_spec: ArraySpec) -> CI: ... def _encode_sync(self, chunk_data: CI, chunk_spec: ArraySpec) -> CO | None: ... +def _codec_supports_sync(codec: object) -> bool: + """Whether `codec` can actually run on a synchronous (no event loop) path. + + Structural membership in `SupportsSyncCodec` is necessary but not always + sufficient: a codec can provide `_decode_sync`/`_encode_sync` whose ability + to run depends on runtime configuration the type system cannot see. + `ShardingCodec` is the canonical case — its sync methods delegate to its + configured inner and index codec chains, so they only work when every codec + in those chains is itself sync-capable. Such codecs opt out dynamically via + a `_sync_capable` attribute/property (absent means capable). + """ + return isinstance(codec, SupportsSyncCodec) and getattr(codec, "_sync_capable", True) + + class BaseCodec[CI: CodecInput, CO: CodecOutput](Metadata): """Generic base class for codecs. diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 2d4d63d400..2a95109e6b 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -14,7 +14,7 @@ ArrayBytesCodecPartialEncodeMixin, Codec, CodecPipeline, - SupportsSyncCodec, + _codec_supports_sync, ) from zarr.abc.store import ( ByteGetter, @@ -1406,8 +1406,30 @@ def _is_complete_shard_write( is_complete_chunk for *_, is_complete_chunk in indexed_chunks ) + @property + def _sync_capable(self) -> bool: + """Dynamic opt-out consulted by `_codec_supports_sync` / `ChunkTransform`. + + This codec structurally satisfies `SupportsSyncCodec`, but every sync + method (`_decode_sync`, `_encode_sync`, `_decode_partial_sync`, + `_encode_partial_sync`) delegates to the inner and index codec chains + through `ChunkTransform`, so it can only run synchronously when every + codec in BOTH chains is itself sync-capable. Reporting False here makes + `ChunkTransform` construction raise, which in turn makes + `FusedCodecPipeline.evolve_from_array_spec` set `sync_transform=None` — + the whole pipeline then declines the sync fast path and routes through + the async paths (partial shard decode / async fallback write), exactly + as it does for an async-only TOP-level codec or a non-sync store. + """ + return self._inner_codecs_sync_capable() and self._index_codecs_sync_capable() + + def _inner_codecs_sync_capable(self) -> bool: + # _codec_supports_sync (not bare isinstance) so a nested sharding codec + # with an async-only inner chain propagates its opt-out outward. + return all(_codec_supports_sync(c) for c in self.codecs) + def _index_codecs_sync_capable(self) -> bool: - return all(isinstance(c, SupportsSyncCodec) for c in self.index_codecs) + return all(_codec_supports_sync(c) for c in self.index_codecs) async def _decode_shard_index( self, index_bytes: Buffer, chunks_per_shard: tuple[int, ...] diff --git a/src/zarr/core/chunk_utils.py b/src/zarr/core/chunk_utils.py index ee42e60cce..d93793f853 100644 --- a/src/zarr/core/chunk_utils.py +++ b/src/zarr/core/chunk_utils.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, cast -from zarr.abc.codec import GetResult, SupportsSyncCodec +from zarr.abc.codec import GetResult, SupportsSyncCodec, _codec_supports_sync from zarr.core.indexing import is_scalar if TYPE_CHECKING: @@ -240,7 +240,13 @@ class ChunkTransform: def __post_init__(self) -> None: from zarr.core.codec_pipeline import codecs_from_list - non_sync = [c for c in self.codecs if not isinstance(c, SupportsSyncCodec)] + # _codec_supports_sync, not a bare isinstance check: a codec can satisfy + # the SupportsSyncCodec protocol structurally yet be unable to run + # synchronously (ShardingCodec whose inner/index chain contains an + # async-only codec). Such codecs opt out via `_sync_capable`, and the + # TypeError here is what makes FusedCodecPipeline.evolve_from_array_spec + # decline the sync fast path and fall back to the async pipeline. + non_sync = [c for c in self.codecs if not _codec_supports_sync(c)] if non_sync: names = ", ".join(type(c).__name__ for c in non_sync) raise TypeError( diff --git a/src/zarr/storage/_fsspec.py b/src/zarr/storage/_fsspec.py index 617980ac19..37d134dd95 100644 --- a/src/zarr/storage/_fsspec.py +++ b/src/zarr/storage/_fsspec.py @@ -3,7 +3,6 @@ import json import warnings from contextlib import suppress -from logging import getLogger from typing import TYPE_CHECKING, Any from packaging.version import parse as parse_version @@ -19,8 +18,6 @@ from zarr.errors import ZarrUserWarning from zarr.storage._utils import _dereference_path -logger = getLogger(__name__) - if TYPE_CHECKING: from collections.abc import AsyncIterator, Iterable @@ -38,26 +35,6 @@ ) -async def _close_fs(fs: AsyncFileSystem) -> None: - """ - Best-effort async close of an fsspec async filesystem owned by FsspecStore. - - For filesystems that expose `set_session()` (e.g. s3fs) the underlying - aiohttp `ClientSession` is closed explicitly, which prevents - "Unclosed client session" `ResourceWarning`s from aiohttp. For all - other filesystem types the call is a no-op (not every implementation - manages an HTTP session directly). - - Note that `set_session()` lazily creates a session if none exists yet, so - closing a store that never performed any I/O may instantiate a session - purely to close it. This is accepted best-effort behavior; fsspec does not - expose a stable, cross-implementation way to test for an existing session. - """ - if hasattr(fs, "set_session"): - session = await fs.set_session() - await session.close() - - def _make_async(fs: AbstractFileSystem) -> AsyncFileSystem: """Convert a sync FSSpec filesystem to an async FFSpec filesystem @@ -126,6 +103,15 @@ class FsspecStore(Store): ZarrUserWarning If the file system (fs) was not created with `asynchronous=True`. + Notes + ----- + Closing the store does not close the underlying filesystem or its network + session. fsspec caches and shares filesystem instances across callers, so + the store cannot know whether it is the only user, and closing a shared + session would break other stores. The filesystem's lifecycle belongs to + whoever created it; use fsspec's own tools (e.g. `clear_instance_cache`) + to release it. + See Also -------- FsspecStore.from_upath @@ -152,9 +138,6 @@ def __init__( self.fs = fs self.path = path self.allowed_exceptions = allowed_exceptions - # True only when this store created fs itself (from_url / from_mapper with new instance). - # Callers who supply their own fs remain responsible for its lifecycle. - self._owns_fs: bool = False if not self.fs.async_impl: raise TypeError("Filesystem needs to support async operations.") @@ -220,17 +203,13 @@ def from_mapper( ------- FsspecStore """ - original_fs = fs_map.fs - fs = _make_async(original_fs) - store = cls( + fs = _make_async(fs_map.fs) + return cls( fs=fs, path=fs_map.root, read_only=read_only, allowed_exceptions=allowed_exceptions, ) - # _make_async returns a new instance when converting sync→async; own it. - store._owns_fs = fs is not original_fs - return store @classmethod def from_url( @@ -272,39 +251,16 @@ def from_url( if not fs.async_impl: fs = _make_async(fs) - store = cls(fs=fs, path=path, read_only=read_only, allowed_exceptions=allowed_exceptions) - store._owns_fs = True - return store + return cls(fs=fs, path=path, read_only=read_only, allowed_exceptions=allowed_exceptions) def with_read_only(self, read_only: bool = False) -> FsspecStore: # docstring inherited - new_store = type(self)( + return type(self)( fs=self.fs, path=self.path, allowed_exceptions=self.allowed_exceptions, read_only=read_only, ) - # The derived store shares the same fs. Transfer ownership so the - # surviving store closes it, and clear ours to avoid a double-close. - # Otherwise the common `from_url(...).with_read_only()` chain would - # drop the only owner (the unreferenced source) and leak the session. - new_store._owns_fs = self._owns_fs - self._owns_fs = False - return new_store - - def close(self) -> None: - # docstring inherited - if self._owns_fs: - from zarr.core.sync import sync as zarr_sync - - # Best-effort: a failure to release the session must not block close(), - # but log it so a genuine regression in the close path stays observable - # rather than silently reverting to the leaking behavior. - try: - zarr_sync(_close_fs(self.fs)) - except Exception: - logger.debug("Failed to close owned filesystem %r", self.fs, exc_info=True) - super().close() async def clear(self) -> None: # docstring inherited diff --git a/tests/test_fused_pipeline.py b/tests/test_fused_pipeline.py index 73c2c6e1c3..489ba310c7 100644 --- a/tests/test_fused_pipeline.py +++ b/tests/test_fused_pipeline.py @@ -8,6 +8,7 @@ import pytest import zarr +from zarr.abc.codec import BytesBytesCodec from zarr.codecs.bytes import BytesCodec from zarr.codecs.gzip import GzipCodec from zarr.codecs.transpose import TransposeCodec @@ -582,6 +583,93 @@ def spy_write_sync(self: Any, *args: Any, **kwargs: Any) -> Any: ) +# --------------------------------------------------------------------------- +# Async-only codecs inside a shard's inner codec chain +# --------------------------------------------------------------------------- + + +class _AsyncOnlyNoopCodec(BytesBytesCodec): # type: ignore[misc,unused-ignore] + """A no-op BB codec implementing ONLY the async codec interface. + + Deliberately does NOT satisfy `SupportsSyncCodec` (no `_decode_sync` / + `_encode_sync`), modelling a third-party codec that predates the sync + protocol. Class-level counters prove the codec actually ran. + """ + + is_fixed_size = True + encode_calls = 0 + decode_calls = 0 + + def to_dict(self) -> dict[str, Any]: + return {"name": "test-async-only-noop", "configuration": {}} + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> _AsyncOnlyNoopCodec: + return cls() + + def compute_encoded_size(self, input_byte_length: int, _spec: Any) -> int: + return input_byte_length + + async def _encode_single(self, chunk_bytes: Any, chunk_spec: Any) -> Any: + type(self).encode_calls += 1 + return chunk_bytes + + async def _decode_single(self, chunk_bytes: Any, chunk_spec: Any) -> Any: + type(self).decode_calls += 1 + return chunk_bytes + + +def test_sharded_roundtrip_with_async_only_inner_codec() -> None: + """A sharded array whose INNER codec chain contains an async-only codec + round-trips under FusedCodecPipeline (full write, partial write, full read, + partial read). + + Regression: the pipeline's top-level guard (evolve_from_array_spec -> + sync_transform=None) only inspected the top-level chain. ShardingCodec + structurally satisfies SupportsSyncCodec, so a sync transform was built and + the sync fast path dove into ShardingCodec's sync shard paths, which raised + TypeError from the inner ChunkTransform. The pipeline must instead decline + the sync fast path and fall back to the async inner pipeline, like + BatchedCodecPipeline. + """ + _AsyncOnlyNoopCodec.encode_calls = 0 + _AsyncOnlyNoopCodec.decode_calls = 0 + + with zarr_config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}): + store = MemoryStore() + arr = zarr.create_array( + store=store, + shape=(16, 16), + shards=(8, 8), + chunks=(4, 4), + dtype="int32", + compressors=[_AsyncOnlyNoopCodec()], + fill_value=-1, + ) + assert isinstance(arr._async_array.codec_pipeline, FusedCodecPipeline) + + data = np.arange(256, dtype="int32").reshape(16, 16) + arr[:] = data # full write + np.testing.assert_array_equal(arr[:], data) # full read + np.testing.assert_array_equal(arr[2:11, 3:14], data[2:11, 3:14]) # partial read + + arr[5:7, 5:13] = 0 # partial write (read-merge-write of existing shards) + data[5:7, 5:13] = 0 + np.testing.assert_array_equal(arr[:], data) + + assert _AsyncOnlyNoopCodec.encode_calls > 0, "async-only inner codec never encoded" + assert _AsyncOnlyNoopCodec.decode_calls > 0, "async-only inner codec never decoded" + + # The stored bytes are valid for the default pipeline too: read them back + # under BatchedCodecPipeline (default codec_pipeline.path). Opening from + # metadata needs the codec name in the registry. + from zarr.registry import register_codec + + register_codec("test-async-only-noop", _AsyncOnlyNoopCodec) + reread = zarr.open_array(store=store, mode="r") + np.testing.assert_array_equal(reread[:], data) + + # --------------------------------------------------------------------------- # AsyncChunkTransform: the async per-chunk codec chain used on the async # fallback path. It is the async mirror of ChunkTransform, so it must produce diff --git a/tests/test_store/test_fsspec.py b/tests/test_store/test_fsspec.py index 898d49ec08..515e1526b6 100644 --- a/tests/test_store/test_fsspec.py +++ b/tests/test_store/test_fsspec.py @@ -276,75 +276,20 @@ async def test_delete_dir_unsupported_deletes(self, store: FsspecStore) -> None: ): await store.delete_dir("test_prefix") - # ── Filesystem lifecycle (ownership) ────────────────────────────────────── + # ── Filesystem lifecycle ────────────────────────────────────────────────── - def test_from_url_owns_filesystem(self, endpoint_url: str) -> None: - """FsspecStore.from_url() creates the async fs; it must own it.""" + async def test_close_marks_store_closed(self, endpoint_url: str) -> None: + """close() must succeed and mark the store not-open.""" store = FsspecStore.from_url( f"s3://{test_bucket_name}/lifecycle/", storage_options={"endpoint_url": endpoint_url, "anon": False}, ) - assert store._owns_fs - store.close() - - async def test_from_url_close_releases_store(self, endpoint_url: str) -> None: - """ - close() on a from_url() store must succeed without error and mark the - store as closed. For the owned filesystem, _close_fs() is invoked to - release the underlying S3 client / aiohttp connection pool. - """ - store = FsspecStore.from_url( - f"s3://{test_bucket_name}/lifecycle/", - storage_options={"endpoint_url": endpoint_url, "anon": False}, - ) - # Materialise the S3 client and connection pool. await store.set("probe", cpu.Buffer.from_bytes(b"x")) store.close() assert not store._is_open - def test_direct_construction_does_not_own_filesystem(self, endpoint_url: str) -> None: - """Direct FsspecStore() must not claim ownership — the caller owns the fs.""" - try: - from fsspec import url_to_fs - except ImportError: - from fsspec.core import url_to_fs - fs, path = url_to_fs( - f"s3://{test_bucket_name}", endpoint_url=endpoint_url, anon=False, asynchronous=True - ) - store = FsspecStore(fs=fs, path=path) - assert not store._owns_fs - - @pytest.mark.skipif( - parse_version(fsspec.__version__) < parse_version("2024.03.01"), - reason="Prior bug in from_upath", - ) - def test_from_upath_does_not_own_filesystem(self, endpoint_url: str) -> None: - """from_upath() uses the UPath's existing fs; the store must not own it.""" - upath = pytest.importorskip("upath") - path = upath.UPath( - f"s3://{test_bucket_name}/foo/bar/", - endpoint_url=endpoint_url, - anon=False, - asynchronous=True, - ) - store = FsspecStore.from_upath(path) - assert not store._owns_fs - - def test_from_mapper_does_not_own_already_async_filesystem(self, endpoint_url: str) -> None: - """from_mapper() with an already-async fs must not claim ownership.""" - s3_filesystem = s3fs.S3FileSystem( - asynchronous=True, - endpoint_url=endpoint_url, - anon=False, - skip_instance_cache=True, - ) - mapper = s3_filesystem.get_mapper(f"s3://{test_bucket_name}/") - store = FsspecStore.from_mapper(mapper) - # _make_async returns the same instance for an already-async fs. - assert not store._owns_fs - def array_roundtrip(store: FsspecStore) -> None: """ @@ -574,47 +519,47 @@ def test_open_s3map_raises(endpoint_url: str) -> None: zarr.open(store=mapper, storage_options={"anon": True}, mode="w", shape=(3, 3)) -async def test_close_fs_closes_s3_client() -> None: - """ - _close_fs() must call set_session() and then close() on the returned - S3 client. This is verified with mocks to avoid a real S3 connection. - """ - from unittest.mock import AsyncMock +async def test_close_does_not_close_filesystem_session() -> None: + """close() must not touch the filesystem's session. - from zarr.storage._fsspec import _close_fs + fsspec caches and shares filesystem instances across callers, so the + session is not the store's to close. HTTP is used because its aiohttp + session is observably closed for good; s3fs transparently reconnects, which + would hide a regression. No request is issued — set_session() only + constructs the session. + """ + pytest.importorskip("aiohttp") + store = FsspecStore.from_url("http://example.com/a") + session = await store.fs.set_session() - mock_client = AsyncMock() - mock_fs = AsyncMock() - mock_fs.set_session = AsyncMock(return_value=mock_client) + store.close() - await _close_fs(mock_fs) + assert not session.closed - mock_fs.set_session.assert_called_once() - mock_client.close.assert_called_once() +async def test_close_does_not_break_a_sibling_store() -> None: + """Closing one store must not close a session another store is using. -async def test_close_fs_no_op_for_fs_without_set_session() -> None: - """_close_fs() must be a no-op for filesystems that don't expose set_session().""" - from unittest.mock import AsyncMock + Two stores from different URLs on one host are handed the same cached + filesystem; a store that closed it on close() would take the sibling's + session down too. This is the regression guard for that bug. + """ + pytest.importorskip("aiohttp") + s1 = FsspecStore.from_url("http://example.com/a") + s2 = FsspecStore.from_url("http://example.com/b") + session = await s2.fs.set_session() - from zarr.storage._fsspec import _close_fs + s1.close() - mock_fs = AsyncMock(spec=[]) # empty spec — no set_session attribute - await _close_fs(mock_fs) # must not raise + assert not session.closed @pytest.mark.skipif( parse_version(fsspec.__version__) < parse_version("2024.12.0"), reason="No AsyncFileSystemWrapper", ) -def test_from_mapper_owns_wrapped_sync_filesystem(tmp_path: pathlib.Path) -> None: - """ - from_mapper() with a sync fs must wrap it in AsyncFileSystemWrapper and - claim ownership so that close() cleans it up. - - The local filesystem is synchronous; _make_async() produces a new - AsyncFileSystemWrapper instance — a different object from the original fs. - """ +def test_from_mapper_wraps_sync_filesystem(tmp_path: pathlib.Path) -> None: + """from_mapper() with a sync fs wraps it in an AsyncFileSystemWrapper.""" import fsspec as _fsspec from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper @@ -622,32 +567,21 @@ def test_from_mapper_owns_wrapped_sync_filesystem(tmp_path: pathlib.Path) -> Non mapper = fs.get_mapper(str(tmp_path)) store = FsspecStore.from_mapper(mapper) assert isinstance(store.fs, AsyncFileSystemWrapper) - assert store._owns_fs @pytest.mark.skipif( parse_version(fsspec.__version__) < parse_version("2024.12.0"), reason="No AsyncFileSystemWrapper", ) -def test_with_read_only_transfers_filesystem_ownership(tmp_path: pathlib.Path) -> None: - """ - with_read_only() must transfer fs ownership to the derived store and clear - it on the source, so the surviving store closes the shared fs exactly once. - - In the common ``from_url(...).with_read_only()`` chain the source store is - immediately unreferenced; if ownership were not transferred, the only owner - would be garbage-collected without close() and the session would leak. - """ +def test_with_read_only_shares_filesystem(tmp_path: pathlib.Path) -> None: + """with_read_only() returns a store sharing the source's filesystem.""" source = FsspecStore.from_url(f"file://{tmp_path}", storage_options={"auto_mkdir": False}) - assert source._owns_fs derived = source.with_read_only(read_only=True) - # Ownership moved to the survivor; the source no longer owns it (no double-close). - assert derived._owns_fs - assert not source._owns_fs - # The derived store shares the same underlying fs. assert derived.fs is source.fs + assert derived.read_only + assert not source.read_only @pytest.mark.parametrize("asynchronous", [True, False])