Zstd Codec on the GPU by akshaysubr · Pull Request #2863 · zarr-developers/zarr-python

akshaysubr · 2025-02-25T16:15:20Z

This PR adds a Zstd codec that runs on the GPU using the nvCOMP 4.2 python APIs.

TODO:

Make fully async
Performance benchmarking
CPU-GPU roundtrip testing
Add unit tests and/or doctests in docstrings
Add docstrings and API docs for any new/modified user-facing classes and functions
New/modified features documented in docs/user-guide/*.rst
Changes documented as a new file in changes/
GitHub Actions have all passed
Test coverage is 100% (Codecov passes)

dstansby · 2025-02-28T09:09:22Z

Thanks for opening this PR! At the moment we do not have any codecs implemented in the zarr-python package, but instead store them in numcodecs. So although it looks like there are some zarr-python specific changes that are needed to support the new GPU codec, the actual codec should be implemented in numcodecs, and then imported in zarr-python.

akshaysubr · 2025-02-28T16:00:06Z

@dstansby My understanding was that numcodecs is the place that python bindings to native codec implementations live and that with v3, the Codec class itself lives in zarr-python. The GPU codecs and python bindings are implemented in nvCOMP and imported through the nvidia-nvcomp-cu12 python package so I'm not sure which part of this would need to go in numcodecs. What did you have in mind?

…hread

TomAugspurger

Looks nice overall. I think the async side of things ended up in a pretty good spot. The code itself is pretty easy to follow (a normal stream synchronize). Having to do that on another host thread is a bit unfortunate, but there's only one synchronize per batch so this should be fine.

I left comments on a few things to clean up that I can help out with if you want @akshaysubr.

TomAugspurger · 2025-07-07T11:38:18Z

+    checksum: bool = False
+
+    def __init__(self, *, level: int = 0, checksum: bool = False) -> None:
+        # TODO: Set CUDA device appropriately here and also set CUDA stream


Agreed with leaving devices / streams as a TODO for now.

I want to enable users to overlap host-to-device memcpys with compute operations (like decode, but their own compute operations as well), but I'm not sure yet what that API will look like.

If you have any thoughts on how best to do this I'd love to hear them, and write them up as an issue.

#3271 for planning on devices and streams.

TomAugspurger · 2025-07-07T11:41:03Z

+        chunks_and_specs: Iterable[tuple[Buffer | None, ArraySpec]],
+    ) -> Iterable[Buffer | None]:
+        return [
+            spec.prototype.buffer.from_array_like(cp.array(a, dtype=np.dtype("b"), copy=False))


This is one spot where @weiji14's idea to use dlpack in #2658 (comment) would help. If NDBuffer new how to consume objects implementing the dlpack protocol, we could (maybe) get rid of the cp.array

TomAugspurger · 2025-07-09T19:30:53Z

Fixed some merge conflicts and changed the _convert_{to,from}_nvcomp_arrays methods to be sync.

TomAugspurger · 2025-07-09T19:38:38Z

        if array_like.ndim != 1:
            raise ValueError("array_like: only 1-dim allowed")
-        if array_like.dtype != np.dtype("B"):
+        if array_like.dtype.itemsize != 1:


The new tests in test_nvcomp.py were failing without this change.

I'd like to get us to a point where we don't care as much about the details of the buffer passed in here. This is an OK start I think.

what exactly does this check for? It's not clear to me why any numpy array that can be viewed as bytes wouldn't be allowed in here

and same for the dimensionality check, since any N-dimensional numpy array can be viewed as a 1D array.

Yeah, I'm not really sure...

I agree that the actual data we store internally here needs to be a byte dtype. Just doing cp.asarray(input).view("b") seems pretty reasonable to me.

i'm not even convinced that we need Buffer / NDBuffer, when Buffer is just a special case of NDBuffer where there's 1 dimension and the data type is bytes

we could even express this formally by:

making NDBuffer generic with two type parameters (number of dimensions and dtype)

having APIs that insist on consuming a Buffer instead insist on consuming NDBuffer[Literal[1], np.uint8]

(super out of scope for this PR ofc)

It's not clear to me why any numpy array that can be viewed as bytes wouldn't be allowed in here

I think this is mainly because NDBuffer objects don't need to be contiguous, but Buffer objects must be contiguous in memory which might be important when we send those out to codecs that expect a contiguous memory slice.

But I agree that we can probably merge those two and make Buffer a specialization of NDBuffer.

codecov · 2025-07-09T19:59:31Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 83.66%. Comparing base (8f14d67) to head (5f9fe39).
⚠️ Report is 352 commits behind head on main.

❗ There is a different number of reports uploaded between BASE (8f14d67) and HEAD (5f9fe39). Click for more details.

HEAD has 1 upload less than BASE

Flag BASE (8f14d67) HEAD (5f9fe39)

tests 7 6

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2863      +/-   ##
==========================================
- Coverage   92.91%   83.66%   -9.25%     
==========================================
  Files          86       98      +12     
  Lines       10707    19415    +8708     
==========================================
+ Hits         9948    16243    +6295     
- Misses        759     3172    +2413

see 77 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

…ed config - decode()/encode() created a cp.cuda.Event() but never called .record() before .synchronize(), so the "wait for GPU completion" was a no-op -- a race, not just wasted asyncio.to_thread overhead. Now records on the stream the nvcomp op was enqueued on before waiting. - _zstd_codec was a cached_property binding device/stream once, forever, even though codec/pipeline instances are cached and reused (e.g. by sharding's pipeline cache). Now built fresh per call from the caller's current device/stream. - level/checksum were parsed and round-tripped through to_dict/from_dict but never passed to nvcomp.Codec(...). checksum now maps to nvcomp's ChecksumPolicy; level has no nvCOMP equivalent, so a nonzero value now raises NotImplementedError instead of being silently dropped. - is_fixed_size was True for a compressor (encoded size is data-dependent, same as the CPU ZstdCodec's is_fixed_size = False); left is_fixed_size True made compute_encoded_size's NotImplementedError reachable by any caller that correctly trusts is_fixed_size before calling it.

# Conflicts: # docs/user-guide/gpu.md # pyproject.toml # src/zarr/codecs/__init__.py # src/zarr/codecs/gpu.py # src/zarr/core/buffer/gpu.py # tests/test_codecs/test_codecs.py # tests/test_codecs/test_nvcomp.py

# Conflicts: # src/zarr/core/buffer/gpu.py # tests/test_api.py

…in CI

…of cast Answers review question on why the cast was needed: get_codec_class returns type[Codec]; the registry lookup exists so enable_gpu() can swap in NvcompZstdCodec. A runtime isinstance check replaces the static cast.

…n gpu.NDBuffer (S101)

…Buffer object-dtype error with main

negin513 · 2026-10-14T19:49:29Z

@akshaysubr I've opened akshaysubr#4 against this PR's gpu-codecs branch: it merges current main (resolving the conflicts here), carries @TomAugspurger's follow-up commits that hadn't been pushed to this branch, and fixes a few NvcompZstdCodec issues (unrecorded sync event, cached device/stream, ignored level/checksum config, is_fixed_size). CPU tests/lint/mypy pass locally; GPU tests still need a run on GPU hardware. Merging it into gpu-codecs will update this PR.

…ailable in CI" This reverts commit 4d4463f.

Per review discussion on zarr-developers#2863 (JeppeKlitgaard, d-v-b, TomAugspurger): the extra installs CUDA-specific packages (cupy-cuda12x, nvidia-nvcomp-cu12), so name it after CUDA rather than GPUs in general, leaving room for other accelerator backends. The 'gpu' extra is removed (extras can't be deprecated). enable_gpu() and zarr.buffer.gpu are unchanged: they predate this PR and are public API.

Raising NotImplementedError in __init__ also fired from from_dict, so any array written by the CPU ZstdCodec with level != 0 could not be opened under enable_gpu(). level is an encode-only parameter; accept and preserve it in metadata as the CPU codec does. Restore the device/stream TODO.

nvCOMP only allows NO_COMPUTE_NO_VERIFY with a RAW bitstream, so mapping checksum to ChecksumPolicy made checksum=True fail on both encode and decode. The CUDA event is still recorded before synchronizing.

Sync gpu-codecs with main, rename the gpu extra to cuda, fix nvCOMP Zstd sync and options

github-actions Bot added the needs release notes Automatically applied to PRs which haven't added release notes label Feb 25, 2025

TomAugspurger reviewed Feb 25, 2025

View reviewed changes

Comment thread src/zarr/codecs/gpu.py Outdated

akshaysubr added 2 commits February 25, 2025 21:30

First working version of Zstd codec on the GPU

49d5ee8

Adding nvcomp to the GPU dependency list

d548adc

akshaysubr force-pushed the gpu-codecs branch from d6608e7 to d548adc Compare February 25, 2025 21:30

Updating codec pipeline batch size for GPU codecs to enable parallelism

a8c0db3

madsbk mentioned this pull request Feb 28, 2025

Multi-GPU setup and python/zarr api rapidsai/kvikio#651

Closed

This was referenced Mar 3, 2025

Support zarr-python 3.x rapidsai/kvikio#646

Merged

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

Open

weiji14 mentioned this pull request Mar 11, 2025

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

Draft

3 of 6 tasks

TomAugspurger reviewed Mar 18, 2025

View reviewed changes

Comment thread src/zarr/codecs/cuda.py

vyasr mentioned this pull request Apr 7, 2025

Evaluate proper home for Zarr bindings in kvikio rapidsai/kvikio#686

Open

weiji14 mentioned this pull request Apr 7, 2025

Request for packaging Python bindings conda-forge/nvcomp-feedstock#21

Closed

TomAugspurger reviewed Jun 9, 2025

View reviewed changes

Comment thread src/zarr/codecs/cuda.py

Comment thread src/zarr/codecs/gpu.py Outdated

Comment thread src/zarr/codecs/cuda.py

akshaysubr added 2 commits July 3, 2025 11:59

Making encode and decode async

69aa274

Removing custom awaitable in favor of event synchronize in an async t…

10e1bc9

…hread

TomAugspurger reviewed Jul 7, 2025

View reviewed changes

TomAugspurger added 2 commits July 9, 2025 12:28

Merge remote-tracking branch 'upstream/main' into gpu-codecs

771c0c1

Sync convert methods

ec07100

TomAugspurger added 2 commits July 9, 2025 12:36

test coverage

d1c37a3

loosen dtype restriction

69ea74e

TomAugspurger reviewed Jul 9, 2025

View reviewed changes

TomAugspurger marked this pull request as ready for review July 9, 2025 19:38

fixed Buffer.__add__

1b85fdc

TomAugspurger added 2 commits July 11, 2025 12:30

Added whatsnew

f5c7814

Merge remote-tracking branch 'upstream/main' into gpu-codecs

7671274

TomAugspurger and others added 22 commits August 25, 2026 15:58

@gpu_test

8a801c7

wip test stuff

ee1ec92

doc updates

3264865

added failing compatibility test

bd34083

added a matching test

1b024fb

Some buffer coverage

576ba5f

coverage

93d6010

update error message

314bc54

private

56b9e00

test fixup

ef2e2f7

doc fix

166c011

lint

fbefa5e

resolve ambiguous link

21dea73

cupy 14 compat

1ad457d

Apply suggestion from @TomAugspurger

58cee0e

Merge akshay/gpu-codecs (PR zarr-developers#2863 tip) into update branch

1acd153

# Conflicts: # docs/user-guide/gpu.md # pyproject.toml # src/zarr/codecs/__init__.py # src/zarr/codecs/gpu.py # src/zarr/core/buffer/gpu.py # tests/test_codecs/test_codecs.py # tests/test_codecs/test_nvcomp.py

Merge remote-tracking branch 'origin/main' into gpu-codecs-update-main

54a8430

# Conflicts: # src/zarr/core/buffer/gpu.py # tests/test_api.py

exclude GPU files from coverage because they require GPU unavailable …

4d4463f

…in CI

lint: drop unused ChunkCoords import, replace assert with TypeError i…

f2722ab

…n gpu.NDBuffer (S101)

docs: remove duplicated Codecs section from config.md merge; align ND…

8287247

…Buffer object-dtype error with main

negin513 and others added 7 commits October 14, 2026 14:46

Revert "exclude GPU files from coverage because they require GPU unav…

fba64c2

…ailable in CI" This reverts commit 4d4463f.

Drop explanatory comment on is_fixed_size

b6dc00d

Set NvcompZstdCodec.is_fixed_size = True

f87dce1

Restore cached _zstd_codec; drop checksum_policy mapping

05d97a1

nvCOMP only allows NO_COMPUTE_NO_VERIFY with a RAW bitstream, so mapping checksum to ChecksumPolicy made checksum=True fail on both encode and decode. The CUDA event is still recorded before synchronizing.

Merge pull request #4 from negin513/gpu-codecs-review-fixes

5f9fe39

Sync gpu-codecs with main, rename the gpu extra to cuda, fix nvCOMP Zstd sync and options

Conversation

akshaysubr commented Feb 25, 2025

Uh oh!

Uh oh!

dstansby commented Feb 28, 2025

Uh oh!

akshaysubr commented Feb 28, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

TomAugspurger left a comment

Choose a reason for hiding this comment

Uh oh!

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

TomAugspurger commented Jul 9, 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Jul 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

negin513 commented Oct 14, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

codecov Bot commented Jul 9, 2025 •

edited

Loading