From 70abb2c84f4cfa0e9817e83600bee42a7a8aa808 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 19 Feb 2025 13:24:34 +0000
Subject: [PATCH 01/21] Support zarr-python 3.x

This implements a zarr-python 3.x compatible
GDSStore.
---
 python/kvikio/kvikio/zarr_v3.py     | 125 ++++++++++++++++++++++++++++
 python/kvikio/pyproject.toml        |   3 +-
 python/kvikio/tests/test_zarr.py    |  10 +--
 python/kvikio/tests/test_zarr_v3.py |  67 +++++++++++++++
 4 files changed, 199 insertions(+), 6 deletions(-)
 create mode 100644 python/kvikio/kvikio/zarr_v3.py
 create mode 100644 python/kvikio/tests/test_zarr_v3.py

diff --git a/python/kvikio/kvikio/zarr_v3.py b/python/kvikio/kvikio/zarr_v3.py
new file mode 100644
index 0000000000..c500c5e188
--- /dev/null
+++ b/python/kvikio/kvikio/zarr_v3.py
@@ -0,0 +1,125 @@
+# 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
+
+
+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():  # type: ignore[misc]
+            nbytes = max(0, file_size - byte_range.offset)
+            file_offset = byte_range.offset
+        case RangeByteRequest():  # type: ignore[misc]
+            nbytes = byte_range.end - byte_range.start
+            file_offset = byte_range.start
+        case SuffixByteRequest():  # type: ignore[misc]
+            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:
+            # TODO: seems like this isn't tested
+            f.write(value.as_array_like(), file_offset=start)
+        return None
+    else:
+        view = memoryview(value.as_numpy_array().tobytes())
+        if exclusive:
+            if path.exists():
+                raise FileExistsError(f"File exists: {path}")
+        mode = "wb"
+        with kvikio.CuFile(path, flags=mode) as f:
+            return f.write(view)
+
+
+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:
+        # docstring inherited
+        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:
+        # docstring inherited
+        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)
diff --git a/python/kvikio/pyproject.toml b/python/kvikio/pyproject.toml
index 53963b4ba5..ebd7bfc7a8 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",
 ] # 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",
@@ -156,3 +156,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/test_zarr.py b/python/kvikio/tests/test_zarr.py
index ba42e35a96..d04af981fe 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.
 
 
@@ -20,10 +20,10 @@
     )
 
 
-@pytest.fixture
-def store(tmp_path):
-    """Fixture that creates a GDS Store"""
-    return kvikio_zarr.GDSStore(tmp_path / "test-file.zarr")
+# @pytest.fixture
+# def store(tmp_path):
+#     """Fixture that creates a GDS Store"""
+#     return kvikio_zarr.GDSStore(tmp_path / "test-file.zarr")
 
 
 def test_direct_store_access(store, xp):
diff --git a/python/kvikio/tests/test_zarr_v3.py b/python/kvikio/tests/test_zarr_v3.py
new file mode 100644
index 0000000000..9eaa2c5330
--- /dev/null
+++ b/python/kvikio/tests/test_zarr_v3.py
@@ -0,0 +1,67 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+import pathlib
+
+import cupy as cp
+import pytest
+import zarr.core.buffer
+import zarr.storage
+from zarr.core.buffer.cpu import Buffer
+from zarr.testing.store import StoreTests
+
+import kvikio.zarr_v3
+
+
+@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_v3.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)
+
+
+class TestKvikIOStore(StoreTests[kvikio.zarr_v3.GDSStore, Buffer]):
+    store_cls = kvikio.zarr_v3.GDSStore
+    buffer_cls = Buffer
+
+    async def get(self, store: kvikio.zarr_v3.GDSStore, key: str) -> Buffer:
+        return self.buffer_cls.from_bytes((store.root / key).read_bytes())
+
+    async def set(
+        self, store: kvikio.zarr_v3.GDSStore, key: str, value: Buffer
+    ) -> 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_v3.GDSStore:
+        return self.store_cls(**store_kwargs)
+
+    @pytest.fixture
+    async def store_not_open(
+        self, store_kwargs: dict[str, str]
+    ) -> kvikio.zarr_v3.GDSStore:
+        return self.store_cls(**store_kwargs)

From a870ff798441d7eccde274953eb673b739811837 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Tue, 25 Feb 2025 14:34:41 +0000
Subject: [PATCH 02/21] Avoid memcopy in set

---
 python/kvikio/kvikio/zarr_v3.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/python/kvikio/kvikio/zarr_v3.py b/python/kvikio/kvikio/zarr_v3.py
index c500c5e188..0b485b6a35 100644
--- a/python/kvikio/kvikio/zarr_v3.py
+++ b/python/kvikio/kvikio/zarr_v3.py
@@ -73,16 +73,18 @@ def _put(
     if start is not None:
         with kvikio.CuFile(path, "r+b") as f:
             # TODO: seems like this isn't tested
+            # https://github.com/zarr-developers/zarr-python/issues/2859
             f.write(value.as_array_like(), file_offset=start)
         return None
     else:
-        view = memoryview(value.as_numpy_array().tobytes())
+        # view = memoryview(value.as_numpy_array().tobytes())
+        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(view)
+            return f.write(buf)
 
 
 class GDSStore(zarr.storage.LocalStore):

From 0d7248513c7a98ac5af4e289f4b32c2ada88cf68 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Tue, 25 Feb 2025 14:54:24 +0000
Subject: [PATCH 03/21] fixup

---
 python/kvikio/kvikio/zarr.py     | 16 ++++++++++++++--
 python/kvikio/tests/test_zarr.py | 10 +++++-----
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/python/kvikio/kvikio/zarr.py b/python/kvikio/kvikio/zarr.py
index 1a8e7f3848..2e3c1d401e 100644
--- a/python/kvikio/kvikio/zarr.py
+++ b/python/kvikio/kvikio/zarr.py
@@ -1,10 +1,11 @@
-# Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 from __future__ import annotations
 
 import contextlib
 import os
 import os.path
+import warnings
 from abc import abstractmethod
 from typing import Any, Literal, Mapping, Optional, Sequence, Union
 
@@ -17,7 +18,18 @@
 import zarr.creation
 import zarr.errors
 import zarr.storage
-import zarr.util
+
+try:
+    import zarr.util
+except ImportError:
+    message = (
+        "Failed to import 'zarr.util' while importing 'kvikio.zarr'.  "
+        "If you have zarr-python>=3.0, try importing 'kvikio.zarr_v3' instead."
+    )
+    warnings.warn(message, UserWarning)
+    raise
+
+
 from numcodecs.abc import Codec
 from numcodecs.compat import ensure_contiguous_ndarray_like
 from numcodecs.registry import register_codec
diff --git a/python/kvikio/tests/test_zarr.py b/python/kvikio/tests/test_zarr.py
index d04af981fe..ba42e35a96 100644
--- a/python/kvikio/tests/test_zarr.py
+++ b/python/kvikio/tests/test_zarr.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
 
@@ -20,10 +20,10 @@
     )
 
 
-# @pytest.fixture
-# def store(tmp_path):
-#     """Fixture that creates a GDS Store"""
-#     return kvikio_zarr.GDSStore(tmp_path / "test-file.zarr")
+@pytest.fixture
+def store(tmp_path):
+    """Fixture that creates a GDS Store"""
+    return kvikio_zarr.GDSStore(tmp_path / "test-file.zarr")
 
 
 def test_direct_store_access(store, xp):

From 4e9a9f23c81ad89689ad71b6cbc9688329b63286 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Tue, 25 Feb 2025 15:54:16 +0000
Subject: [PATCH 04/21] fixup

---
 .pre-commit-config.yaml      | 2 +-
 python/kvikio/pyproject.toml | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index d7185880e0..a5b5d58d14 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/python/kvikio/pyproject.toml b/python/kvikio/pyproject.toml
index ebd7bfc7a8..21594bcf8b 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",
+    "zarr>=2.0.0,<3.0.0a0",
 ] # 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",

From 4e08f8cefec6dade11818cc5d3317a71ed62e5d7 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 26 Feb 2025 11:57:03 -0800
Subject: [PATCH 05/21] nvtx

---
 python/kvikio/kvikio/zarr_v3.py | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/python/kvikio/kvikio/zarr_v3.py b/python/kvikio/kvikio/zarr_v3.py
index 0b485b6a35..5b7efb522a 100644
--- a/python/kvikio/kvikio/zarr_v3.py
+++ b/python/kvikio/kvikio/zarr_v3.py
@@ -2,6 +2,7 @@
 # See file LICENSE for terms.
 
 import asyncio
+import contextlib
 import os
 from pathlib import Path
 
@@ -17,6 +18,19 @@
 
 import kvikio
 
+try:
+    import nvtx
+except ImportError:
+    HAS_NVTX = False
+else:
+    HAS_NVTX = True
+
+
+if HAS_NVTX:
+    annotate = nvtx.annotate
+else:
+    annotate = contextlib.nullcontext
+
 
 def _get(
     path: Path, prototype: BufferPrototype, byte_range: ByteRequest | None
@@ -104,8 +118,14 @@ async def get(
             await self._open()
         assert isinstance(key, str)
         path = self.root / key
+
+        if HAS_NVTX:
+            kwargs = {"message": "kvikio.zarr.get", "domain": "Zarr"}
+        else:
+            kwargs = {}
         try:
-            return await asyncio.to_thread(_get, path, prototype, byte_range)
+            with annotate(**kwargs):
+                return await asyncio.to_thread(_get, path, prototype, byte_range)
         except (FileNotFoundError, IsADirectoryError, NotADirectoryError):
             return None
 
@@ -124,4 +144,11 @@ async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None:
                 f"instance of {type(value)} instead."
             )
         path = self.root / key
-        await asyncio.to_thread(_put, path, value, start=None, exclusive=exclusive)
+
+        if HAS_NVTX:
+            kwargs = {"message": "kvikio.zarr.get", "domain": "Zarr"}
+        else:
+            kwargs = {}
+
+        with annotate(**kwargs):
+            await asyncio.to_thread(_put, path, value, start=None, exclusive=exclusive)

From 9ad5c8ac67178bfc778dcaa3e55077b3b10598c6 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 26 Feb 2025 11:57:31 -0800
Subject: [PATCH 06/21] updated deps

---
 conda/environments/all_cuda-118_arch-aarch64.yaml | 2 +-
 conda/environments/all_cuda-118_arch-x86_64.yaml  | 2 +-
 conda/environments/all_cuda-128_arch-aarch64.yaml | 2 +-
 conda/environments/all_cuda-128_arch-x86_64.yaml  | 2 +-
 dependencies.yaml                                 | 2 +-
 python/kvikio/pyproject.toml                      | 2 +-
 6 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/conda/environments/all_cuda-118_arch-aarch64.yaml b/conda/environments/all_cuda-118_arch-aarch64.yaml
index 8471e29a1e..e0dafa09df 100644
--- a/conda/environments/all_cuda-118_arch-aarch64.yaml
+++ b/conda/environments/all_cuda-118_arch-aarch64.yaml
@@ -38,5 +38,5 @@ dependencies:
 - sphinx-click
 - sphinx_rtd_theme
 - sysroot_linux-aarch64=2.28
-- zarr>=2.0.0,<3.0.0a0
+- zarr>=2.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 a42d41f746..5a297e8094 100644
--- a/conda/environments/all_cuda-118_arch-x86_64.yaml
+++ b/conda/environments/all_cuda-118_arch-x86_64.yaml
@@ -40,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
 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 4f184129bc..8e904c91b2 100644
--- a/conda/environments/all_cuda-128_arch-aarch64.yaml
+++ b/conda/environments/all_cuda-128_arch-aarch64.yaml
@@ -38,5 +38,5 @@ dependencies:
 - sphinx-click
 - sphinx_rtd_theme
 - sysroot_linux-aarch64=2.28
-- zarr>=2.0.0,<3.0.0a0
+- zarr>=2.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 d3df1235b6..f97dabf576 100644
--- a/conda/environments/all_cuda-128_arch-x86_64.yaml
+++ b/conda/environments/all_cuda-128_arch-x86_64.yaml
@@ -38,5 +38,5 @@ dependencies:
 - sphinx-click
 - sphinx_rtd_theme
 - sysroot_linux-64=2.28
-- zarr>=2.0.0,<3.0.0a0
+- zarr>=2.0.0
 name: all_cuda-128_arch-x86_64
diff --git a/dependencies.yaml b/dependencies.yaml
index e951fe78cd..f8694838b6 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
           # See https://github.com/zarr-developers/numcodecs/pull/475
           - numcodecs !=0.12.0
           - packaging
diff --git a/python/kvikio/pyproject.toml b/python/kvikio/pyproject.toml
index 21594bcf8b..01ea15e9ed 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",
 ] # 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",

From 7bbfe9098f28605b92a9c2c35aa639bf9f2ed20d Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 26 Feb 2025 20:05:21 +0000
Subject: [PATCH 07/21] temporarily disable zarr

---
 python/kvikio/pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/kvikio/pyproject.toml b/python/kvikio/pyproject.toml
index 01ea15e9ed..9afbb6299a 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",
+    # "zarr>=2.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",

From e8c8f54133a9e52b998dc88be181fa5160b1ece2 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Fri, 28 Feb 2025 18:09:42 +0000
Subject: [PATCH 08/21] fixups

---
 .../environments/all_cuda-118_arch-aarch64.yaml |  1 +
 .../environments/all_cuda-118_arch-x86_64.yaml  |  1 +
 .../environments/all_cuda-128_arch-aarch64.yaml |  1 +
 .../environments/all_cuda-128_arch-x86_64.yaml  |  1 +
 dependencies.yaml                               |  1 +
 python/kvikio/kvikio/zarr.py                    |  7 +++++--
 python/kvikio/kvikio/zarr_v3.py                 |  2 +-
 python/kvikio/pyproject.toml                    |  9 ++++++++-
 python/kvikio/tests/conftest.py                 |  5 +++--
 python/kvikio/tests/test_nvcomp_codec.py        | 17 ++++++++++++++++-
 python/kvikio/tests/test_zarr_v3.py             |  8 +++++---
 11 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/conda/environments/all_cuda-118_arch-aarch64.yaml b/conda/environments/all_cuda-118_arch-aarch64.yaml
index e0dafa09df..4a3bcb2d13 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
 - python>=3.10,<3.13
 - rangehttpserver
diff --git a/conda/environments/all_cuda-118_arch-x86_64.yaml b/conda/environments/all_cuda-118_arch-x86_64.yaml
index 5a297e8094..1f41b02a7b 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
 - python>=3.10,<3.13
 - rangehttpserver
diff --git a/conda/environments/all_cuda-128_arch-aarch64.yaml b/conda/environments/all_cuda-128_arch-aarch64.yaml
index 8e904c91b2..80f78d9c90 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
 - python>=3.10,<3.13
 - rangehttpserver
diff --git a/conda/environments/all_cuda-128_arch-x86_64.yaml b/conda/environments/all_cuda-128_arch-x86_64.yaml
index f97dabf576..092974063f 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
 - python>=3.10,<3.13
 - rangehttpserver
diff --git a/dependencies.yaml b/dependencies.yaml
index f8694838b6..b81575d819 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -421,6 +421,7 @@ dependencies:
           - rapids-dask-dependency==25.4.*,>=0.0.0a0
           - pytest
           - pytest-cov
+          - pytest-asyncio
           - rangehttpserver
           - boto3>=1.21.21
       - output_types: [requirements, pyproject]
diff --git a/python/kvikio/kvikio/zarr.py b/python/kvikio/kvikio/zarr.py
index 2e3c1d401e..31dcd6b943 100644
--- a/python/kvikio/kvikio/zarr.py
+++ b/python/kvikio/kvikio/zarr.py
@@ -50,7 +50,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
@@ -369,7 +369,10 @@ def open_cupy_array(
                 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
diff --git a/python/kvikio/kvikio/zarr_v3.py b/python/kvikio/kvikio/zarr_v3.py
index 5b7efb522a..f2bf113309 100644
--- a/python/kvikio/kvikio/zarr_v3.py
+++ b/python/kvikio/kvikio/zarr_v3.py
@@ -146,7 +146,7 @@ async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None:
         path = self.root / key
 
         if HAS_NVTX:
-            kwargs = {"message": "kvikio.zarr.get", "domain": "Zarr"}
+            kwargs = {"message": "kvikio.zarr.set", "domain": "Zarr"}
         else:
             kwargs = {}
 
diff --git a/python/kvikio/pyproject.toml b/python/kvikio/pyproject.toml
index 9afbb6299a..ddbdbde9c7 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",
+    "zarr>=2.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",
     "rangehttpserver",
     "rapids-dask-dependency==25.4.*,>=0.0.0a0",
@@ -108,6 +109,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.py",
+    "python/kvikio/tests/test_nvcomp_codec.py",
+]
 
 [project.entry-points."numcodecs.codecs"]
 nvcomp_batch = "kvikio.nvcomp_codec:NvCompBatchCodec"
diff --git a/python/kvikio/tests/conftest.py b/python/kvikio/tests/conftest.py
index c1cc77026e..17137ac36a 100644
--- a/python/kvikio/tests/conftest.py
+++ b/python/kvikio/tests/conftest.py
@@ -1,9 +1,10 @@
-# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
 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_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_v3.py b/python/kvikio/tests/test_zarr_v3.py
index 9eaa2c5330..a66d711bbc 100644
--- a/python/kvikio/tests/test_zarr_v3.py
+++ b/python/kvikio/tests/test_zarr_v3.py
@@ -7,7 +7,7 @@
 import pytest
 import zarr.core.buffer
 import zarr.storage
-from zarr.core.buffer.cpu import Buffer
+from zarr.core.buffer import Buffer
 from zarr.testing.store import StoreTests
 
 import kvikio.zarr_v3
@@ -58,10 +58,12 @@ def store_kwargs(self, tmpdir: pathlib.Path) -> dict[str, str]:
 
     @pytest.fixture
     async def store(self, store_kwargs: dict[str, str]) -> kvikio.zarr_v3.GDSStore:
-        return self.store_cls(**store_kwargs)
+        # 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_v3.GDSStore:
-        return self.store_cls(**store_kwargs)
+        # ignore Argument 1 has incompatible type "**Dict[str, str]"; expected "bool"
+        return self.store_cls(**store_kwargs)  # type: ignore[arg-type]

From 47b9bd21949af073e7701fc872510afdd97dcf44 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Fri, 28 Feb 2025 20:29:32 +0000
Subject: [PATCH 09/21] Refactor

Moves zarr.py to _zarr_v2. Moves zarr_v3 to _zarr_v3.
---
 python/kvikio/kvikio/zarr/__init__.py         | 10 ++++++++
 .../{zarr.py => zarr/_zarr_python_2.py}       |  6 ++---
 .../{zarr_v3.py => zarr/_zarr_python_3.py}    |  0
 python/kvikio/pyproject.toml                  |  2 +-
 python/kvikio/tests/test_zarr_v3.py           | 23 +++++++++++--------
 5 files changed, 28 insertions(+), 13 deletions(-)
 create mode 100644 python/kvikio/kvikio/zarr/__init__.py
 rename python/kvikio/kvikio/{zarr.py => zarr/_zarr_python_2.py} (98%)
 rename python/kvikio/kvikio/{zarr_v3.py => zarr/_zarr_python_3.py} (100%)

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 98%
rename from python/kvikio/kvikio/zarr.py
rename to python/kvikio/kvikio/zarr/_zarr_python_2.py
index 31dcd6b943..6b494a4eef 100644
--- a/python/kvikio/kvikio/zarr.py
+++ b/python/kvikio/kvikio/zarr/_zarr_python_2.py
@@ -364,7 +364,7 @@ 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,
@@ -385,7 +385,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(),
@@ -415,7 +415,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_v3.py b/python/kvikio/kvikio/zarr/_zarr_python_3.py
similarity index 100%
rename from python/kvikio/kvikio/zarr_v3.py
rename to python/kvikio/kvikio/zarr/_zarr_python_3.py
diff --git a/python/kvikio/pyproject.toml b/python/kvikio/pyproject.toml
index ddbdbde9c7..bd43bf0038 100644
--- a/python/kvikio/pyproject.toml
+++ b/python/kvikio/pyproject.toml
@@ -112,7 +112,7 @@ ignore_missing_imports = true
 exclude = [
     # we type check against zarr-python 3.x
     # and ignore modules using 2.x
-    "python/kvikio/kvikio/zarr.py",
+    "python/kvikio/kvikio/zarr/_zarr_python_2.py",
     "python/kvikio/tests/test_nvcomp_codec.py",
 ]
 
diff --git a/python/kvikio/tests/test_zarr_v3.py b/python/kvikio/tests/test_zarr_v3.py
index a66d711bbc..a3016b5eed 100644
--- a/python/kvikio/tests/test_zarr_v3.py
+++ b/python/kvikio/tests/test_zarr_v3.py
@@ -7,10 +7,12 @@
 import pytest
 import zarr.core.buffer
 import zarr.storage
-from zarr.core.buffer import Buffer
+from zarr.core.buffer.gpu import Buffer
 from zarr.testing.store import StoreTests
 
-import kvikio.zarr_v3
+import kvikio.zarr
+
+pytest.importorskip("zarr", minversion="3.0.0")
 
 
 @pytest.mark.asyncio
@@ -21,7 +23,7 @@ async def test_basic(tmp_path: pathlib.Path) -> None:
     arr[5:] = 1
 
     assert await src.exists("a/zarr.json")
-    store = kvikio.zarr_v3.GDSStore(tmp_path, read_only=True)
+    store = kvikio.zarr.GDSStore(tmp_path, read_only=True)
     assert await store.exists("a/zarr.json")
 
     # regular works
@@ -36,15 +38,18 @@ async def test_basic(tmp_path: pathlib.Path) -> None:
         cp.testing.assert_array_equal(result, expected)
 
 
-class TestKvikIOStore(StoreTests[kvikio.zarr_v3.GDSStore, Buffer]):
-    store_cls = kvikio.zarr_v3.GDSStore
+class TestKvikIOStore(StoreTests[kvikio.zarr.GDSStore, Buffer]):
+    store_cls = kvikio.zarr.GDSStore
     buffer_cls = Buffer
 
-    async def get(self, store: kvikio.zarr_v3.GDSStore, key: str) -> 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_v3.GDSStore, key: str, value: Buffer
+        self,
+        store: kvikio.zarr.GDSStore,
+        key: str,
+        value: Buffer,  # type: ignore[override]
     ) -> None:
         parent = (store.root / key).parent
         if not parent.exists():
@@ -57,13 +62,13 @@ def store_kwargs(self, tmpdir: pathlib.Path) -> dict[str, str]:
         return kwargs
 
     @pytest.fixture
-    async def store(self, store_kwargs: dict[str, str]) -> kvikio.zarr_v3.GDSStore:
+    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_v3.GDSStore:
+    ) -> kvikio.zarr.GDSStore:
         # ignore Argument 1 has incompatible type "**Dict[str, str]"; expected "bool"
         return self.store_cls(**store_kwargs)  # type: ignore[arg-type]

From 8bccefb585e08b0be6a5889720cb72a4c1915ecb Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Fri, 28 Feb 2025 20:38:04 +0000
Subject: [PATCH 10/21] docs

---
 docs/source/zarr.rst | 30 +++++++++++++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

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

From e2e7bcad1d4144678448393d00750fce849cc25d Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Fri, 28 Feb 2025 21:04:32 +0000
Subject: [PATCH 11/21] cleanup

---
 python/kvikio/kvikio/zarr/_zarr_python_2.py | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/python/kvikio/kvikio/zarr/_zarr_python_2.py b/python/kvikio/kvikio/zarr/_zarr_python_2.py
index 6b494a4eef..edd7f275f5 100644
--- a/python/kvikio/kvikio/zarr/_zarr_python_2.py
+++ b/python/kvikio/kvikio/zarr/_zarr_python_2.py
@@ -5,7 +5,6 @@
 import contextlib
 import os
 import os.path
-import warnings
 from abc import abstractmethod
 from typing import Any, Literal, Mapping, Optional, Sequence, Union
 
@@ -18,18 +17,6 @@
 import zarr.creation
 import zarr.errors
 import zarr.storage
-
-try:
-    import zarr.util
-except ImportError:
-    message = (
-        "Failed to import 'zarr.util' while importing 'kvikio.zarr'.  "
-        "If you have zarr-python>=3.0, try importing 'kvikio.zarr_v3' instead."
-    )
-    warnings.warn(message, UserWarning)
-    raise
-
-
 from numcodecs.abc import Codec
 from numcodecs.compat import ensure_contiguous_ndarray_like
 from numcodecs.registry import register_codec

From 7b88b3890ae29c11cecb817f9f2b5fede960817b Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Fri, 28 Feb 2025 21:07:51 +0000
Subject: [PATCH 12/21] cleanup

---
 python/kvikio/kvikio/zarr/_zarr_python_3.py | 40 +++------------------
 1 file changed, 5 insertions(+), 35 deletions(-)

diff --git a/python/kvikio/kvikio/zarr/_zarr_python_3.py b/python/kvikio/kvikio/zarr/_zarr_python_3.py
index f2bf113309..25766cf3b3 100644
--- a/python/kvikio/kvikio/zarr/_zarr_python_3.py
+++ b/python/kvikio/kvikio/zarr/_zarr_python_3.py
@@ -2,7 +2,6 @@
 # See file LICENSE for terms.
 
 import asyncio
-import contextlib
 import os
 from pathlib import Path
 
@@ -18,19 +17,6 @@
 
 import kvikio
 
-try:
-    import nvtx
-except ImportError:
-    HAS_NVTX = False
-else:
-    HAS_NVTX = True
-
-
-if HAS_NVTX:
-    annotate = nvtx.annotate
-else:
-    annotate = contextlib.nullcontext
-
 
 def _get(
     path: Path, prototype: BufferPrototype, byte_range: ByteRequest | None
@@ -42,13 +28,13 @@ def _get(
         case None:
             nbytes = file_size
             file_offset = 0
-        case OffsetByteRequest():  # type: ignore[misc]
+        case OffsetByteRequest():
             nbytes = max(0, file_size - byte_range.offset)
             file_offset = byte_range.offset
-        case RangeByteRequest():  # type: ignore[misc]
+        case RangeByteRequest():
             nbytes = byte_range.end - byte_range.start
             file_offset = byte_range.start
-        case SuffixByteRequest():  # type: ignore[misc]
+        case SuffixByteRequest():
             nbytes = byte_range.suffix
             file_offset = max(0, file_size - byte_range.suffix)
         case _:
@@ -86,12 +72,9 @@ def _put(
     path.parent.mkdir(parents=True, exist_ok=True)
     if start is not None:
         with kvikio.CuFile(path, "r+b") as f:
-            # TODO: seems like this isn't tested
-            # https://github.com/zarr-developers/zarr-python/issues/2859
             f.write(value.as_array_like(), file_offset=start)
         return None
     else:
-        # view = memoryview(value.as_numpy_array().tobytes())
         buf = value.as_array_like()
         if exclusive:
             if path.exists():
@@ -111,7 +94,6 @@ async def get(
         prototype: BufferPrototype | None = None,
         byte_range: ByteRequest | None = None,
     ) -> Buffer | None:
-        # docstring inherited
         if prototype is None:
             prototype = default_buffer_prototype()
         if not self._is_open:
@@ -119,18 +101,12 @@ async def get(
         assert isinstance(key, str)
         path = self.root / key
 
-        if HAS_NVTX:
-            kwargs = {"message": "kvikio.zarr.get", "domain": "Zarr"}
-        else:
-            kwargs = {}
         try:
-            with annotate(**kwargs):
-                return await asyncio.to_thread(_get, path, prototype, byte_range)
+            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:
-        # docstring inherited
         return await self._set(key, value)
 
     async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None:
@@ -145,10 +121,4 @@ async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None:
             )
         path = self.root / key
 
-        if HAS_NVTX:
-            kwargs = {"message": "kvikio.zarr.set", "domain": "Zarr"}
-        else:
-            kwargs = {}
-
-        with annotate(**kwargs):
-            await asyncio.to_thread(_put, path, value, start=None, exclusive=exclusive)
+        await asyncio.to_thread(_put, path, value, start=None, exclusive=exclusive)

From 5cbe53e80bf2f278e65f557a13445b0a0b096939 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Fri, 28 Feb 2025 21:35:38 +0000
Subject: [PATCH 13/21] Fixed v3 skip

---
 python/kvikio/tests/test_zarr_v3.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/python/kvikio/tests/test_zarr_v3.py b/python/kvikio/tests/test_zarr_v3.py
index a3016b5eed..9e694655b3 100644
--- a/python/kvikio/tests/test_zarr_v3.py
+++ b/python/kvikio/tests/test_zarr_v3.py
@@ -5,16 +5,18 @@
 
 import cupy as cp
 import pytest
-import zarr.core.buffer
-import zarr.storage
-from zarr.core.buffer.gpu import Buffer
-from zarr.testing.store import StoreTests
 
 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)

From 55e2fa89c1d4de0b5d6d3d9830869aff364754af Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Mon, 3 Mar 2025 14:17:25 +0000
Subject: [PATCH 14/21] zarr dependency

---
 conda/environments/all_cuda-118_arch-aarch64.yaml | 2 +-
 conda/environments/all_cuda-118_arch-x86_64.yaml  | 2 +-
 conda/environments/all_cuda-128_arch-aarch64.yaml | 2 +-
 conda/environments/all_cuda-128_arch-x86_64.yaml  | 2 +-
 conda/recipes/kvikio/meta.yaml                    | 2 +-
 dependencies.yaml                                 | 2 +-
 python/kvikio/pyproject.toml                      | 2 +-
 7 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/conda/environments/all_cuda-118_arch-aarch64.yaml b/conda/environments/all_cuda-118_arch-aarch64.yaml
index 715ac12452..3327871d53 100644
--- a/conda/environments/all_cuda-118_arch-aarch64.yaml
+++ b/conda/environments/all_cuda-118_arch-aarch64.yaml
@@ -40,5 +40,5 @@ dependencies:
 - sphinx-click
 - sphinx_rtd_theme
 - sysroot_linux-aarch64=2.28
-- zarr>=2.0.0
+- zarr>=2.0.0,<4.0.0a0
 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 faf463a106..5ad9c646c9 100644
--- a/conda/environments/all_cuda-118_arch-x86_64.yaml
+++ b/conda/environments/all_cuda-118_arch-x86_64.yaml
@@ -42,5 +42,5 @@ dependencies:
 - sphinx-click
 - sphinx_rtd_theme
 - sysroot_linux-64=2.28
-- zarr>=2.0.0
+- zarr>=2.0.0,<4.0.0a0
 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 896a9b723f..e4444b955f 100644
--- a/conda/environments/all_cuda-128_arch-aarch64.yaml
+++ b/conda/environments/all_cuda-128_arch-aarch64.yaml
@@ -40,5 +40,5 @@ dependencies:
 - sphinx-click
 - sphinx_rtd_theme
 - sysroot_linux-aarch64=2.28
-- zarr>=2.0.0
+- zarr>=2.0.0,<4.0.0a0
 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 dc0d518486..627d8cfd18 100644
--- a/conda/environments/all_cuda-128_arch-x86_64.yaml
+++ b/conda/environments/all_cuda-128_arch-x86_64.yaml
@@ -40,5 +40,5 @@ dependencies:
 - sphinx-click
 - sphinx_rtd_theme
 - sysroot_linux-64=2.28
-- zarr>=2.0.0
+- zarr>=2.0.0,<4.0.0a0
 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 d706318409..c60ed417c4 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
+          - zarr>=2.0.0,<4.0.0a0
           # See https://github.com/zarr-developers/numcodecs/pull/475
           - numcodecs !=0.12.0
           - packaging
diff --git a/python/kvikio/pyproject.toml b/python/kvikio/pyproject.toml
index c70aec6418..3df8170cc3 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",
+    "zarr>=2.0.0,<4.0.0a0",
 ] # 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",

From 5eeb0556bae307a5c9a238dbc64c70ca5ee10eab Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Mon, 3 Mar 2025 15:26:47 +0000
Subject: [PATCH 15/21] Skip fixes

---
 python/kvikio/kvikio/zarr/_zarr_python_3.py |  7 +++++++
 python/kvikio/tests/test_benchmarks.py      | 15 +++++++++++++++
 python/kvikio/tests/test_zarr.py            |  9 ++++++++-
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/python/kvikio/kvikio/zarr/_zarr_python_3.py b/python/kvikio/kvikio/zarr/_zarr_python_3.py
index 25766cf3b3..3f65c8c5f8 100644
--- a/python/kvikio/kvikio/zarr/_zarr_python_3.py
+++ b/python/kvikio/kvikio/zarr/_zarr_python_3.py
@@ -122,3 +122,10 @@ async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None:
         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/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_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):

From c2a59bfadbbe0c895de5e1855ed564156843260d Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Mon, 3 Mar 2025 15:28:37 +0000
Subject: [PATCH 16/21] Note LocalFile similarity

---
 python/kvikio/kvikio/zarr/_zarr_python_3.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/python/kvikio/kvikio/zarr/_zarr_python_3.py b/python/kvikio/kvikio/zarr/_zarr_python_3.py
index 3f65c8c5f8..69d008ecf9 100644
--- a/python/kvikio/kvikio/zarr/_zarr_python_3.py
+++ b/python/kvikio/kvikio/zarr/_zarr_python_3.py
@@ -17,6 +17,10 @@
 
 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

From e8ec00dd4d6741d1aec167d60db48ac3043fda6a Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Mon, 3 Mar 2025 15:30:28 +0000
Subject: [PATCH 17/21] Add note about StoreTests

---
 python/kvikio/tests/test_zarr_v3.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/python/kvikio/tests/test_zarr_v3.py b/python/kvikio/tests/test_zarr_v3.py
index 9e694655b3..e44de36d22 100644
--- a/python/kvikio/tests/test_zarr_v3.py
+++ b/python/kvikio/tests/test_zarr_v3.py
@@ -40,6 +40,11 @@ async def test_basic(tmp_path: pathlib.Path) -> None:
         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

From 75a5cb5514175f35595322b8e760627c71cbd5fb Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Mon, 3 Mar 2025 08:41:01 -0800
Subject: [PATCH 18/21] Avoid zarr alpha dep

In https://github.com/rapidsai/kvikio/actions/runs/13633004880/job/38105532469?pr=646#step:10:88
we pick up an alpha release of zarr-python 3. There aren't any official releases
of zarr-python 3.x that support python 3.10, but some alphas did.

I'm not 100% why pip decided to pick that alpha release, but it's possible
that the upper bound on zarr-python<4.0.0a0 allowed pip to select
an alpha release. By removing that `.a0` I'm hopeful that pip
will only consider official releases, and pick the latest 2.x release
on python 3.10
---
 python/kvikio/pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/kvikio/pyproject.toml b/python/kvikio/pyproject.toml
index 3df8170cc3..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,<4.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",

From 6bd40796b9ed784184af28966f171188efca49d7 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Mon, 3 Mar 2025 09:27:24 -0800
Subject: [PATCH 19/21] fix in dependencies.yaml

---
 conda/environments/all_cuda-118_arch-aarch64.yaml | 2 +-
 conda/environments/all_cuda-118_arch-x86_64.yaml  | 2 +-
 conda/environments/all_cuda-128_arch-aarch64.yaml | 2 +-
 conda/environments/all_cuda-128_arch-x86_64.yaml  | 2 +-
 dependencies.yaml                                 | 2 +-
 5 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/conda/environments/all_cuda-118_arch-aarch64.yaml b/conda/environments/all_cuda-118_arch-aarch64.yaml
index 3327871d53..dea7a9f2e1 100644
--- a/conda/environments/all_cuda-118_arch-aarch64.yaml
+++ b/conda/environments/all_cuda-118_arch-aarch64.yaml
@@ -40,5 +40,5 @@ dependencies:
 - sphinx-click
 - sphinx_rtd_theme
 - sysroot_linux-aarch64=2.28
-- zarr>=2.0.0,<4.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 5ad9c646c9..c21a7d9a5d 100644
--- a/conda/environments/all_cuda-118_arch-x86_64.yaml
+++ b/conda/environments/all_cuda-118_arch-x86_64.yaml
@@ -42,5 +42,5 @@ dependencies:
 - sphinx-click
 - sphinx_rtd_theme
 - sysroot_linux-64=2.28
-- zarr>=2.0.0,<4.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 e4444b955f..7834ffcab5 100644
--- a/conda/environments/all_cuda-128_arch-aarch64.yaml
+++ b/conda/environments/all_cuda-128_arch-aarch64.yaml
@@ -40,5 +40,5 @@ dependencies:
 - sphinx-click
 - sphinx_rtd_theme
 - sysroot_linux-aarch64=2.28
-- zarr>=2.0.0,<4.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 627d8cfd18..e38e3729a3 100644
--- a/conda/environments/all_cuda-128_arch-x86_64.yaml
+++ b/conda/environments/all_cuda-128_arch-x86_64.yaml
@@ -40,5 +40,5 @@ dependencies:
 - sphinx-click
 - sphinx_rtd_theme
 - sysroot_linux-64=2.28
-- zarr>=2.0.0,<4.0.0a0
+- zarr>=2.0.0,<4.0.0
 name: all_cuda-128_arch-x86_64
diff --git a/dependencies.yaml b/dependencies.yaml
index c60ed417c4..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,<4.0.0a0
+          - zarr>=2.0.0,<4.0.0
           # See https://github.com/zarr-developers/numcodecs/pull/475
           - numcodecs !=0.12.0
           - packaging

From 1ca12e2e9a1e624d46bc36b387cd60e30896b8f9 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Mon, 3 Mar 2025 19:00:47 +0000
Subject: [PATCH 20/21] Added example skip

---
 python/kvikio/tests/test_examples.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

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

From 172eeb7090f75cd3ed47c6ffd51328866fcb7b3e Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Mon, 3 Mar 2025 19:05:10 +0000
Subject: [PATCH 21/21] compat in benchmarks

---
 python/kvikio/kvikio/benchmarks/single_node_io.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

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