Use Zarr v2.13.0a2

which include zarr-developers/zarr-python#934
madsbk · Sep 9, 2022 · 91cb3e8 · 91cb3e8
1 parent f8f4063
commit 91cb3e8
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 40 deletions.
diff --git a/python/benchmarks/single-node-io.py b/python/benchmarks/single-node-io.py
@@ -3,7 +3,6 @@
 
 import argparse
 import contextlib
-import functools
 import os
 import os.path
 import pathlib
@@ -193,30 +192,30 @@ def run_posix(args):
     return read_time, write_time
 
 
-def run_zarr(store_type, args):
+def run_zarr(args):
     """Use the Zarr API"""
 
     import zarr
-    import zarr.cupy
 
     import kvikio.zarr
 
+    dir_path = args.dir / "zarr"
+
+    if not hasattr(zarr.Array, "meta_array"):
+        RuntimeError("Requires Zarr v2.13.0+ for CuPy support")
+
     a = cupy.arange(args.nbytes, dtype="uint8")
 
-    # Retrieve the store and compressor to use based on `store_type`
-    shutil.rmtree(str(args.dir / store_type), ignore_errors=True)
-    store, compressor = {
-        "gds": (kvikio.zarr.GDSStore(args.dir / store_type), None),
-        "posix": (
-            zarr.DirectoryStore(args.dir / store_type),
-            zarr.cupy.CuPyCPUCompressor(),
-        ),
-    }[store_type]
+    shutil.rmtree(str(dir_path), ignore_errors=True)
 
     # Write
     t0 = clock()
     z = zarr.array(
-        a, chunks=False, compressor=compressor, store=store, meta_array=cupy.empty(())
+        a,
+        chunks=False,
+        compressor=None,
+        store=kvikio.zarr.GDSStore(dir_path),
+        meta_array=cupy.empty(()),
     )
     write_time = clock() - t0
 
@@ -231,8 +230,7 @@ def run_zarr(store_type, args):
 
 API = {
     "cufile": run_cufile,
-    "zarr-gds": functools.partial(run_zarr, "gds"),
-    "zarr-posix": functools.partial(run_zarr, "posix"),
+    "zarr": run_zarr,
     "posix": run_posix,
     "cufile-mfma": run_cufile_multiple_files_multiple_arrays,
     "cufile-mf": run_cufile_multiple_files,

diff --git a/python/tests/test_benchmarks.py b/python/tests/test_benchmarks.py
@@ -21,21 +21,16 @@
         "cufile-mfma",
         "cufile-mf",
         "cufile-ma",
-        "zarr-gds",
-        "zarr-posix",
+        "zarr",
     ],
 )
 def test_single_node_io(run_cmd, tmp_path, api):
     """Test benchmarks/single-node-io.py"""
 
     if "zarr" in api:
-        pytest.importorskip(
-            "zarr.cupy",
-            reason=(
-                "To use Zarr arrays with GDS directly, Zarr needs CuPy support: "
-                "<https://github.com/zarr-developers/zarr-python/pull/934>"
-            ),
-        )
+        zarr = pytest.importorskip("zarr")
+        if not hasattr(zarr.Array, "meta_array"):
+            pytest.skip("Requires Zarr v2.13.0+ for CuPy support")
 
     retcode = run_cmd(
         cmd=[

diff --git a/python/tests/test_zarr.py b/python/tests/test_zarr.py
@@ -9,6 +9,12 @@
 GDSStore = pytest.importorskip("kvikio.zarr").GDSStore
 
 
+cupy_support = pytest.mark.skipif(
+    not hasattr(zarr.Array, "meta_array"),
+    reason="Requires Zarr v2.13.0+ for CuPy support",
+)
+
+
 @pytest.fixture
 def store(tmp_path):
     """Fixture that creates a GDS Store"""
@@ -29,17 +35,10 @@ def test_direct_store_access(store, array_type):
     cupy.testing.assert_array_equal(a, b)
 
 
+@cupy_support
 def test_array(store):
     """Test Zarr array"""
 
-    pytest.importorskip(
-        "zarr.cupy",
-        reason=(
-            "To use Zarr arrays with GDS directly, Zarr needs CuPy support: "
-            "<https://github.com/zarr-developers/zarr-python/pull/934>"
-        ),
-    )
-
     a = cupy.arange(100)
     z = zarr.array(
         a, chunks=10, compressor=None, store=store, meta_array=cupy.empty(())
@@ -50,17 +49,10 @@ def test_array(store):
     cupy.testing.assert_array_equal(a, z[:])
 
 
+@cupy_support
 def test_group(store):
     """Test Zarr group"""
 
-    pytest.importorskip(
-        "zarr.cupy",
-        reason=(
-            "To use Zarr arrays with GDS directly, Zarr needs CuPy support: "
-            "<https://github.com/zarr-developers/zarr-python/pull/934>"
-        ),
-    )
-
     g = zarr.open_group(store, meta_array=cupy.empty(()))
     g.ones("data", shape=(10, 11), dtype=int, compressor=None)
     a = g["data"]