linumpy.gpu.nvcomp_zstd#

GPU-accelerated zstd codec for zarr-python, backed by nvCOMP.

This is a vendored stand-in for upstream zarr-python PR #2863 (zarr.codecs.gpu.NvcompZstdCodec), which is approved but not yet released. The codec implementation matches that PR closely so we can swap to the upstream class once it lands without changing call sites.

Why this exists#

With zarr.config.enable_gpu(), zarr-python (≤3.2) reads zstd-compressed chunks by:

GDSStore (or any store) yields a CPU buffer.
The default zarr.codecs.zstd.ZstdCodec decodes on the CPU (numcodecs).
The decoded chunk is copied H→D into a cupy buffer.

That CPU decode + H→D copy is the bottleneck for compressed tiles. This module replaces step 2: chunks are uploaded to the GPU and decoded with nvidia.nvcomp.Codec("Zstd"), so the result already lives on device.

Usage#

The codec is registered under the name "zstd" (same as the default), so any existing zarr file with zstd-compressed chunks will use it once registered. Registration happens lazily from linumpy.gpu.zarr_io whenever a GPU read path is taken, so CPU-only workflows are unaffected.

References

Classes#

NvcompZstdCodec

zstd codec that decodes/encodes on an NVIDIA GPU via nvCOMP.

Functions#

`register_nvcomp_zstd`()	Register `NvcompZstdCodec` under the name `zstd`.
`gpu_zstd_config`()	Return the zarr config overrides that route the `zstd` codec through nvCOMP.

Module Contents#

class linumpy.gpu.nvcomp_zstd.NvcompZstdCodec(*, level=0, checksum=False)[source]#

Bases: zarr.abc.codec.BytesBytesCodec

zstd codec that decodes/encodes on an NVIDIA GPU via nvCOMP.

Parameters:

level (int)
checksum (bool)

is_fixed_size = True[source]#

level: int = 0[source]#

checksum: bool = False[source]#

classmethod from_dict(data)[source]#

Construct codec from its serialised metadata representation.

Parameters:: data (dict[str, zarr.core.common.JSON])
Return type:: Self

to_dict()[source]#

Return the codec’s metadata dict (named zstd for on-disk compatibility).

Return type:: dict[str, zarr.core.common.JSON]

async decode(chunks_and_specs)[source]#

Decode a batch of zstd-compressed chunks on the GPU via nvCOMP.

Parameters:: chunks_and_specs (collections.abc.Iterable[tuple[zarr.core.buffer.Buffer | None, zarr.core.array_spec.ArraySpec]])
Return type:: collections.abc.Iterable[zarr.core.buffer.Buffer | None]

async encode(chunks_and_specs)[source]#

Encode a batch of chunks with zstd on the GPU via nvCOMP.

Parameters:: chunks_and_specs (collections.abc.Iterable[tuple[zarr.core.buffer.Buffer | None, zarr.core.array_spec.ArraySpec]])
Return type:: collections.abc.Iterable[zarr.core.buffer.Buffer | None]

compute_encoded_size(input_byte_length, chunk_spec)[source]#

Raise NotImplementedError — encoded size is data-dependent for zstd.

Parameters:

input_byte_length (int)
chunk_spec (zarr.core.array_spec.ArraySpec)

Return type:

int

linumpy.gpu.nvcomp_zstd.register_nvcomp_zstd()[source]#

Register NvcompZstdCodec under the name zstd.

Idempotent. Returns True if registration succeeded (or was already done), False if the required GPU dependencies (cupy, nvidia.nvcomp) are not importable.

Note that registration alone is not enough to make zarr use the GPU codec — zarr 3.2 selects the concrete class via zarr.config["codecs.zstd"] and the default points at the CPU codec. Use gpu_zstd_config() (or rely on linumpy.gpu.zarr_io.gpu_zarr_context()) to flip the config for a scoped block.

Return type:: bool

linumpy.gpu.nvcomp_zstd.gpu_zstd_config()[source]#

Return the zarr config overrides that route the zstd codec through nvCOMP.

Pass the result to zarr.config.set(...) to enable GPU-side zstd decoding for the duration of the set context. Caller is responsible for calling register_nvcomp_zstd() first (otherwise zarr cannot resolve the class name).

Return type:: dict[str, str]