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]