Support zarr-python 3.x

.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",

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

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

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

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

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

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

docs/source/zarr.rst

@@ -7,7 +7,35 @@ Zarr
 `Zarr-Python <https://zarr.readthedocs.io/en/stable/>`_ 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) <https://developer.nvidia.com/blog/gpudirect-storage/>`_ 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 <https://zarr.readthedocs.io/en/stable/user-guide/gpu.html#>`__ 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.
 
 

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):

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

python/kvikio/kvikio/zarr.py → 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,

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

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"

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