Support zarr-python 3.x #646

TomAugspurger · 2025-02-25T03:40:51Z

This adds support for zarr-python 3.x to kvikio.

zarr-python 3.x made large changes to the store API and the way arrays (and their location) are specified, but minimal changes to the main Group and Array APIs. So in this context "support" for zarr-python 3.x means just a new Store class that uses kvikio for file I/O.

In particular, we don't need a replacement for anything else like open_cupy_array or compressors: these are built into zarr-python 3.x now.

Closes #582.

This implements a zarr-python 3.x compatible GDSStore.

copy-pr-bot · 2025-02-25T03:40:56Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

TomAugspurger · 2025-02-25T15:54:55Z

I added a few notes / questions to the original post about how we want to handle backwards compatibility and packaging, if anyone has thoughts there.

Moves zarr.py to _zarr_v2. Moves zarr_v3 to _zarr_v3.

TomAugspurger · 2025-02-28T21:03:26Z

I've updated this to make the kvikio.zarr import depend on the version of zarr installed at runtime. https://github.com/rapidsai/kvikio/pull/646/files#diff-094f0da1167f8d0545aa3c38e2cdda455a6da2862cc1707f6e78c06ce8466dd1R10 has the documentation

If the user has zarr-python 2.x installed, then there's no change.

If you have zarr-python 3.x, then your code might continue to work if you're just using kvikio.zarr.GDSStore and passing that to parts of zarr-python's API that haven't changed. But most of that is outside of our control here. If you were using something like kvikio.zarr.open_cupy_array then that will break, but mostly because it isn't necessary anymore.

madsbk

Looks great. One question, can we get CI to run both tests?
Maybe pip install Zarr v3 in run_pytests.sh or something?

dependencies.yaml

python/kvikio/kvikio/zarr/_zarr_python_3.py

wence- · 2025-03-03T12:55:39Z

It looks like zarr-python 3 is python >= 3.11, whereas I think we're >= 3.10. That might complicate things a bit.

Per their migration guidelines, v2 is only likely to be supported until the summer.

Given that, I would be OK with only supporting v3 and just not supporting this feature with py3.10.

madsbk · 2025-03-03T12:56:42Z

Given that, I would be OK with only supporting v3 and just not supporting this feature with py3.10.

Is also fine with me.

TomAugspurger · 2025-03-03T14:20:33Z

I'm looking into the versions resolved during CI now. Currently, all jobs are using zarr-python 2.x

The 11.8 / 3.10 job picked up zarr 2.x
- list: https://github.com/rapidsai/kvikio/actions/runs/13596937483/job/38016439897#step:9:493
- test: https://github.com/rapidsai/kvikio/actions/runs/13596937483/job/38016439897#step:9:1938
- The 12.8 / 3.12 job also picked up zarr 2.x for some reason
  - list: https://github.com/rapidsai/kvikio/actions/runs/13596937483/job/38016440081#step:9:512

Assuming we can get the 3.12 job using zarr 3.x, I think we'll be in a good spot for CI / testing.

wence- · 2025-03-03T12:57:43Z

python/kvikio/kvikio/zarr/__init__.py

+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


FWIW, if we're doing this, I would prefer that we have:

import kvikio.zarr # old zarr-python 2 interface, for backwards compat, warns about deprecation import kvikio.zarr2 # move the zarr-python 2 code here import kvikio.zarr3 # new zarr-python 3 interface

This way, a user's scripts don't suddenly break because they run on them on a machine with a different environment.

We have two strategies here: try to preserve source-code compatibility or not.

So under option 1, we have the potential for some user's code working with either zarr-python 2.x or 3.x, but if and only if the user limits themselves to kvikio.zarr.GDSStore and the public Zarr API that's compatible between v2 and v3 (essentially, passing in store=store to most of the group and array APIs).

If you're using any of the other features of kvikio.zarr like open_cupy_array or any of the compressors then that will break if you update to zarr-python 3.x, because we don't attempt to provide source compatibility there. And if you're using one of the zarr-python APIs that changed between zarr-python 2.x and 3.x, then stuff will break but that's outside of kvikio's hands.

We also get the (IMO) nicer name of kvikio.zarr. I dislike kvikio.zarr3 since it's not immediately clear whether the 3 refers to zarr format version or the zarr-python version (it's the zarr-python version, which can write files with zarr-format=2 or 3).

I don't have a strong preference between the two strategies. My initial implementation would have required users to import the new module name with zarr-python 3.x, but I very slightly preferred trying to preserve source compatibility for that (potentially narrow) subset of use cases where source compatibility is possible to preserve.

Hmm, I guess I see that. I'm kind of on the fence because presumably at some point in the future we will only support zarr3 and then there'll be another breakage.

What do we think is our best migration strategy, if at some point in the future kvikio.zarr will mean kvikio.zarr3?

at some point in the future we will only support zarr3 and then there'll be another breakage.

What would cause the second breakage? Moving people back to kvikio.zarr, if we do the kvikio.zarr2 / kvikio.zarr3 split? IMO that wouldn't be worth the extra churn.

Here's an attempt at a summary. Under the PR as currently written, we have

time import behavior

previous kvikio.zarr only zarr-v2 supported

today kvikio.zarr zarr v2 or v3 supported. API depends on version at runtime. Some code will work with either, some code will break.

future kvikio.zarr only zarr v3 supported.

Under the split zarr.v2, zarr.v3 proposal

time import behavior

previous kvikio.zarr only zarr-v2 supported

today kvikio.zarr2 or kvikio.zarr3 zarr v2 or v3 supported. Users must update source to use the version they want.

future kvikio.zarr (or kvikio.zarr3?) only zarr v3 supported. Users must update their source if we deprecate kvikio.zarr3, or do nothing if we leave it as kvikio.zarr3.

I'm still leaning slightly towards just using kvikio.zarr thanks to that potential for maintaining source-compatibility, but won't push strongly for it.

I'm still leaning slightly towards just using kvikio.zarr thanks to that potential for maintaining source-compatibility, but won't push strongly for it.

Agree, but I have no strong opinion either :)
@wence- ?

Yeah, that's a convincing argument in favour of the approach here. Let's go with that

python/kvikio/kvikio/zarr/_zarr_python_3.py

wence- · 2025-03-03T13:02:16Z

python/kvikio/kvikio/zarr/_zarr_python_3.py

+    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


question: Should we figure out a way to hook this up with the normal asyncio stuff?

Maybe... When I went down a rabbit hole investigating zarr-developers/zarr-python#2863 (comment) I looked into this a bit and came across some discussions around supporting io_uring for asyncio on local files (python/cpython#85443, https://discuss.python.org/t/asyncio-for-files/31077). Those haven't gone anywhere.

In theory we ought to be able to do stuff with selectors and file descriptors. I'll do a bit more research and will open an issue.

#79 (comment) has a bit of a writeup with my findings if you're curious, but the tl/dr is that this seems pretty challenging (at least given my understanding of the C++ / CUDA side of things).

OK, thanks. It looks like if we would want to do that we need to return honest-to-goodness python futures from pread on the python side. There's some work changing how the C++ futures are launched that might help things here.

What I would be worried about getting right is that setting the python future status requires the gil so locking order is going to be tricky.

python/kvikio/kvikio/zarr/_zarr_python_3.py

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

TomAugspurger · 2025-03-05T16:40:40Z

CI stuff should be sorted out in 172eeb7.

The 3.10 environments are resolving to zarr-python 2.18.3 and running the tests in test_zarr.

The 3.12 tests are getting the latest version of zarr and are running the tests in [test_zarr_v3](https://github.com/rapidsai/kvikio/actions/runs/13638438229/job/38123429784#step:9:1955).

madsbk

Looks good

TomAugspurger · 2025-03-07T16:52:46Z

dependencies.yaml

        packages:
          - rapids-dask-dependency==25.4.*,>=0.0.0a0
          - pytest
+          - pytest-asyncio


This was needed to run the tests from zarr's StoreTests.

bdice

Packaging changes look good.

Before merging, could you rewrite the PR description to reflect the final state rather than the wide range of design options discussed in the early stages of this PR? It will go into the commit message and will also help future readers of this PR.

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

Support zarr-python 3.x

70abb2c

This implements a zarr-python 3.x compatible GDSStore.

madsbk added improvement Improves an existing functionality non-breaking Introduces a non-breaking change labels Feb 25, 2025

TomAugspurger added 3 commits February 25, 2025 14:34

Avoid memcopy in set

a870ff7

fixup

0d72485

fixup

4e9a9f2

TomAugspurger added 3 commits February 26, 2025 11:57

nvtx

4e08f8c

updated deps

9ad5c8a

temporarily disable zarr

7bbfe90

TomAugspurger mentioned this pull request Feb 27, 2025

Summary: Benchmark kvikio.CuFile Store pangeo-data/ncar-hackathon-xarray-on-gpus#23

Open

TomAugspurger added 4 commits February 28, 2025 18:09

fixups

e8c8f54

Refactor

47b9bd2

Moves zarr.py to _zarr_v2. Moves zarr_v3 to _zarr_v3.

docs

8bccefb

Merge remote-tracking branch 'upstream/branch-25.04' into tom/zarr-v3

91852e8

TomAugspurger added 2 commits February 28, 2025 21:04

cleanup

e2e7bca

cleanup

7b88b38

TomAugspurger marked this pull request as ready for review February 28, 2025 21:09

TomAugspurger requested review from a team as code owners February 28, 2025 21:09

TomAugspurger requested a review from KyleFromNVIDIA February 28, 2025 21:09

Fixed v3 skip

5cbe53e

TomAugspurger changed the title ~~[WIP]: Support zarr-python 3.x~~ Support zarr-python 3.x Feb 28, 2025

madsbk requested changes Mar 3, 2025

View reviewed changes

dependencies.yaml Outdated Show resolved Hide resolved

python/kvikio/kvikio/zarr/_zarr_python_3.py Show resolved Hide resolved

python/kvikio/kvikio/zarr/_zarr_python_3.py Show resolved Hide resolved

python/kvikio/kvikio/zarr/_zarr_python_3.py Show resolved Hide resolved

TomAugspurger added 2 commits March 3, 2025 14:17

zarr dependency

55e2fa8

Merge remote-tracking branch 'upstream/branch-25.04' into tom/zarr-v3

30fddab

wence- reviewed Mar 3, 2025

View reviewed changes

TomAugspurger added 9 commits March 3, 2025 15:26

Skip fixes

5eeb055

Note LocalFile similarity

c2a59bf

Add note about StoreTests

e8ec00d

fix in dependencies.yaml

6bd4079

Merge remote-tracking branch 'upstream/branch-25.04' into tom/zarr-v3

bc92a2c

Added example skip

1ca12e2

compat in benchmarks

172eeb7

Merge remote-tracking branch 'upstream/branch-25.04' into tom/zarr-v3

101d1f7

fixup

6e2e2e1

madsbk approved these changes Mar 7, 2025

View reviewed changes

TomAugspurger commented Mar 7, 2025

View reviewed changes

bdice approved these changes Mar 7, 2025

View reviewed changes

rapids-bot bot merged commit 5c710f3 into rapidsai:branch-25.04 Mar 7, 2025
61 checks passed

rapidsai deleted a comment from TomAugspurger Mar 7, 2025

TomAugspurger mentioned this pull request Mar 7, 2025

Tracking upstream changes pangeo-data/ncar-hackathon-xarray-on-gpus#28

Open

weiji14 added a commit to xarray-contrib/cupy-xarray that referenced this pull request Mar 11, 2025

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.

weiji14 mentioned this pull request Mar 11, 2025

Kvikio backend entrypoint with Zarr v3 xarray-contrib/cupy-xarray#70

Draft

6 tasks

TomAugspurger mentioned this pull request Mar 21, 2025

obstore-based Store implementation zarr-developers/zarr-python#1661

Merged

7 tasks

weiji14 mentioned this pull request Apr 2, 2025

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

Open

weiji14 mentioned this pull request Apr 9, 2025

Update to stable version of kvikio=25.04.00 pangeo-data/ncar-hackathon-xarray-on-gpus#33

Merged

time	import	behavior
previous	kvikio.zarr	only zarr-v2 supported
today	kvikio.zarr	zarr v2 or v3 supported. API depends on version at runtime. Some code will work with either, some code will break.
future	kvikio.zarr	only zarr v3 supported.

Support zarr-python 3.x #646

Support zarr-python 3.x #646

Uh oh!

Conversation

TomAugspurger commented Feb 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot bot commented Feb 25, 2025

Uh oh!

TomAugspurger commented Feb 25, 2025

Uh oh!

TomAugspurger commented Feb 28, 2025

Uh oh!

madsbk left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wence- commented Mar 3, 2025

Uh oh!

madsbk commented Mar 3, 2025

Uh oh!

TomAugspurger commented Mar 3, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

TomAugspurger commented Mar 5, 2025

Uh oh!

madsbk left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bdice left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

TomAugspurger commented Feb 25, 2025 •

edited

Loading