diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 772a2dcd95..bc68e3d370 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -36,7 +36,7 @@ repos: rev: 'v1.3.0' hooks: - id: mypy - additional_dependencies: [types-cachetools] + additional_dependencies: [types-cachetools, zarr] args: ["--config-file=python/kvikio/pyproject.toml", "python/kvikio/kvikio", "python/kvikio/tests", diff --git a/conda/environments/all_cuda-118_arch-aarch64.yaml b/conda/environments/all_cuda-118_arch-aarch64.yaml index 43dd9a4764..dea7a9f2e1 100644 --- a/conda/environments/all_cuda-118_arch-aarch64.yaml +++ b/conda/environments/all_cuda-118_arch-aarch64.yaml @@ -28,6 +28,7 @@ dependencies: - packaging - pre-commit - pytest +- pytest-asyncio - pytest-cov - pytest-timeout - python>=3.10,<3.13 @@ -39,5 +40,5 @@ dependencies: - sphinx-click - sphinx_rtd_theme - sysroot_linux-aarch64=2.28 -- zarr>=2.0.0,<3.0.0a0 +- zarr>=2.0.0,<4.0.0 name: all_cuda-118_arch-aarch64 diff --git a/conda/environments/all_cuda-118_arch-x86_64.yaml b/conda/environments/all_cuda-118_arch-x86_64.yaml index 5d42a932a3..c21a7d9a5d 100644 --- a/conda/environments/all_cuda-118_arch-x86_64.yaml +++ b/conda/environments/all_cuda-118_arch-x86_64.yaml @@ -30,6 +30,7 @@ dependencies: - packaging - pre-commit - pytest +- pytest-asyncio - pytest-cov - pytest-timeout - python>=3.10,<3.13 @@ -41,5 +42,5 @@ dependencies: - sphinx-click - sphinx_rtd_theme - sysroot_linux-64=2.28 -- zarr>=2.0.0,<3.0.0a0 +- zarr>=2.0.0,<4.0.0 name: all_cuda-118_arch-x86_64 diff --git a/conda/environments/all_cuda-128_arch-aarch64.yaml b/conda/environments/all_cuda-128_arch-aarch64.yaml index 16a54ffc36..7834ffcab5 100644 --- a/conda/environments/all_cuda-128_arch-aarch64.yaml +++ b/conda/environments/all_cuda-128_arch-aarch64.yaml @@ -28,6 +28,7 @@ dependencies: - packaging - pre-commit - pytest +- pytest-asyncio - pytest-cov - pytest-timeout - python>=3.10,<3.13 @@ -39,5 +40,5 @@ dependencies: - sphinx-click - sphinx_rtd_theme - sysroot_linux-aarch64=2.28 -- zarr>=2.0.0,<3.0.0a0 +- zarr>=2.0.0,<4.0.0 name: all_cuda-128_arch-aarch64 diff --git a/conda/environments/all_cuda-128_arch-x86_64.yaml b/conda/environments/all_cuda-128_arch-x86_64.yaml index df936afae7..e38e3729a3 100644 --- a/conda/environments/all_cuda-128_arch-x86_64.yaml +++ b/conda/environments/all_cuda-128_arch-x86_64.yaml @@ -28,6 +28,7 @@ dependencies: - packaging - pre-commit - pytest +- pytest-asyncio - pytest-cov - pytest-timeout - python>=3.10,<3.13 @@ -39,5 +40,5 @@ dependencies: - sphinx-click - sphinx_rtd_theme - sysroot_linux-64=2.28 -- zarr>=2.0.0,<3.0.0a0 +- zarr>=2.0.0,<4.0.0 name: all_cuda-128_arch-x86_64 diff --git a/conda/recipes/kvikio/meta.yaml b/conda/recipes/kvikio/meta.yaml index 01c98fd528..4465ac3f01 100644 --- a/conda/recipes/kvikio/meta.yaml +++ b/conda/recipes/kvikio/meta.yaml @@ -73,7 +73,7 @@ requirements: - python - numpy >=1.23,<3.0a0 - cupy >=12.0.0 - - zarr >=2.0.0,<3.0.0a0 + - zarr >=2.0.0,<4.0.0a0 # See https://github.com/zarr-developers/numcodecs/pull/475 - numcodecs !=0.12.0 - packaging diff --git a/dependencies.yaml b/dependencies.yaml index 3da0c3fdc2..1eae7775f4 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -393,7 +393,7 @@ dependencies: - output_types: [conda, requirements, pyproject] packages: - numpy>=1.23,<3.0a0 - - zarr>=2.0.0,<3.0.0a0 + - zarr>=2.0.0,<4.0.0 # See https://github.com/zarr-developers/numcodecs/pull/475 - numcodecs !=0.12.0 - packaging @@ -420,6 +420,7 @@ dependencies: packages: - rapids-dask-dependency==25.4.*,>=0.0.0a0 - pytest + - pytest-asyncio - pytest-cov - pytest-timeout - rangehttpserver diff --git a/docs/source/zarr.rst b/docs/source/zarr.rst index 82e6186026..019eff2767 100644 --- a/docs/source/zarr.rst +++ b/docs/source/zarr.rst @@ -7,7 +7,35 @@ Zarr `Zarr-Python `_ is the official Python package for reading and writing Zarr arrays. Its main feature is a NumPy-like array that translates array operations into file IO seamlessly. KvikIO provides a GPU backend to Zarr-Python that enables `GPUDirect Storage (GDS) `_ seamlessly. -The following is an example of how to use the convenience function :py:meth:`kvikio.zarr.open_cupy_array` +KvikIO supports either zarr-python 2.x or zarr-python 3.x. +However, the API provided in :mod:`kvikio.zarr` differs based on which version of zarr you have, following the differences between zarr-python 2.x and zarr-python 3.x. + + +Zarr Python 3.x +--------------- + +Zarr-python includes native support for reading Zarr chunks into device memory if you `configure Zarr `__ to use GPUs. +You can use any store, but KvikIO provides :py:class:`kvikio.zarr.GDSStore` to efficiently load data directly into GPU memory. + +.. code-block:: python + + >>> import zarr + >>> from kvikio.zarr import GDSStore + >>> zarr.config.enable_gpu() + >>> store = GDSStore(root="data.zarr") + >>> z = zarr.create_array( + ... store=store, shape=(100, 100), chunks=(10, 10), dtype="float32", overwrite=True + ... ) + >>> type(z[:10, :10]) + cupy.ndarray + + + +Zarr Python 2.x +--------------- + + +The following uses zarr-python 2.x, and is an example of how to use the convenience function :py:meth:`kvikio.zarr.open_cupy_array` to create a new Zarr array and how to open an existing Zarr array. diff --git a/python/kvikio/kvikio/benchmarks/single_node_io.py b/python/kvikio/kvikio/benchmarks/single_node_io.py index f5fc9057d1..e3b152cbaf 100644 --- a/python/kvikio/kvikio/benchmarks/single_node_io.py +++ b/python/kvikio/kvikio/benchmarks/single_node_io.py @@ -25,7 +25,13 @@ def get_zarr_compressors() -> Dict[str, Any]: import kvikio.zarr except ImportError: return {} - return {c.__name__.lower(): c for c in kvikio.zarr.nvcomp_compressors} + try: + compressors = kvikio.zarr.nvcomp_compressors + except AttributeError: + # zarr-python 3.x + return {} + else: + return {c.__name__.lower(): c for c in compressors} def create_data(nbytes): diff --git a/python/kvikio/kvikio/zarr/__init__.py b/python/kvikio/kvikio/zarr/__init__.py new file mode 100644 index 0000000000..7ec22c275a --- /dev/null +++ b/python/kvikio/kvikio/zarr/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +from importlib import metadata as _metadata + +from packaging.version import Version as _Version, parse as _parse + +if _parse(_metadata.version("zarr")) >= _Version("3.0.0"): + from ._zarr_python_3 import * # noqa: F401,F403 +else: + from ._zarr_python_2 import * # type: ignore[assignment] # noqa: F401,F403 diff --git a/python/kvikio/kvikio/zarr.py b/python/kvikio/kvikio/zarr/_zarr_python_2.py similarity index 96% rename from python/kvikio/kvikio/zarr.py rename to python/kvikio/kvikio/zarr/_zarr_python_2.py index a0d6dd8fb3..38df925f0a 100644 --- a/python/kvikio/kvikio/zarr.py +++ b/python/kvikio/kvikio/zarr/_zarr_python_2.py @@ -17,7 +17,6 @@ import zarr.creation import zarr.errors import zarr.storage -import zarr.util from numcodecs.abc import Codec from numcodecs.compat import ensure_contiguous_ndarray_like from numcodecs.registry import register_codec @@ -38,7 +37,7 @@ supported = parse(zarr.__version__) >= parse(MINIMUM_ZARR_VERSION) -class GDSStore(zarr.storage.DirectoryStore): +class GDSStore(zarr.storage.DirectoryStore): # type: ignore[name-defined] """GPUDirect Storage (GDS) class using directories and files. This class works like `zarr.storage.DirectoryStore` but implements @@ -338,12 +337,15 @@ def open_cupy_array( # In order to handle "a", we start by trying to open the file in read mode. try: ret = zarr.open_array( - store=kvikio.zarr.GDSStore(path=store), + store=kvikio.zarr.GDSStore(path=store), # type: ignore[call-arg] mode="r+", meta_array=meta_array, **kwargs, ) - except (zarr.errors.ContainsGroupError, zarr.errors.ArrayNotFoundError): + except ( + zarr.errors.ContainsGroupError, + zarr.errors.ArrayNotFoundError, # type: ignore[attr-defined] + ): # If we are reading, this is a genuine error. if mode in ("r", "r+"): raise @@ -356,7 +358,7 @@ def open_cupy_array( compat_lz4 = CompatCompressor.lz4() if ret.compressor == compat_lz4.cpu: ret = zarr.open_array( - store=kvikio.zarr.GDSStore( + store=kvikio.zarr.GDSStore( # type: ignore[call-arg] path=store, compressor_config_overwrite=compat_lz4.cpu.get_config(), decompressor_config_overwrite=compat_lz4.gpu.get_config(), @@ -386,7 +388,7 @@ def open_cupy_array( decompressor_config_overwrite = None return zarr.open_array( - store=kvikio.zarr.GDSStore( + store=kvikio.zarr.GDSStore( # type: ignore[call-arg] path=store, compressor_config_overwrite=compressor_config_overwrite, decompressor_config_overwrite=decompressor_config_overwrite, diff --git a/python/kvikio/kvikio/zarr/_zarr_python_3.py b/python/kvikio/kvikio/zarr/_zarr_python_3.py new file mode 100644 index 0000000000..69d008ecf9 --- /dev/null +++ b/python/kvikio/kvikio/zarr/_zarr_python_3.py @@ -0,0 +1,135 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# See file LICENSE for terms. + +import asyncio +import os +from pathlib import Path + +import zarr.storage +from zarr.abc.store import ( + ByteRequest, + OffsetByteRequest, + RangeByteRequest, + SuffixByteRequest, +) +from zarr.core.buffer import Buffer, BufferPrototype +from zarr.core.buffer.core import default_buffer_prototype + +import kvikio + +# The GDSStore implementation follows the `LocalStore` implementation +# at https://github.com/zarr-developers/zarr-python/blob/main/src/zarr/storage/_local.py +# with differences coming swapping in `cuFile` for the stdlib open file object. + + +def _get( + path: Path, prototype: BufferPrototype, byte_range: ByteRequest | None +) -> Buffer: + file_size = os.path.getsize(path) + file_offset: int + + match byte_range: + case None: + nbytes = file_size + file_offset = 0 + case OffsetByteRequest(): + nbytes = max(0, file_size - byte_range.offset) + file_offset = byte_range.offset + case RangeByteRequest(): + nbytes = byte_range.end - byte_range.start + file_offset = byte_range.start + case SuffixByteRequest(): + nbytes = byte_range.suffix + file_offset = max(0, file_size - byte_range.suffix) + case _: + # This isn't allowed by mypy, but the tests assert we raise + # something here. + raise TypeError(f"Unexpected byte_range, got {byte_range}") + + # kvikio doesn't support reading past the end of a file. Some zarr tests + # rely on this behavior: to "read" 3 bytes out of a 0 byte file, or to + # "seek" past the end of a file with file_offset. The semantics seem to + # be roughly the same as slicing an empty bytestring. + + nbytes = min(nbytes, file_size) + file_offset = min(file_offset, file_size) + + raw = prototype.nd_buffer.create(shape=(nbytes,), dtype="b").as_ndarray_like() + buf = prototype.buffer.from_array_like(raw) + + with kvikio.CuFile(path) as f: + # Note: this currently creates an IOFuture and then blocks + # on reading it. The blocking read means this is in a regular + # sync function, and so this must be run in a threadpool. + future = f.pread(raw, size=nbytes, file_offset=file_offset) + future.get() # blocks + + return buf + + +def _put( + path: Path, + value: Buffer, + start: int | None = None, + exclusive: bool = False, +) -> int | None: + path.parent.mkdir(parents=True, exist_ok=True) + if start is not None: + with kvikio.CuFile(path, "r+b") as f: + f.write(value.as_array_like(), file_offset=start) + return None + else: + buf = value.as_array_like() + if exclusive: + if path.exists(): + raise FileExistsError(f"File exists: {path}") + mode = "wb" + with kvikio.CuFile(path, flags=mode) as f: + return f.write(buf) + + +class GDSStore(zarr.storage.LocalStore): + def __repr__(self) -> str: + return f"{type(self).__name__}('{self}')" + + async def get( + self, + key: str, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + if prototype is None: + prototype = default_buffer_prototype() + if not self._is_open: + await self._open() + assert isinstance(key, str) + path = self.root / key + + try: + return await asyncio.to_thread(_get, path, prototype, byte_range) + except (FileNotFoundError, IsADirectoryError, NotADirectoryError): + return None + + async def set(self, key: str, value: Buffer) -> None: + return await self._set(key, value) + + async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None: + if not self._is_open: + await self._open() + self._check_writable() + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"LocalStore.set(): `value` must be a Buffer instance. Got an " + f"instance of {type(value)} instead." + ) + path = self.root / key + + await asyncio.to_thread(_put, path, value, start=None, exclusive=exclusive) + + +# Matching the check that zarr.__version__ > 2.15 that's +# part of the public API for our zarr 2.x support +# This module is behind a check that zarr.__version__ > 3 +# so we can just assume it's already checked and supported. +supported = True diff --git a/python/kvikio/pyproject.toml b/python/kvikio/pyproject.toml index 5548e579e3..dc5207c69e 100644 --- a/python/kvikio/pyproject.toml +++ b/python/kvikio/pyproject.toml @@ -25,7 +25,7 @@ dependencies = [ "numpy>=1.23,<3.0a0", "nvidia-nvcomp==4.2.0.11", "packaging", - "zarr>=2.0.0,<3.0.0a0", + "zarr>=2.0.0,<4.0.0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. classifiers = [ "Intended Audience :: Developers", @@ -44,6 +44,7 @@ test = [ "cuda-python>=11.8.5,<12.0a0", "moto[server]>=4.0.8", "pytest", + "pytest-asyncio", "pytest-cov", "pytest-timeout", "rangehttpserver", @@ -109,6 +110,12 @@ skip = [ [tool.mypy] ignore_missing_imports = true +exclude = [ + # we type check against zarr-python 3.x + # and ignore modules using 2.x + "python/kvikio/kvikio/zarr/_zarr_python_2.py", + "python/kvikio/tests/test_nvcomp_codec.py", +] [project.entry-points."numcodecs.codecs"] nvcomp_batch = "kvikio.nvcomp_codec:NvCompBatchCodec" @@ -157,3 +164,4 @@ filterwarnings = [ markers = [ "cufile: tests to skip if cuFile isn't available e.g. run with `pytest -m 'not cufile'`" ] +asyncio_mode = "auto" diff --git a/python/kvikio/tests/conftest.py b/python/kvikio/tests/conftest.py index 120f862ae6..07636095eb 100644 --- a/python/kvikio/tests/conftest.py +++ b/python/kvikio/tests/conftest.py @@ -4,6 +4,7 @@ import contextlib import multiprocessing as mp import subprocess +from multiprocessing.connection import Connection from typing import Iterable import pytest @@ -13,7 +14,7 @@ mp = mp.get_context("spawn") # type: ignore -def command_server(conn): +def command_server(conn: Connection) -> None: """Server to run commands given through `conn`""" while True: # Get the next command to run diff --git a/python/kvikio/tests/test_benchmarks.py b/python/kvikio/tests/test_benchmarks.py index 6707a86efc..8450fdfc25 100644 --- a/python/kvikio/tests/test_benchmarks.py +++ b/python/kvikio/tests/test_benchmarks.py @@ -7,6 +7,7 @@ from pathlib import Path import pytest +from packaging.version import parse import kvikio @@ -34,9 +35,16 @@ def test_single_node_io(run_cmd, tmp_path, api): if "zarr" in api: kz = pytest.importorskip("kvikio.zarr") + import zarr + if not kz.supported: pytest.skip(f"requires Zarr >={kz.MINIMUM_ZARR_VERSION}") + if parse(zarr.__version__) >= parse("3.0.0"): + pytest.skip( + "requires Zarr<3", + ) + retcode = run_cmd( cmd=[ sys.executable or "python", @@ -65,9 +73,16 @@ def test_zarr_io(run_cmd, tmp_path, api): """Test benchmarks/zarr_io.py""" kz = pytest.importorskip("kvikio.zarr") + import zarr + if not kz.supported: pytest.skip(f"requires Zarr >={kz.MINIMUM_ZARR_VERSION}") + if parse(zarr.__version__) >= parse("3.0.0"): + pytest.skip( + "requires Zarr<3", + ) + retcode = run_cmd( cmd=[ sys.executable or "python", diff --git a/python/kvikio/tests/test_examples.py b/python/kvikio/tests/test_examples.py index 07be1fc156..f32485b6c4 100644 --- a/python/kvikio/tests/test_examples.py +++ b/python/kvikio/tests/test_examples.py @@ -1,4 +1,4 @@ -# Copyright (c) 2021-2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved. # See file LICENSE for terms. import os @@ -6,6 +6,7 @@ from pathlib import Path import pytest +from packaging.version import parse import kvikio @@ -24,7 +25,9 @@ def test_zarr_cupy_nvcomp(tmp_path, monkeypatch): """Test examples/zarr_cupy_nvcomp.py""" # `examples/zarr_cupy_nvcomp.py` requires the Zarr submodule - pytest.importorskip("kvikio.zarr") + zarr = pytest.importorskip("zarr") + if parse(zarr.__version__) >= parse("3.0.0"): + pytest.skip(reason="Requires zarr<3") monkeypatch.syspath_prepend(str(examples_path)) import_module("zarr_cupy_nvcomp").main(tmp_path / "test-file") diff --git a/python/kvikio/tests/test_nvcomp_codec.py b/python/kvikio/tests/test_nvcomp_codec.py index be8cb3d922..29e50ad64b 100644 --- a/python/kvikio/tests/test_nvcomp_codec.py +++ b/python/kvikio/tests/test_nvcomp_codec.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2023-2025, NVIDIA CORPORATION. All rights reserved. # See file LICENSE for terms. import itertools as it @@ -7,6 +7,8 @@ import cupy as cp import numcodecs import numpy as np +import packaging +import packaging.version import pytest import zarr from numpy.testing import assert_equal @@ -24,6 +26,13 @@ SUPPORTED_CODECS = [LZ4_ALGO, GDEFLATE_ALGO, SNAPPY_ALGO, ZSTD_ALGO, DEFLATE_ALGO] +def skip_if_zarr_v3(): + return pytest.mark.skipif( + packaging.version.parse(zarr.__version__) >= packaging.version.Version("3.0.0"), + reason="zarr 3.x not supported.", + ) + + def _get_codec(algo: str, **kwargs): codec_args = {"id": NVCOMP_CODEC_ID, "algorithm": algo, "options": kwargs} return numcodecs.registry.get_codec(codec_args) @@ -68,6 +77,7 @@ def test_basic(algo: str, shape): @pytest.mark.parametrize("algo", SUPPORTED_CODECS) +@skip_if_zarr_v3() def test_basic_zarr(algo: str, shape_chunks): shape, chunks = shape_chunks @@ -118,6 +128,7 @@ def test_batch_comp_decomp(algo: str, chunk_sizes, out: str): @pytest.mark.parametrize("algo", SUPPORTED_CODECS) +@skip_if_zarr_v3() def test_comp_decomp(algo: str, shape_chunks): shape, chunks = shape_chunks @@ -149,6 +160,7 @@ def test_comp_decomp(algo: str, shape_chunks): ("gdeflate", {"algo": 1}), # low-throughput, high compression ratio algo ], ) +@skip_if_zarr_v3() def test_codec_options(algo, options): codec = NvCompBatchCodec(algo, options) @@ -162,6 +174,7 @@ def test_codec_options(algo, options): assert_equal(z[:], data[:]) +@skip_if_zarr_v3() def test_codec_invalid_options(): # There are currently only 3 supported algos in Gdeflate codec = NvCompBatchCodec(GDEFLATE_ALGO, options={"algo": 10}) @@ -179,6 +192,7 @@ def test_codec_invalid_options(): ("zstd", ZSTD_ALGO), ], ) +@skip_if_zarr_v3() def test_cpu_comp_gpu_decomp(cpu_algo, gpu_algo): cpu_codec = numcodecs.registry.get_codec({"id": cpu_algo}) gpu_codec = _get_codec(gpu_algo) @@ -203,6 +217,7 @@ def test_cpu_comp_gpu_decomp(cpu_algo, gpu_algo): assert_equal(z1[:], z2[:]) +@skip_if_zarr_v3() def test_lz4_codec_header(shape_chunks): shape, chunks = shape_chunks diff --git a/python/kvikio/tests/test_zarr.py b/python/kvikio/tests/test_zarr.py index ba42e35a96..a793e2568e 100644 --- a/python/kvikio/tests/test_zarr.py +++ b/python/kvikio/tests/test_zarr.py @@ -1,4 +1,4 @@ -# Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved. # See file LICENSE for terms. @@ -6,6 +6,7 @@ import numpy import pytest +from packaging.version import parse cupy = pytest.importorskip("cupy") zarr = pytest.importorskip("zarr") @@ -19,6 +20,12 @@ allow_module_level=True, ) +if parse(zarr.__version__) >= parse("3.0.0"): + pytest.skip( + "requires Zarr<3", + allow_module_level=True, + ) + @pytest.fixture def store(tmp_path): diff --git a/python/kvikio/tests/test_zarr_v3.py b/python/kvikio/tests/test_zarr_v3.py new file mode 100644 index 0000000000..e44de36d22 --- /dev/null +++ b/python/kvikio/tests/test_zarr_v3.py @@ -0,0 +1,81 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# See file LICENSE for terms. + +import pathlib + +import cupy as cp +import pytest + +import kvikio.zarr + +pytest.importorskip("zarr", minversion="3.0.0") + + +import zarr.core.buffer # noqa: E402 +import zarr.storage # noqa: E402 +from zarr.core.buffer.gpu import Buffer # noqa: E402 +from zarr.testing.store import StoreTests # noqa: E402 + + +@pytest.mark.asyncio +async def test_basic(tmp_path: pathlib.Path) -> None: + src = zarr.storage.LocalStore(tmp_path) + arr = zarr.create_array(src, name="a", shape=(10,), dtype="u4", zarr_format=3) + arr[:5] = 0 + arr[5:] = 1 + + assert await src.exists("a/zarr.json") + store = kvikio.zarr.GDSStore(tmp_path, read_only=True) + assert await store.exists("a/zarr.json") + + # regular works + zarr.open_array(src, path="a", zarr_format=3) + + # read + with zarr.config.enable_gpu(): + arr = zarr.open_array(store, path="a", zarr_format=3) + result = arr[:] + assert isinstance(result, cp.ndarray) + expected = cp.array([0] * 5 + [1] * 5, dtype="uint32") + cp.testing.assert_array_equal(result, expected) + + +# These tests use zarr-python's StoreTests +# https://zarr.readthedocs.io/en/stable/api/zarr/testing/store/index.html#zarr.testing.store.StoreTests +# which provide a set of unit tests each implementation is expected to pass. + + +class TestKvikIOStore(StoreTests[kvikio.zarr.GDSStore, Buffer]): + store_cls = kvikio.zarr.GDSStore + buffer_cls = Buffer + + async def get(self, store: kvikio.zarr.GDSStore, key: str) -> Buffer: + return self.buffer_cls.from_bytes((store.root / key).read_bytes()) + + async def set( + self, + store: kvikio.zarr.GDSStore, + key: str, + value: Buffer, # type: ignore[override] + ) -> None: + parent = (store.root / key).parent + if not parent.exists(): + parent.mkdir(parents=True) + (store.root / key).write_bytes(value.to_bytes()) + + @pytest.fixture + def store_kwargs(self, tmpdir: pathlib.Path) -> dict[str, str]: + kwargs = {"root": str(tmpdir)} + return kwargs + + @pytest.fixture + async def store(self, store_kwargs: dict[str, str]) -> kvikio.zarr.GDSStore: + # ignore Argument 1 has incompatible type "**Dict[str, str]"; expected "bool" + return self.store_cls(**store_kwargs) # type: ignore[arg-type] + + @pytest.fixture + async def store_not_open( + self, store_kwargs: dict[str, str] + ) -> kvikio.zarr.GDSStore: + # ignore Argument 1 has incompatible type "**Dict[str, str]"; expected "bool" + return self.store_cls(**store_kwargs) # type: ignore[arg-type]