diff --git a/python/benchmarks/single-node-io.py b/python/benchmarks/single-node-io.py index f8255bec46..4b56090f77 100644 --- 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 index 772aa3f92b..21ebb2f737 100644 --- 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: " - "" - ), - ) + 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 index 6907e3be08..757c2bd22a 100644 --- 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: " - "" - ), - ) - 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: " - "" - ), - ) - g = zarr.open_group(store, meta_array=cupy.empty(())) g.ones("data", shape=(10, 11), dtype=int, compressor=None) a = g["data"]