diff --git a/packages/zarr-indexing/changes/4348.bugfix.md b/packages/zarr-indexing/changes/4348.bugfix.md new file mode 100644 index 0000000000..69a7daf9bc --- /dev/null +++ b/packages/zarr-indexing/changes/4348.bugfix.md @@ -0,0 +1 @@ +Delegate LazyArray source tokenization to Dask, honoring its registered normalizers, source hooks, and deterministic-token requirements. Remove local content-hashing and UUID fallbacks. Dask remains optional for indexing and reading. diff --git a/packages/zarr-indexing/docs/guide/integrations.md b/packages/zarr-indexing/docs/guide/integrations.md index a0bc721175..f2399e1f16 100644 --- a/packages/zarr-indexing/docs/guide/integrations.md +++ b/packages/zarr-indexing/docs/guide/integrations.md @@ -243,3 +243,14 @@ that implementation or reproduce its full worker/GPU lifecycle. ยท **API:** [API reference](../api/index.md) + +## Dask tokenization + +`LazyArray.__dask_tokenize__()` combines Dask's token for the wrapped source +with the serialized view transform. Dask owns source hashing, registered +normalizers, custom source hooks, and deterministic-token requirements. +Tokenization may read or hash source values. Dask is optional for indexing +and reading, but required when requesting a Dask token. + +The reader and partitioning are omitted because they must preserve values. +Changing a source after graph construction does not update existing Dask keys. diff --git a/packages/zarr-indexing/pyproject.toml b/packages/zarr-indexing/pyproject.toml index f2e286ab42..f9dff85ac8 100644 --- a/packages/zarr-indexing/pyproject.toml +++ b/packages/zarr-indexing/pyproject.toml @@ -99,9 +99,8 @@ extend = "../../pyproject.toml" target-version = "py312" [tool.ruff.lint.per-file-ignores] -# Chunk discovery and __dask_tokenize__ deliberately catch Exception: a -# foreign source's attributes or token hooks may fail, so these paths provide -# fallback metadata or tokens when ordinary exceptions occur. Configured here (not as +# Chunk discovery deliberately catches Exception: a foreign source's +# attributes may fail, so this path provides fallback metadata. Configured here (not as # noqa comments) because different ruff versions # have differed on whether these rules fire; RUF100 can remove unused noqa comments. "src/zarr_indexing/lazy_array.py" = ["BLE001", "S110"] diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index ed27d07f20..7326fc5256 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -12,8 +12,8 @@ ``` Selection construction does not read source values: `result()`, `__array__`, -and eager `__getitem__` perform reads. Tokenization can also read or hash source -data, depending on the source and tokenization path. `.lazy` operations inspect +and eager `__getitem__` perform reads. Source tokenization is delegated to +Dask and may inspect source values. `.lazy` operations inspect selection metadata and may copy or process supplied index arrays. Composition does not accumulate wrapper layers: a view of a view is still a single transform and retains its reader. @@ -138,7 +138,6 @@ import json import math import operator -import uuid from collections.abc import Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Protocol, cast @@ -173,10 +172,6 @@ __all__ = ["LazyArray", "Partition"] -# Above this declared byte count, the no-Dask token fallback adds a fresh UUID -# instead of digesting contents. See `_wrapped_token`. -_TOKEN_DIGEST_LIMIT = 1 << 20 - def _invoke_reader( reader: Reader, @@ -525,66 +520,6 @@ def _validate_prepared_parts(parts: Sequence[Partition], out_shape: tuple[int, . raise ValueError("prepared parts do not tile the view exactly") -# --------------------------------------------------------------------------- # -# Tokenization -# --------------------------------------------------------------------------- # - - -def _wrapped_token(array: Any) -> Any: - """A token for the wrapped array. - - In order of preference: the array's own `__dask_tokenize__`; - `dask.base.tokenize` when dask is importable (imported lazily โ€” this package - never requires it); otherwise a local fallback that digests the contents of - a small array. - - Tokens can differ depending on whether Dask is available and on the source - hook. This fallback does not provide a portable content identifier. - - Above `_TOKEN_DIGEST_LIMIT`, or when conversion is unavailable, the local - fallback adds a fresh UUID on each call, so repeated calls normally differ. - The returned tuple is still equal to itself. Below the limit, conversion - can read the source, and the hash is of its NumPy buffer bytes. Object-array - buffer bytes contain object references, not a recursive content snapshot. - """ - hook = getattr(array, "__dask_tokenize__", None) - if hook is not None: - try: - return hook() - # A failing source hook falls through to the remaining tokenization paths. - except Exception: # pragma: no cover - a hook that refuses to run - pass - try: - # dask is an optional peer, never a dependency of this package, so it is - # imported here and its absence is ordinary. - from dask.base import tokenize # pyright: ignore[reportMissingImports] - except ImportError: - pass - else: - return tokenize(array) - - shape = tuple(int(s) for s in getattr(array, "shape", ())) - dtype = getattr(array, "dtype", None) - structural = (type(array).__qualname__, shape, str(dtype)) - # A fresh identifier per call when contents cannot be identified. It - # is the shape and dtype that would otherwise be mistaken for an identity, - # so they are kept alongside it for a reader looking at a graph. - unidentified = (*structural, "unidentified", uuid.uuid4().hex) - - # Decide whether to digest the contents from the *declared* size. Measuring - # it by converting first would read the whole array โ€” a multi-gigabyte store - # pulled into memory by a token call, which is the opposite of the point. - itemsize = getattr(dtype, "itemsize", None) - if not isinstance(itemsize, int) or itemsize * math.prod(shape) > _TOKEN_DIGEST_LIMIT: - return unidentified - try: - contents = np.ascontiguousarray(array) - # Failed NumPy conversion leaves this source unidentified. - except Exception: - return unidentified - return (*structural, hashlib.sha256(contents.tobytes()).hexdigest()) - - # --------------------------------------------------------------------------- # # The wrapper # --------------------------------------------------------------------------- # @@ -1249,15 +1184,18 @@ def __dask_tokenize__(self) -> Any: tokens produce equal tokens; arbitrary semantically equivalent mappings are not guaranteed to serialize identically. - Source tokenization can read or hash data and need not be deterministic - on every fallback path; see `_wrapped_token`. The reader and partitioning - are omitted under the contract that they preserve values. Cache users - must also account for source mutation and the source's token semantics. + Dask tokenizes the wrapped source using its normal dispatch and + determinism policy. This may read or hash source values. Dask is + imported only when this method is called and is otherwise optional. + The reader and partitioning are omitted because they must preserve + values. Mutating a source does not update keys in existing Dask graphs. """ + from dask.base import tokenize # pyright: ignore[reportMissingImports] + canonical = json.dumps(self._transform.to_json(), sort_keys=True) return ( type(self).__qualname__, - _wrapped_token(self._array), + tokenize(self._array), hashlib.sha256(canonical.encode()).hexdigest(), ) diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 445e243aa9..5bb55ae076 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -1776,6 +1776,7 @@ def test_nonfirst_partition_transform_directly_addresses_its_array() -> None: def test_partition_token_encodes_its_public_global_transform() -> None: + pytest.importorskip("dask.base") source = np.arange(8) base = LazyArray.from_numpy(source) partition_view = list(base.with_parts((4,)).parts())[1].view @@ -1907,6 +1908,7 @@ def test_with_parts_validates_strictly(parts: Any, match: str) -> None: def test_dask_token_is_deterministic_and_discriminating() -> None: """Same data and same view token alike; a different selection differs.""" + pytest.importorskip("dask.base") data = reference() base = LazyArray(data) assert base.__dask_tokenize__() == LazyArray(reference()).__dask_tokenize__() @@ -1923,6 +1925,7 @@ def test_dask_token_is_deterministic_and_discriminating() -> None: def test_reader_and_partitioning_do_not_change_dask_identity() -> None: + pytest.importorskip("dask.base") base = LazyArray(reference()) token = base.__dask_tokenize__() assert base.with_reader(numpy_reader).__dask_tokenize__() == token @@ -2000,7 +2003,6 @@ def test_pickle_round_trip() -> None: view = LazyArray(reference()).with_parts((2, 2, 2)).lazy[1:6, ::2].lazy.oindex[[3, 0, 0], :, :] restored = pickle.loads(pickle.dumps(view)) assert restored.shape == view.shape - assert restored.__dask_tokenize__() == view.__dask_tokenize__() np.testing.assert_array_equal(np.asarray(restored.result()), np.asarray(view.result())) @@ -2010,7 +2012,7 @@ def test_pickle_round_trip() -> None: def test_dask_from_array_roundtrip() -> None: - """A `LazyArray` is a drop-in dask source โ€” no translation ceremony.""" + """Dask can tokenize and read a wrapper over a Zarr source.""" da = pytest.importorskip("dask.array") source = make_source("zarr") @@ -2435,25 +2437,6 @@ def test_a_masked_source_keeps_its_mask_when_the_view_is_empty(parts: Any) -> No assert np.asarray(got).shape == (3, 0), parts -def test_a_large_array_without_dask_refuses_to_claim_equality( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Without Dask, arrays above the digest limit receive distinct fallback tokens.""" - import sys - - monkeypatch.setitem(sys.modules, "dask.base", None) - big = np.zeros(1 << 19, dtype=np.int64) - other = big.copy() - other[0] = 1 - - assert LazyArray(big).__dask_tokenize__() != LazyArray(other).__dask_tokenize__() - assert LazyArray(big).__dask_tokenize__() != LazyArray(big).__dask_tokenize__() - - # Below the limit the contents are digested, so equal data still tokens alike. - small = np.zeros(8, dtype=np.int64) - assert LazyArray(small).__dask_tokenize__() == LazyArray(small.copy()).__dask_tokenize__() - - # --------------------------------------------------------------------------- # Completeness and partition spellings # --------------------------------------------------------------------------- @@ -2576,3 +2559,63 @@ def test_fancy_composition_over_an_empty_axis() -> None: scalar = composed.lazy.vindex[..., np.array(1)] assert scalar.shape == (2, 0) assert np.asarray(scalar.result()).shape == (2, 0) + + +@pytest.mark.parametrize("kind", ["numpy", "object", "masked", "registered", "hook"]) +def test_source_token_uses_dask_policy(kind: str) -> None: + dask_base = pytest.importorskip("dask.base") + + class RegisteredArray(ForeignArray): + pass + + class VersionedArray(ForeignArray): + def __dask_tokenize__(self) -> Any: + return ("versioned-source", 1) + + dask_base.normalize_token.register(RegisteredArray, lambda source: ("registered-source", 1)) + data = np.arange(4) + sources = { + "numpy": data, + "object": data.astype(object), + "masked": np.ma.masked_greater(data, 2), + "registered": RegisteredArray(data, None), + "hook": VersionedArray(data, None), + } + source = sources[kind] + assert LazyArray(source).__dask_tokenize__()[1] == dask_base.tokenize(source) + + +def test_source_token_preserves_dask_determinism_requirement() -> None: + dask_base = pytest.importorskip("dask.base") + dask_tokenize = pytest.importorskip("dask.tokenize") + + class UnserializableArray(ForeignArray): + def __reduce_ex__(self, protocol: int) -> Any: + raise TypeError("cannot serialize source") + + source = UnserializableArray(np.arange(4), None) + for value in (source, LazyArray(source)): + with pytest.raises(dask_tokenize.TokenizationError): + dask_base.tokenize(value, ensure_deterministic=True) + + +def test_source_token_preserves_hook_failure() -> None: + pytest.importorskip("dask.base") + + class RefusingArray(ForeignArray): + def __dask_tokenize__(self) -> Any: + raise RuntimeError("source version unavailable") + + with pytest.raises(RuntimeError, match="source version unavailable"): + LazyArray(RefusingArray(np.arange(4), None)).__dask_tokenize__() + + +def test_dask_is_only_required_for_tokenization(monkeypatch: pytest.MonkeyPatch) -> None: + import sys + + monkeypatch.setitem(sys.modules, "dask.base", None) + data = np.arange(4) + view = LazyArray(data).lazy[1:] + np.testing.assert_array_equal(view.result(), data[1:]) + with pytest.raises(ModuleNotFoundError): + view.__dask_tokenize__()