Kvikio backend entrypoint with Zarr v3 #70

weiji14 · 2025-03-11T01:59:57Z

Branched off the work at #10, this PR registers a kvikio backend that allows reading data from Zarr (v3) stores directly into GPU memory using NVIDIA GPU Direct Storage (GDS) without going through host (CPU) memory.

Preview demo at https://cupy-xarray--70.org.readthedocs.build/70/source/kvikio.html

Requires some patches from un-released versions:

kvikio>=25.04.00a (Support zarr-python 3.x rapidsai/kvikio#646)
~~xarray>=2025.1.3.dev22+g0184702f (Generalize lazy backend indexing a little more pydata/xarray#10078)~~ Now in xarray=2025.03.0

Also requires:

zarr >=3.0.3 (containing Update and document GPU buffer handling zarr-developers/zarr-python#2751)

TODO:

Wait for stable release of kvikio>=25.04 and xarray>=2025.3.0
Figure out where to insert zarr.config.enable_gpu() into the backend to read into cupy.ndarray by default?
Write more unit tests
Cleanup API and tutorial documentation

Limitations:

Assumes that Zarr store is uncompressed. GPU-based decompression via nvCOMP is not supported yet (wait for Zstd Codec on the GPU zarr-developers/zarr-python#2863)

References:

for more information, see https://pre-commit.ci

* main: Min xarray >= 0.19.0 (#25) Fix broken dask_array_type import (#24)

for more information, see https://pre-commit.ci

* upstream/main: Documentation Updates 📖 (#35) [pre-commit.ci] pre-commit autoupdate (#37) [pre-commit.ci] pre-commit autoupdate (#34) [pre-commit.ci] pre-commit autoupdate (#32) Update .pre-commit-config.yaml (#28) Expand installation doc (#27)

Allow it to be rendered under the User Guide section.

Will need it for the kvikio.zarr docs later.

Create new section in the API documentation page for the kvikIO engine. Added more docstrings to the kvikio.py file, and fixed some imports so things render nicely on the API page. Also added an intersphinx link to the kvikio docs at https://docs.rapids.ai/api/kvikio/stable.

Fixes error like `TypeError: ZarrArrayWrapper.__init__() takes 2 positional arguments but 3 were given`.

for more information, see https://pre-commit.ci

Fix improper merge handling on d684dad

Fix `TypeError: Implicit conversion to a NumPy array is not allowed. Please use `.get()` to construct a NumPy array explicitly` on https://github.com/pydata/xarray/blob/v2024.11.0/xarray/core/indexing.py#L578

for more information, see https://pre-commit.ci

Need patches from rapidsai/kvikio#646 and zarr-developers/zarr-python#2751.

Can directly rely on upstream xarray's ZarrStore.open_store_variable method since Zarr v3 compatibility was added in pydata/xarray#9552.

UserWarning: The `compressor` argument is deprecated. Use `compressors` instead.

Only difference is to pass the Zarr store's root filepath to kvikio.zarr.GDSStore instead of xarray.backends.zarr.ZarrStore.

Dev version containing patch at pydata/xarray#10078 that fixes `TypeError: NumpyIndexingAdapter only wraps np.ndarray. Trying to wrap <class 'cupy.ndarray'>`.

Needed to return cupy.ndarray instead of numpy.ndarray objects. Should find a better place to put this `zarr.config.enable_gpu()` call later.

Rearranged some cells so the air-temperature.zarr store is created first. Added a couple of `zarr.config.enable_gpu()` statements in to ensure arrays are read to cupy.ndarray. Removed the flox section at the end.

weiji14 · 2025-03-11T02:06:22Z

cupy_xarray/tests/test_kvikio.py

+
+@pytest.mark.parametrize("indexer", [slice(None), slice(2, 4), 2, [2, 3, 5]])
+def test_lazy_indexing(indexer, store):
+    with zarr.config.enable_gpu(), xr.open_dataset(store, engine="kvikio") as ds:


Ideally, xr.open_dataset(store, engine="kvikio") should have zarr.config.enable_gpu() set already, so GPU-backed cupy arrays are returned when a user accesses the arrays. Unless we expect there to be users who wants to use the kvikio engine while returning CPU-backed numpy arrays. Default should probably be cupy though?

I'm not sure... The (v3) kvikio.zarr.GDSStore is written to accept either. I'm tempted to leave this up to the user (so make them configure Zarr correctly) rather than assuming they want to use GPU memory, but maybe engine="kvikio" is sufficient indication that they want stuff on the GPU.

Yeah, kvikio.zarr.GDSStore can return numpy (CPU) arrays too, and I guess there might be a case where someone wants nvCOMP decompression on the GPU (once zarr-developers/zarr-python#2863 is implemented), but have a numpy array returned in CPU memory? But I do feel that engine="kvikio" should default to putting things on GPU/cupy, while having a way to toggle it off as needed.

Looking at xr.open_dataset, I'm thinking if it might make sense to pass something into the backend_kwargs parameter to indicate the user's preference for GPU or CPU outputs. Is the from_array_kwargs or chunked_array_type parameters (experimental API) the intended use for this actually?

I am inclined to agree that kvikio.zarr.GDSStore should have the ability to not enable the GPU zarr configs even if someone wants to get data on the GPU since they can try and use a different BufferPrototype that return torch tensors instead of cupy arrays, for example.

But for ease of use, it might be better to default to zarr.config.enable_gpu() and that can be overridden by passing in a different context manager? So xr.open_dataset(store, engine="kvikio") will default to using cupy arrays, but xr.open_dataset(store, engine="kvikio", zarr_context=contextlib.nullcontext()) would return numpy arrays and xr.open_dataset(store, engine="kvikio", zarr_context=MyCustomTorchContext()) would return torch tensors?

I think this is similar to the discussion at zarr-developers/zarr-python#2473. That zarr_context parameter is exactly what I'm thinking of, I'm trying to find out if there's a way to implement this right now by setting/monkeypatching the default prototype in kvikio.zarr.GDSStore's get method here:

https://github.com/rapidsai/kvikio/blob/a4170fc098e80d339a42c5da9a605796eb864c9f/python/kvikio/kvikio/zarr/_zarr_python_3.py#L98-L102

Currently the prototype: BufferPrototype is set dynamically by calling zarr.core.buffer.core.default_buffer_prototype(). But I'm wondering if we can have a pre-determined default buffer prototype that is set when the GDSStore instance is created?

jacobtomlinson · 2025-03-11T10:07:08Z

cc @jakirkham @TomAugspurger @madsbk @ncclementi who may be interested in this

TomAugspurger

High-level question: is are these two functionally equivalent?

ds = xr.open_dataset(kvikio.zarr.GDSStore("data.zarr"), engine="zarr")
ds = xr.open_dataset("data.zarr", engine="kvikio")

Do people have a preference on what to recommend to users?

TomAugspurger · 2025-03-11T16:37:08Z

cupy_xarray/tests/test_kvikio.py

+
+@pytest.mark.parametrize("indexer", [slice(None), slice(2, 4), 2, [2, 3, 5]])
+def test_lazy_indexing(indexer, store):
+    with zarr.config.enable_gpu(), xr.open_dataset(store, engine="kvikio") as ds:


I'm not sure... The (v3) kvikio.zarr.GDSStore is written to accept either. I'm tempted to leave this up to the user (so make them configure Zarr correctly) rather than assuming they want to use GPU memory, but maybe engine="kvikio" is sufficient indication that they want stuff on the GPU.

weiji14 · 2025-03-11T18:12:24Z

cupy_xarray/kvikio.py

+        filename_or_obj = _normalize_path(filename_or_obj)
+        if not store:
+            store = ZarrStore.open_group(
+                store=kvikio.zarr.GDSStore(root=filename_or_obj),


High-level question: is are these two functionally equivalent?

ds = xr.open_dataset(kvikio.zarr.GDSStore("data.zarr"), engine="zarr") ds = xr.open_dataset("data.zarr", engine="kvikio")

Do people have a preference on what to recommend to users?

Yes, those two are functionally the same. The xr.open_dataset(kvikio.zarr.GDSStore("data.zarr"), engine="zarr") style works today (with xarray ~~patch~~ =2025.03.0 and kvikio=25.04.00a), while xr.open_dataset("data.zarr", engine="kvikio") is more just a convenience syntax. I was hoping to also sneak in the zarr.config.enable_gpu() in the backend somewhere to get cupy arrays by default (but we can discuss that in the https://github.com/xarray-contrib/cupy-xarray/pull/70/files#r1988251560 thread).

weiji14 · 2025-03-11T18:33:45Z

cupy_xarray/tests/test_kvikio.py

+
+@pytest.mark.parametrize("indexer", [slice(None), slice(2, 4), 2, [2, 3, 5]])
+def test_lazy_indexing(indexer, store):
+    with zarr.config.enable_gpu(), xr.open_dataset(store, engine="kvikio") as ds:


Yeah, kvikio.zarr.GDSStore can return numpy (CPU) arrays too, and I guess there might be a case where someone wants nvCOMP decompression on the GPU (once zarr-developers/zarr-python#2863 is implemented), but have a numpy array returned in CPU memory? But I do feel that engine="kvikio" should default to putting things on GPU/cupy, while having a way to toggle it off as needed.

Looking at xr.open_dataset, I'm thinking if it might make sense to pass something into the backend_kwargs parameter to indicate the user's preference for GPU or CPU outputs. Is the from_array_kwargs or chunked_array_type parameters (experimental API) the intended use for this actually?

Bumps [xarray](https://github.com/pydata/xarray) from 2025.1.3.dev22+g0184702f to 2025.03.0. - [Release notes](https://github.com/pydata/xarray/releases) - [Changelog](https://github.com/pydata/xarray/blob/main/HOW_TO_RELEASE.md) - [Commits](pydata/xarray@0184702...v2025.03.0) Use stable version of xarray, also with patch pydata/xarray#10081

jakirkham · 2025-03-25T23:04:22Z

cc @akshaysubr

Using functools.partial to override default buffer protocol to be GPU buffer instead of CPU buffer. Not quite working as expected, but hopefully gets a point across.

dcherian and others added 30 commits August 2, 2022 15:17

Add Kvikio backend entrypoint

9deadb7

Add demo notebook

aa2dc91

Update kvikio notebook

7fb4b94

Merge branch 'main' into kvikio-entrypoint

743fe7d

Merge branch 'main' into kvikio-entrypoint

5d501e4

[pre-commit.ci] auto fixes from pre-commit.com hooks

facf5f7

for more information, see https://pre-commit.ci

Update cupy_xarray/kvikio.py

f3f5189

Merge branch 'main' into kvikio-entrypoint

9c98d19

Merge branch 'main' into kvikio-entrypoint

dd8bc57

* main: Min xarray >= 0.19.0 (#25) Fix broken dask_array_type import (#24)

Add url, description.

d2da1e4

Working

b87c3c2

Updated notebook

87cb74e

[pre-commit.ci] auto fixes from pre-commit.com hooks

d7394ef

for more information, see https://pre-commit.ci

Add tests

ca0cf45

Merge branch 'main' into kvikio-entrypoint

97260d6

Move kvikio notebook under docs/source

5d27b26

Allow it to be rendered under the User Guide section.

Add zarr as a dependency in ci/doc.yml

85491d7

Will need it for the kvikio.zarr docs later.

Fix input argument into CupyZarrArrayWrapper

95efa18

Fixes error like `TypeError: ZarrArrayWrapper.__init__() takes 2 positional arguments but 3 were given`.

Merge branch 'main' into kvikio-entrypoint

d684dad

[pre-commit.ci] auto fixes from pre-commit.com hooks

ae2a7f1

for more information, see https://pre-commit.ci

Re-add kvikio backend entrypoint to pyproject.toml

15fbafd

Fix improper merge handling on d684dad

Fix C408 and E402

f3df115

Use get_duck_array instead of get_array

4e1857a

Fix `TypeError: Implicit conversion to a NumPy array is not allowed. Please use `.get()` to construct a NumPy array explicitly` on https://github.com/pydata/xarray/blob/v2024.11.0/xarray/core/indexing.py#L578

Fix SIM108 Use ternary operator

7345b61

[pre-commit.ci] auto fixes from pre-commit.com hooks

e2b410e

for more information, see https://pre-commit.ci

Install nightly version of kvikio=25.04.00a and zarr>=3.0.5

7dd78e9

Need patches from rapidsai/kvikio#646 and zarr-developers/zarr-python#2751.

Remove custom open_store_variable method from GDSZarrStore class

cb77678

Can directly rely on upstream xarray's ZarrStore.open_store_variable method since Zarr v3 compatibility was added in pydata/xarray#9552.

Fix UserWarning compressor -> compressors

0262151

UserWarning: The `compressor` argument is deprecated. Use `compressors` instead.

weiji14 added 4 commits March 11, 2025 13:21

Reuse logic from xarray.backends.zarr.ZarrStore.open_group

e26ed24

Only difference is to pass the Zarr store's root filepath to kvikio.zarr.GDSStore instead of xarray.backends.zarr.ZarrStore.

Install xarray=2025.1.3.dev22+g0184702f

f185b44

Dev version containing patch at pydata/xarray#10078 that fixes `TypeError: NumpyIndexingAdapter only wraps np.ndarray. Trying to wrap <class 'cupy.ndarray'>`.

Add zarr.config.enable_gpu() context manager to test_lazy_indexing

1a52ce5

Needed to return cupy.ndarray instead of numpy.ndarray objects. Should find a better place to put this `zarr.config.enable_gpu()` call later.

Refresh kvikIO demo notebook

789a9b6

Rearranged some cells so the air-temperature.zarr store is created first. Added a couple of `zarr.config.enable_gpu()` statements in to ensure arrays are read to cupy.ndarray. Removed the flox section at the end.

weiji14 added the enhancement New feature or request label Mar 11, 2025

weiji14 added this to the 0.2.0 milestone Mar 11, 2025

weiji14 self-assigned this Mar 11, 2025

weiji14 commented Mar 11, 2025

View reviewed changes

Merge branch 'main' into kvikio-backend

3894d29

TomAugspurger reviewed Mar 11, 2025

View reviewed changes

weiji14 commented Mar 11, 2025

View reviewed changes

Try overriding default prototype to be GPU buffer

1e205ec

Using functools.partial to override default buffer protocol to be GPU buffer instead of CPU buffer. Not quite working as expected, but hopefully gets a point across.

weiji14 mentioned this pull request Apr 2, 2025

Should we statically associate Store instances with Buffer types? zarr-developers/zarr-python#2473

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Kvikio backend entrypoint with Zarr v3 #70

Kvikio backend entrypoint with Zarr v3 #70

weiji14 commented Mar 11, 2025 •

edited

Loading

weiji14 Mar 11, 2025

TomAugspurger Mar 11, 2025

weiji14 Mar 11, 2025

akshaysubr Mar 31, 2025

weiji14 Apr 1, 2025 •

edited

Loading

jacobtomlinson commented Mar 11, 2025

TomAugspurger left a comment

TomAugspurger Mar 11, 2025

weiji14 Mar 11, 2025 •

edited

Loading

weiji14 Mar 11, 2025

jakirkham commented Mar 25, 2025

Kvikio backend entrypoint with Zarr v3 #70

Are you sure you want to change the base?

Kvikio backend entrypoint with Zarr v3 #70

Conversation

weiji14 commented Mar 11, 2025 • edited Loading

weiji14 Mar 11, 2025

Choose a reason for hiding this comment

TomAugspurger Mar 11, 2025

Choose a reason for hiding this comment

weiji14 Mar 11, 2025

Choose a reason for hiding this comment

akshaysubr Mar 31, 2025

Choose a reason for hiding this comment

weiji14 Apr 1, 2025 • edited Loading

Choose a reason for hiding this comment

jacobtomlinson commented Mar 11, 2025

TomAugspurger left a comment

Choose a reason for hiding this comment

TomAugspurger Mar 11, 2025

Choose a reason for hiding this comment

weiji14 Mar 11, 2025 • edited Loading

Choose a reason for hiding this comment

weiji14 Mar 11, 2025

Choose a reason for hiding this comment

jakirkham commented Mar 25, 2025

weiji14 commented Mar 11, 2025 •

edited

Loading

weiji14 Apr 1, 2025 •

edited

Loading

weiji14 Mar 11, 2025 •

edited

Loading