diff --git a/changes/2863.feature.md b/changes/2863.feature.md new file mode 100644 index 0000000000..09443879d5 --- /dev/null +++ b/changes/2863.feature.md @@ -0,0 +1,11 @@ +Added GPU-accelerated Zstd Codec + +This adds support for decoding with the Zstd Codec on NVIDIA GPUs using the +nvidia-nvcomp library. + +With `zarr.config.enable_gpu()`, buffers will be decoded using the GPU +and the output will reside in device memory. + +The optional-dependency extra for GPU support is now named `cuda` (``pip install "zarr[cuda]"``) +instead of `gpu`, since it installs CUDA-specific packages (CuPy and nvCOMP). The `gpu` extra +has been removed. diff --git a/docs/user-guide/config.md b/docs/user-guide/config.md index d1a70a14b0..2cdbcff8ae 100644 --- a/docs/user-guide/config.md +++ b/docs/user-guide/config.md @@ -1,4 +1,4 @@ -# Runtime configuration +# Runtime configuration {#config-runtime-configuration} [`zarr.config`][] is responsible for managing the configuration of zarr and is based on the [donfig](https://github.com/pytroll/donfig) Python library. @@ -50,6 +50,24 @@ first register the implementations in the registry and then select them in the c For example, an implementation of the bytes codec in a class `'custompackage.NewBytesCodec'`, requires the value of `codecs.bytes` to be `'custompackage.NewBytesCodec'`. +## Codecs + +Zarr and zarr-python split the logical codec definition from the implementation. +The Zarr metadata serialized in the store specifies just the codec name and +configuration. To resolve the specific implementation, a Python class, that's +used at runtime to encode or decode data, zarr-python looks up the codec name +in the codec registry. + +For example, after calling `zarr.config.enable_gpu()`, an nvcomp-based +codec will be used: + +```python exec="true" session="config" source="above" result="ansi" +with zarr.config.enable_gpu(): + print(zarr.config.get('codecs.zstd')) +``` + +## Default Configuration + This is the current default configuration: ```python exec="true" session="config" source="above" result="ansi" diff --git a/docs/user-guide/gpu.md b/docs/user-guide/gpu.md index 26d1c114b0..8a105fbc34 100644 --- a/docs/user-guide/gpu.md +++ b/docs/user-guide/gpu.md @@ -2,26 +2,17 @@ Zarr can use GPUs to accelerate your workload by running `zarr.config.enable_gpu()`. -!!! note - `zarr-python` currently supports reading the ndarray data into device (GPU) - memory as the final stage of the codec pipeline. Data will still be read into - or copied to host (CPU) memory for encoding and decoding. - - In the future, codecs will be available for compressing and decompressing data on - the GPU, avoiding the need to move data between the host and device for - compression and decompression. - ## Installation Zarr's GPU support requires [CuPy](https://cupy.dev), which in turn requires a -CUDA-compatible NVIDIA GPU. CuPy can be installed alongside Zarr with the `gpu` +CUDA-compatible NVIDIA GPU. CuPy and nvCOMP can be installed alongside Zarr with the `cuda` extra (see [Installation](installation.md) for the other optional dependency groups): ```console -pip install "zarr[gpu]" +pip install "zarr[cuda]" ``` -This installs the `cupy-cuda12x` package. If you need a CuPy build for a different +This installs the `cupy-cuda12x` and `nvidia-nvcomp-cu12` packages. If you need a CuPy build for a different CUDA version, see the [CuPy installation guide](https://docs.cupy.dev/en/stable/install.html) and install the appropriate package yourself. @@ -44,6 +35,9 @@ assert isinstance(z[:10, :10], cp.ndarray) Note that the arrays returned by reads are of type `cupy.ndarray` rather than NumPy arrays. +For supported codecs, data will be decoded using the GPU via the [nvcomp] library. +See [runtime-configuration][config-runtime-configuration] for more. Issues and feature requests for NVIDIA nvCOMP can be reported on the nvcomp [issue tracker]. + `zarr.config.enable_gpu()` returns a [donfig](https://donfig.readthedocs.io/en/latest/) `ConfigSet`, which can be used as a context manager to enable GPU support for a limited scope: @@ -58,3 +52,6 @@ Under the hood, `enable_gpu()` selects the GPU-backed buffer classes `zarr.buffer.gpu.Buffer` and `zarr.buffer.gpu.NDBuffer` via the `buffer` and `ndbuffer` configuration keys. See [Custom array buffers](extending.md#custom-array-buffers) for more on Zarr's buffer classes, including how to implement your own. + +[nvcomp]: https://docs.nvidia.com/cuda/nvcomp/samples/python_samples.html +[issue tracker]: https://github.com/NVIDIA/CUDALibrarySamples/issues diff --git a/docs/user-guide/gpu.rst b/docs/user-guide/gpu.rst new file mode 100644 index 0000000000..cd4f6c5eaf --- /dev/null +++ b/docs/user-guide/gpu.rst @@ -0,0 +1,34 @@ +.. _user-guide-gpu: + +Using GPUs with Zarr +==================== + +Zarr can use GPUs to accelerate your workload by running +:meth:`zarr.config.enable_gpu`. + +Reading data into device memory +------------------------------- + +:meth:`zarr.config.enable_gpu` configures Zarr to use GPU memory for the data +buffers used internally by Zarr. + +.. code-block:: python + + >>> import zarr + >>> import cupy as cp # doctest: +SKIP + >>> zarr.config.enable_gpu() # doctest: +SKIP + >>> store = zarr.storage.MemoryStore() # doctest: +SKIP + >>> z = zarr.create_array( # doctest: +SKIP + ... store=store, shape=(100, 100), chunks=(10, 10), dtype="float32", + ... ) + >>> type(z[:10, :10]) # doctest: +SKIP + cupy.ndarray + +Note that the output type is a ``cupy.ndarray`` rather than a NumPy array. + +For supported codecs, data will be decoded using the GPU via the `nvcomp`_ +library. See :ref:`user-guide-config` for more. Isseus and feature requests +for NVIDIA nvCOMP can be reported in the `nvcomp issue tracker`_. + +.. _nvcomp: https://docs.nvidia.com/cuda/nvcomp/samples/python_samples.html +.. _nvcomp issue tracker: https://github.com/NVIDIA/CUDALibrarySamples/issues diff --git a/docs/user-guide/installation.md b/docs/user-guide/installation.md index a7487e83c8..90c5b61f45 100644 --- a/docs/user-guide/installation.md +++ b/docs/user-guide/installation.md @@ -21,10 +21,10 @@ pip install zarr ``` There are a number of optional dependency groups you can install for extra functionality. -These can be installed using `pip install "zarr[]"`, e.g. `pip install "zarr[gpu]"` +These can be installed using `pip install "zarr[]"`, e.g. `pip install "zarr[cuda]"` - `remote`: support for reading/writing to remote data stores (fsspec, obstore) -- `gpu`: support for GPUs (cupy) +- `cuda`: support for NVIDIA CUDA GPUs (cupy, nvidia-nvcomp) - `cli`: support for the `zarr` [command-line interface](cli.md) (typer) - `optional`: support for path-like access to local and remote stores (universal-pathlib) - `cast-value-rs`: support for the `cast_value` codec (cast-value-rs) diff --git a/pyproject.toml b/pyproject.toml index c68450474e..36c5ef012d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,8 +80,9 @@ remote = [ "fsspec>=2023.10.0", "obstore>=0.5.1", ] -gpu = [ +cuda = [ "cupy-cuda12x; sys_platform != 'darwin'", + "nvidia-nvcomp-cu12; sys_platform != 'darwin'", ] cast-value-rs = ["cast-value-rs>=0.4.2"] cli = ["typer"] @@ -216,7 +217,7 @@ extra-dependencies = [ # docs example is never executed on GPU hardware. "pytest-examples", ] -features = ["gpu"] +features = ["cuda"] [[tool.hatch.envs.gputest.matrix]] python = ["3.12", "3.13"] diff --git a/src/zarr/codecs/__init__.py b/src/zarr/codecs/__init__.py index 9a1b47b351..54d76f9b67 100644 --- a/src/zarr/codecs/__init__.py +++ b/src/zarr/codecs/__init__.py @@ -4,6 +4,7 @@ from zarr.codecs.bytes import BytesCodec, Endian from zarr.codecs.cast_value import CastValue from zarr.codecs.crc32c_ import Crc32cCodec +from zarr.codecs.cuda import NvcompZstdCodec from zarr.codecs.gzip import GzipCodec from zarr.codecs.numcodecs import ( BZ2, @@ -44,6 +45,7 @@ "Crc32cCodec", "Endian", "GzipCodec", + "NvcompZstdCodec", "ScaleOffset", "ShardingCodec", "ShardingCodecIndexLocation", diff --git a/src/zarr/codecs/cuda.py b/src/zarr/codecs/cuda.py new file mode 100644 index 0000000000..6c687b2f4f --- /dev/null +++ b/src/zarr/codecs/cuda.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from functools import cached_property +from typing import TYPE_CHECKING + +import numpy as np + +from zarr.abc.codec import BytesBytesCodec +from zarr.core.common import JSON, parse_named_configuration +from zarr.registry import register_codec + +if TYPE_CHECKING: + from collections.abc import Iterable + from typing import Self + + from zarr.core.array_spec import ArraySpec + from zarr.core.buffer import Buffer + +try: + import cupy as cp +except ImportError: # pragma: no cover + cp = None + +try: + from nvidia import nvcomp +except ImportError: # pragma: no cover + nvcomp = None + + +def _parse_zstd_level(data: JSON) -> int: + if isinstance(data, int): + if data >= 23: + raise ValueError(f"Value must be less than or equal to 22. Got {data} instead.") + return data + raise TypeError(f"Got value with type {type(data)}, but expected an int.") + + +def _parse_checksum(data: JSON) -> bool: + if isinstance(data, bool): + return data + raise TypeError(f"Expected bool. Got {type(data)}.") + + +@dataclass(frozen=True) +class NvcompZstdCodec(BytesBytesCodec): + is_fixed_size = True + + level: int = 0 + checksum: bool = False + + def __init__(self, *, level: int = 0, checksum: bool = False) -> None: + # TODO: Set CUDA device appropriately here and also set CUDA stream + + level_parsed = _parse_zstd_level(level) + checksum_parsed = _parse_checksum(checksum) + + object.__setattr__(self, "level", level_parsed) + object.__setattr__(self, "checksum", checksum_parsed) + + @classmethod + def from_dict(cls, data: dict[str, JSON]) -> Self: + _, configuration_parsed = parse_named_configuration(data, "zstd") + return cls(**configuration_parsed) # type: ignore[arg-type] + + def to_dict(self) -> dict[str, JSON]: + return { + "name": "zstd", + "configuration": {"level": self.level, "checksum": self.checksum}, + } + + @cached_property + def _zstd_codec(self) -> nvcomp.Codec: + device = cp.cuda.Device() # Select the current default device + stream = cp.cuda.get_current_stream() # Use the current default stream + # Note: this returns an array with dtype=np.dtype("int8") + return nvcomp.Codec( + algorithm="Zstd", + bitstream_kind=nvcomp.BitstreamKind.RAW, + device_id=device.id, + cuda_stream=stream.ptr, + ) + + def _convert_to_nvcomp_arrays( + self, + chunks_and_specs: Iterable[tuple[Buffer | None, ArraySpec]], + ) -> tuple[list[nvcomp.Array], list[int]]: + none_indices = [i for i, (b, _) in enumerate(chunks_and_specs) if b is None] + filtered_inputs = [b.as_array_like() for b, _ in chunks_and_specs if b is not None] + # TODO: add CUDA stream here + return nvcomp.as_arrays(filtered_inputs), none_indices + + def _convert_from_nvcomp_arrays( + self, + arrays: Iterable[nvcomp.Array], + chunks_and_specs: Iterable[tuple[Buffer | None, ArraySpec]], + ) -> Iterable[Buffer | None]: + result: list[Buffer | None] = [] + + for a, (_, spec) in zip(arrays, chunks_and_specs, strict=True): + if a is None: + result.append(None) + else: + a2 = cp.array(a, dtype=a.dtype, copy=False) + if a2.dtype != np.dtype("B"): + a2 = a2.view(dtype=np.dtype("B")) + result.append(spec.prototype.buffer.from_array_like(a2)) + + return result + + async def decode( + self, + chunks_and_specs: Iterable[tuple[Buffer | None, ArraySpec]], + ) -> Iterable[Buffer | None]: + """Decodes a batch of chunks. + Chunks can be None in which case they are ignored by the codec. + + Parameters + ---------- + chunks_and_specs : Iterable[tuple[Buffer | None, ArraySpec]] + Ordered set of encoded chunks with their accompanying chunk spec. + + Returns + ------- + Iterable[Buffer | None] + """ + chunks_and_specs = list(chunks_and_specs) + + # Convert to nvcomp arrays + filtered_inputs, none_indices = self._convert_to_nvcomp_arrays(chunks_and_specs) + + outputs = self._zstd_codec.decode(filtered_inputs) if len(filtered_inputs) > 0 else [] + + # Record event for synchronization + event = cp.cuda.Event() + event.record(cp.cuda.get_current_stream()) + # Wait for decode to complete in a separate async thread + await asyncio.to_thread(event.synchronize) + + for index in none_indices: + outputs.insert(index, None) + + return self._convert_from_nvcomp_arrays(outputs, chunks_and_specs) + + async def encode( + self, + chunks_and_specs: Iterable[tuple[Buffer | None, ArraySpec]], + ) -> Iterable[Buffer | None]: + """Encodes a batch of chunks. + Chunks can be None in which case they are ignored by the codec. + + Parameters + ---------- + chunks_and_specs : Iterable[tuple[Buffer | None, ArraySpec]] + Ordered set of to-be-encoded chunks with their accompanying chunk spec. + + Returns + ------- + Iterable[Buffer | None] + """ + # TODO: Make this actually async + chunks_and_specs = list(chunks_and_specs) + + # Convert to nvcomp arrays + filtered_inputs, none_indices = self._convert_to_nvcomp_arrays(chunks_and_specs) + + outputs = self._zstd_codec.encode(filtered_inputs) if len(filtered_inputs) > 0 else [] + + # Record event for synchronization + event = cp.cuda.Event() + event.record(cp.cuda.get_current_stream()) + # Wait for encode to complete in a separate async thread + await asyncio.to_thread(event.synchronize) + + for index in none_indices: + outputs.insert(index, None) + + return self._convert_from_nvcomp_arrays(outputs, chunks_and_specs) + + def compute_encoded_size(self, _input_byte_length: int, _chunk_spec: ArraySpec) -> int: + raise NotImplementedError + + +register_codec("zstd", NvcompZstdCodec) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 06f7338755..526c3acdff 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -28,7 +28,6 @@ from zarr.codecs._v2 import V2Codec from zarr.codecs.bytes import BytesCodec from zarr.codecs.vlen_utf8 import VLenBytesCodec, VLenUTF8Codec -from zarr.codecs.zstd import ZstdCodec from zarr.core._info import ArrayInfo from zarr.core._json import buffer_to_json_object from zarr.core.array_spec import ArrayConfig, ArrayConfigLike, ArraySpec, parse_array_config @@ -154,6 +153,7 @@ _parse_array_array_codec, _parse_array_bytes_codec, _parse_bytes_bytes_codec, + get_codec_class, get_pipeline_class, ) from zarr.storage._common import StorePath, make_store_path @@ -5113,9 +5113,15 @@ def default_compressors_v3(dtype: ZDType[Any, Any]) -> tuple[BytesBytesCodec, .. """ Given a data type, return the default compressors for that data type. - This is just a tuple containing `ZstdCodec` + This is just a tuple containing an instance of the default "zstd" codec class. + + The codec is resolved through the registry (rather than hard-coding ``ZstdCodec``) + so that configuration overrides -- e.g. ``zarr.config.enable_gpu()`` mapping + "zstd" to ``NvcompZstdCodec`` -- take effect. ``get_codec_class`` returns + ``type[Codec]``, so the instance is narrowed to ``BytesBytesCodec`` with a + runtime check. """ - return (ZstdCodec(),) + return (_parse_bytes_bytes_codec(get_codec_class("zstd")()),) def default_serializer_v3(dtype: ZDType[Any, Any]) -> ArrayBytesCodec: diff --git a/src/zarr/core/buffer/gpu.py b/src/zarr/core/buffer/gpu.py index 6f0e4aec12..ae454bab61 100644 --- a/src/zarr/core/buffer/gpu.py +++ b/src/zarr/core/buffer/gpu.py @@ -8,9 +8,6 @@ cast, ) -import numpy as np -import numpy.typing as npt - from zarr.core.buffer import core from zarr.core.buffer.core import ArrayLike, BufferPrototype, NDArrayLike from zarr.errors import ZarrUserWarning @@ -23,6 +20,8 @@ from collections.abc import Iterable from typing import Self + import numpy.typing as npt + from zarr.core.common import BytesLike try: @@ -54,14 +53,14 @@ class Buffer(core.Buffer): def __init__(self, array_like: ArrayLike) -> None: if cp is None: - raise ImportError( + raise ImportError( # pragma: no cover "Cannot use zarr.buffer.gpu.Buffer without cupy. Please install cupy." ) if array_like.ndim != 1: raise ValueError("array_like: only 1-dim allowed") - if array_like.dtype != np.dtype("B"): - raise ValueError("array_like: only byte dtype allowed") + if array_like.dtype.itemsize != 1: + raise ValueError("array_like: only dtypes with itemsize=1 allowed") if not hasattr(array_like, "__cuda_array_interface__"): # Slow copy based path for arrays that don't support the __cuda_array_interface__ @@ -108,12 +107,13 @@ def as_numpy_array(self) -> npt.NDArray[Any]: return cast("npt.NDArray[Any]", cp.asnumpy(self._data)) def combine(self, others: Iterable[core.Buffer]) -> Self: - data = [cp.asanyarray(self._data)] + left = cp.asanyarray(self._data) + data = [left] for other in others: - other_array = other.as_array_like() - gpu_other = Buffer(other_array) - gpu_other_array = gpu_other.as_array_like() - data.append(cp.asanyarray(gpu_other_array)) + other_array = cp.asanyarray(other.as_array_like()) + if other_array.dtype != left.dtype: + other_array = other_array.view(left.dtype) + data.append(other_array) return self.__class__(cp.concatenate(data)) @@ -144,7 +144,7 @@ class NDBuffer(core.NDBuffer): def __init__(self, array: NDArrayLike) -> None: if cp is None: - raise ImportError( + raise ImportError( # pragma: no cover "Cannot use zarr.buffer.gpu.NDBuffer without cupy. Please install cupy." ) diff --git a/src/zarr/core/config.py b/src/zarr/core/config.py index c0ebb31353..de35992736 100644 --- a/src/zarr/core/config.py +++ b/src/zarr/core/config.py @@ -66,7 +66,15 @@ def enable_gpu(self) -> ConfigSet: Configure Zarr to use GPUs where possible. """ return self.set( - {"buffer": "zarr.buffer.gpu.Buffer", "ndbuffer": "zarr.buffer.gpu.NDBuffer"} + { + "buffer": "zarr.buffer.gpu.Buffer", + "ndbuffer": "zarr.buffer.gpu.NDBuffer", + "codecs": {"zstd": "zarr.codecs.cuda.NvcompZstdCodec"}, + "codec_pipeline": { + "path": "zarr.core.codec_pipeline.BatchedCodecPipeline", + "batch_size": 65536, + }, + } ) diff --git a/tests/test_api.py b/tests/test_api.py index f874bd45b7..daa4ffb8cb 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -16,6 +16,7 @@ import numpy.typing as npt + from zarr.abc.codec import Codec from zarr.abc.store import Store from zarr.core.common import JSON, AccessModeLiteral, MemoryOrder, ZarrFormat from zarr.types import AnyArray @@ -45,6 +46,7 @@ save_array, save_group, ) +from zarr.codecs import NvcompZstdCodec from zarr.core.buffer import NDArrayLike, default_buffer_prototype from zarr.core.sync import sync from zarr.errors import ( @@ -1860,14 +1862,15 @@ def test_api_exports() -> None: assert zarr.api.asynchronous.__all__ == zarr.api.synchronous.__all__ -@gpu_test +@gpu_test # type: ignore[misc,unused-ignore] @pytest.mark.parametrize( "store", ["local", "memory", "zip"], indirect=True, ) @pytest.mark.parametrize("zarr_format", [None, 2, 3]) -def test_gpu_basic(store: Store, zarr_format: ZarrFormat | None) -> None: +@pytest.mark.parametrize("codec", ["auto", NvcompZstdCodec()]) +def test_gpu_basic(store: Store, zarr_format: ZarrFormat | None, codec: str | Codec) -> None: import cupy as cp if zarr_format == 2: @@ -1875,7 +1878,7 @@ def test_gpu_basic(store: Store, zarr_format: ZarrFormat | None) -> None: # array to bytes. compressors = None else: - compressors = "auto" + compressors = codec with zarr.config.enable_gpu(): src = cp.random.uniform(size=(100, 100)) # allocate on the device diff --git a/tests/test_buffer.py b/tests/test_buffer.py index 307abfe499..50239611d2 100644 --- a/tests/test_buffer.py +++ b/tests/test_buffer.py @@ -188,6 +188,19 @@ def test_numpy_buffer_prototype() -> None: ndbuffer.as_scalar() +@gpu_test +def test_gpu_buffer_raises() -> None: + import cupy as cp + + arr = cp.empty((10, 10), dtype="B") + with pytest.raises(ValueError, match="array_like: only 1-dim allowed"): + gpu.Buffer(arr) + + arr = cp.arange(12, dtype="int32") + with pytest.raises(ValueError, match="array_like: only dtypes"): + gpu.Buffer(arr) + + @gpu_test def test_gpu_buffer_prototype() -> None: buffer = gpu.buffer_prototype.buffer.create_zero_length() diff --git a/tests/test_codecs/test_codecs.py b/tests/test_codecs/test_codecs.py index 8b4585503c..694c095f8e 100644 --- a/tests/test_codecs/test_codecs.py +++ b/tests/test_codecs/test_codecs.py @@ -17,12 +17,14 @@ GzipCodec, ShardingCodec, TransposeCodec, + ZstdCodec, ) from zarr.core.buffer import default_buffer_prototype from zarr.core.indexing import BasicSelection, decode_morton, morton_order_coords from zarr.core.metadata.v3 import ArrayV3Metadata from zarr.dtype import UInt8 from zarr.errors import ZarrUserWarning +from zarr.registry import register_codec from zarr.storage import MemoryStore, StorePath if TYPE_CHECKING: @@ -484,3 +486,22 @@ def test_resolve_metadata_only_mutates_shape(codec: Codec) -> None: assert spec_out.dtype == spec_in.dtype, f"{name} changed dtype" assert spec_out.fill_value == spec_in.fill_value, f"{name} changed fill_value" assert spec_out.config == spec_in.config, f"{name} changed config" + + +def test_uses_default_codec() -> None: + class MyZstdCodec(ZstdCodec): + pass + + register_codec("zstd", MyZstdCodec) + + with zarr.config.set( + {"codecs": {"zstd": f"{MyZstdCodec.__module__}.{MyZstdCodec.__qualname__}"}} + ): + a = zarr.create_array( + StorePath(zarr.storage.MemoryStore(), path="mycodec"), + shape=(10, 10), + chunks=(10, 10), + dtype="int32", + ) + assert a.metadata.zarr_format == 3 + assert isinstance(a.metadata.codecs[-1], MyZstdCodec) diff --git a/tests/test_codecs/test_nvcomp.py b/tests/test_codecs/test_nvcomp.py new file mode 100644 index 0000000000..142b0e51c1 --- /dev/null +++ b/tests/test_codecs/test_nvcomp.py @@ -0,0 +1,194 @@ +import contextlib +import typing +from collections.abc import Iterator + +import numpy as np +import pytest + +import zarr +from zarr.abc.store import Store +from zarr.buffer.gpu import buffer_prototype +from zarr.codecs import NvcompZstdCodec +from zarr.core.array_spec import ArrayConfig, ArraySpec +from zarr.storage import StorePath +from zarr.testing.utils import gpu_test + +if typing.TYPE_CHECKING: + from zarr.core.common import JSON + + +# the type-ignores here are here thanks to not reliably having GPU +# libraries in the pre-commit mypy environment. + + +@gpu_test # type: ignore[untyped-decorator,unused-ignore] +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize( + "checksum", + [ + False, + ], +) +@pytest.mark.parametrize( + "selection", + [ + (slice(None), slice(None)), # everything + (slice(4, None), slice(4, None)), # top-left chunk is empty + ], +) +def test_nvcomp_zstd(store: Store, checksum: bool, selection: tuple[slice, slice]) -> None: + import cupy as cp + + with zarr.config.enable_gpu(): + data = cp.arange(0, 256, dtype="uint16").reshape((16, 16)) + + a = zarr.create_array( + StorePath(store, path="nvcomp_zstd"), + shape=data.shape, + chunks=(4, 4), + dtype=data.dtype, + fill_value=0, + compressors=NvcompZstdCodec(level=0, checksum=checksum), + ) + + a[*selection] = data[*selection] + + if selection == (slice(None), slice(None)): + cp.testing.assert_array_equal(data[*selection], a[*selection]) + cp.testing.assert_array_equal(data[:, :], a[:, :]) + else: + assert a.nchunks_initialized < a.nchunks + expected = cp.full(data.shape, a.fill_value) + expected[*selection] = data[*selection] + cp.testing.assert_array_equal(expected[*selection], a[*selection]) + cp.testing.assert_array_equal(expected[:, :], a[:, :]) + + +@gpu_test # type: ignore[untyped-decorator,unused-ignore] +@pytest.mark.parametrize("host_encode", [True, False]) +def test_gpu_codec_compatibility(host_encode: bool) -> None: + # Ensure that the we can decode CPU-encoded data with the GPU + # and GPU-encoded data with the CPU + import cupy as cp + + @contextlib.contextmanager + def gpu_context() -> Iterator[None]: + with zarr.config.enable_gpu(): + yield + + if host_encode: + # CPU encode, GPU decode + write_ctx: contextlib.AbstractContextManager[None] = contextlib.nullcontext() + read_ctx: contextlib.AbstractContextManager[None] = gpu_context() + write_data = np.arange(16, dtype="int32").reshape(4, 4) + read_data = cp.array(write_data) + xp = cp + # MemoryStore holds Buffers; We write a CPU buffer, but read a GPU buffer + # which emits a warning. + expected_warning: pytest.WarningsRecorder | contextlib.AbstractContextManager[None] = ( + pytest.warns(zarr.errors.ZarrUserWarning) + ) + else: + # GPU encode, CPU decode + write_ctx = gpu_context() + read_ctx = contextlib.nullcontext() + write_data = cp.arange(16, dtype="int32").reshape(4, 4) + read_data = write_data.get() + xp = np + expected_warning = contextlib.nullcontext() + + store = zarr.storage.MemoryStore() + + with write_ctx: + z = zarr.create_array( + store=store, + shape=write_data.shape, + chunks=(4, 4), + dtype=write_data.dtype, + ) + z[:] = write_data + + with read_ctx, expected_warning: + # We need to reopen z, because `z.codec_pipeline` is set at creation + z = zarr.open_array(store=store, mode="r") + result = z[:] + assert isinstance(result, type(read_data)) + xp.testing.assert_array_equal(result, read_data) + + +@gpu_test # type: ignore[untyped-decorator,unused-ignore] +def test_invalid_raises() -> None: + with pytest.raises(ValueError): + NvcompZstdCodec(level=100, checksum=False) + + with pytest.raises(TypeError): + NvcompZstdCodec(level="100", checksum=False) # type: ignore[arg-type,unused-ignore] + + with pytest.raises(TypeError): + NvcompZstdCodec(checksum="False") # type: ignore[arg-type,unused-ignore] + + # a nonzero level is accepted and preserved in metadata so that arrays + # written by the CPU ZstdCodec can be read on the GPU + assert NvcompZstdCodec(level=5, checksum=False).level == 5 + + +@gpu_test # type: ignore[untyped-decorator,unused-ignore] +def test_uses_default_codec() -> None: + with zarr.config.enable_gpu(): + a = zarr.create_array( + StorePath(zarr.storage.MemoryStore(), path="nvcomp_zstd"), + shape=(10, 10), + chunks=(10, 10), + dtype="int32", + ) + assert a.metadata.zarr_format == 3 + assert isinstance(a.metadata.codecs[-1], NvcompZstdCodec) + + +@gpu_test # type: ignore[untyped-decorator,unused-ignore] +def test_nvcomp_from_dict() -> None: + config: dict[str, JSON] = { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": False, + }, + } + codec = NvcompZstdCodec.from_dict(config) + assert codec.level == 0 + assert codec.checksum is False + + +@gpu_test # type: ignore[untyped-decorator,unused-ignore] +def test_compute_encoded_chunk_size() -> None: + codec = NvcompZstdCodec(level=0, checksum=False) + with pytest.raises(NotImplementedError): + codec.compute_encoded_size( + _input_byte_length=0, + _chunk_spec=ArraySpec( + shape=(10, 10), + dtype=zarr.core.dtype.Int32(), + fill_value=0, + config=ArrayConfig(order="C", write_empty_chunks=False), + prototype=buffer_prototype, + ), + ) + + +@gpu_test # type: ignore[untyped-decorator,unused-ignore] +async def test_nvcomp_zstd_encode_none() -> None: + codec = NvcompZstdCodec(level=0, checksum=False) + chunks_and_specs = [ + ( + None, + ArraySpec( + shape=(10, 10), + dtype=zarr.core.dtype.Int32(), + fill_value=0, + config=ArrayConfig(order="C", write_empty_chunks=False), + prototype=buffer_prototype, + ), + ) + ] + result = await codec.encode(chunks_and_specs) + assert result == [None] diff --git a/uv.lock b/uv.lock index b14033c5b4..91dd404e37 100644 --- a/uv.lock +++ b/uv.lock @@ -2138,6 +2138,29 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/e5/f5c38d4899cd8edd5b522071c5658202077d113c3225484d779cf257970e/numpydoc-1.11.0-py3-none-any.whl", hash = "sha256:493ad2ce0d8275ca5da484ec37dbdbbc1334c3a6968dffd10ec729b51753e15e", size = 71464, upload-time = "2026-09-15T16:31:07.521Z" }, ] +[[package]] +name = "nvidia-libnvcomp-cu12" +version = "5.3.0.16" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/19/e4/d3351dbc6c3321ce10b540092644ef68d6330ad5a908cc3372f36b70778f/nvidia_libnvcomp_cu12-5.3.0.16-py3-none-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3b6f06cefc44465f6c0e406a7b4bc80ad399511623ff3308372cb05940edf942", size = 41293691, upload-time = "2026-07-14T21:00:39.197Z" }, + { url = "https://files.pythonhosted.org/packages/b7/20/e025aa3fa2ee33ee6a4871aa5a9e10b05472913adaff649ba21f2c98291d/nvidia_libnvcomp_cu12-5.3.0.16-py3-none-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:535944be2e37ead56ad1ef80e9dd58a3126fc2ccf238ddffc9d1cb7a48f5300a", size = 41460080, upload-time = "2026-07-14T21:01:14.445Z" }, + { url = "https://files.pythonhosted.org/packages/f3/7a/1ba038e15e7256019ecb23d52d6310fe9e1d24fe92074ee89c066ca0ac58/nvidia_libnvcomp_cu12-5.3.0.16-py3-none-win_amd64.whl", hash = "sha256:6b64abe9b7e64f5b2ceee7adf62af0ff6e4307a7cb0b20f51e9bf6d4387ec5d9", size = 30337687, upload-time = "2026-07-14T21:01:38.87Z" }, +] + +[[package]] +name = "nvidia-nvcomp-cu12" +version = "5.3.0.16" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-libnvcomp-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/b5/73/48cc25be9e6724820ad5b0bc15c687f1dd9807dd51516aeaa2dbf7e4ad28/nvidia_nvcomp_cu12-5.3.0.16-py3-none-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2cf53e55af520431d5855a417e2095b0b13a758cbdda8fe4fb58a35728081fb2", size = 3526553, upload-time = "2026-07-14T21:06:14.072Z" }, + { url = "https://files.pythonhosted.org/packages/26/07/b0b7e7862699b87679d919e39a35a71dbca88dcc656458be08e88fa1acf9/nvidia_nvcomp_cu12-5.3.0.16-py3-none-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c2cb2bd40b2a97381a46b3645f8e1b401325b6828c1f227b7685c9ee86412749", size = 3709563, upload-time = "2026-07-14T21:06:38.429Z" }, + { url = "https://files.pythonhosted.org/packages/58/04/8617da111cfea182a73f659b61ac68f32378313b26712998d22a066c52c3/nvidia_nvcomp_cu12-5.3.0.16-py3-none-win_amd64.whl", hash = "sha256:997dfec874c554e104e22136aca1a14a81bb33186301adf301236c08406da8c2", size = 1885029, upload-time = "2026-07-14T21:07:00.656Z" }, +] + [[package]] name = "obstore" version = "0.11.1" @@ -3655,8 +3678,9 @@ cast-value-rs = [ cli = [ { name = "typer" }, ] -gpu = [ +cuda = [ { name = "cupy-cuda12x", marker = "sys_platform != 'darwin'" }, + { name = "nvidia-nvcomp-cu12", marker = "sys_platform != 'darwin'" }, ] optional = [ { name = "universal-pathlib" }, @@ -3761,20 +3785,21 @@ test = [ [package.metadata] requires-dist = [ { name = "cast-value-rs", marker = "extra == 'cast-value-rs'", specifier = ">=0.4.2" }, - { name = "cupy-cuda12x", marker = "sys_platform != 'darwin' and extra == 'gpu'" }, + { name = "cupy-cuda12x", marker = "sys_platform != 'darwin' and extra == 'cuda'" }, { name = "donfig", specifier = ">=0.8" }, { name = "fsspec", marker = "extra == 'remote'", specifier = ">=2023.10.0" }, { name = "google-crc32c", specifier = ">=1.5" }, { name = "msgspec", specifier = ">=0.19" }, { name = "numcodecs", specifier = ">=0.16" }, { name = "numpy", specifier = ">=2" }, + { name = "nvidia-nvcomp-cu12", marker = "sys_platform != 'darwin' and extra == 'cuda'" }, { name = "obstore", marker = "extra == 'remote'", specifier = ">=0.5.1" }, { name = "packaging", specifier = ">=22.0" }, { name = "typer", marker = "extra == 'cli'" }, { name = "typing-extensions", specifier = ">=4.14" }, { name = "universal-pathlib", marker = "extra == 'optional'" }, ] -provides-extras = ["cast-value-rs", "cli", "gpu", "optional", "remote"] +provides-extras = ["cast-value-rs", "cli", "cuda", "optional", "remote"] [package.metadata.requires-dev] dev = [