linumpy.gpu.zarr_io#

High-level zarr → GPU loading with automatic backend selection.

Public entry points:

read_zarr_to_gpu() — load an entire zarr array onto the GPU using the fastest available path (kvikio / GDS, falling back to zarr.config.enable_gpu).
gpu_zarr_context() — context manager that flips zarr into GPU mode for its duration, so subsequent zarr.open_array(...) calls return arrays whose slicing materialises directly into cupy.ndarray. Use this for tile-by-tile / per-slab access patterns where loading the whole volume at once is wasteful.

Selection order (when prefer='auto'):

kvikio (GPUDirect Storage, native mode) — chunks DMA’d directly from NVMe into GPU memory. Requires kvikio installed, GDS in native mode, and an uncompressed zarr v2/v3.
zarr.config.enable_gpu() — host I/O with on-host decode then a single H→D copy. Works for any zarr (compressed or not). The fallback when GDS is unavailable, in compat mode, or the array is compressed.

Backend implementations live in their own modules:

linumpy.gpu.kvikio_zarr — kvikio / GDS reader.

Reference numbers on a 16 GiB float32 zarr v3 (RTX A6000, ext4, GDS native) warm cache: kvikio ~9.9 GiB/s, zarr-gpu ~7.1 GiB/s.

Attributes#

Backend

Functions#

`read_zarr_via_zarr_gpu`(array_path)	Load a zarr array onto the GPU using `zarr.config.enable_gpu()`.
`read_zarr_to_gpu`(array_path, *[, prefer])	Load a zarr array onto the GPU using the fastest available path.
`gpu_zarr_context`()	Context manager that puts zarr into GPU mode for arbitrary slice reads.

Module Contents#

linumpy.gpu.zarr_io.Backend[source]#

linumpy.gpu.zarr_io.read_zarr_via_zarr_gpu(array_path)[source]#

Load a zarr array onto the GPU using zarr.config.enable_gpu().

Host I/O with on-host decode, then a single H→D copy. Works for any zarr array (including compressed) and is the recommended fallback when GDS is unavailable or stuck in compat mode.

Parameters:: array_path (str | pathlib.Path) – Path to the zarr array directory.
Returns:: Device-resident array.
Return type:: cupy.ndarray

linumpy.gpu.zarr_io.read_zarr_to_gpu(array_path, *, prefer='auto')[source]#

Load a zarr array onto the GPU using the fastest available path.

Selection order (when prefer='auto'):

kvikio / GDS — only if kvikio is in native or auto mode AND the array is uncompressed v2/v3.
zarr.config.enable_gpu() — works for any zarr.

Parameters:

array_path (str | pathlib.Path) – Path to the zarr array directory.
prefer (Backend) – 'auto' (default), 'kvikio', or 'zarr-gpu'. Forcing a path will raise if that path is unavailable for this array.

Returns:

Device-resident array of shape and dtype matching the zarr metadata.

Return type:

cupy.ndarray

linumpy.gpu.zarr_io.gpu_zarr_context()[source]#

Context manager that puts zarr into GPU mode for arbitrary slice reads.

Inside this context, any subsequent zarr.open_array(...) returns an array whose slicing produces cupy.ndarray results — chunks are decoded on host then transferred to device on each vol[slice] operation. This is the right pattern for tile-by-tile or per-slab work where loading the full volume at once is wasteful.

Outside this context, zarr falls back to its normal numpy-backed mode.

Examples

>>> from linumpy.gpu.zarr_io import gpu_zarr_context
>>> from linumpy.io.zarr import read_omezarr
>>> with gpu_zarr_context():
...     vol, _ = read_omezarr(path, level=0)
...     for tile_region in regions:
...         tile_gpu = vol[tile_region]  # already cupy

Raises:: RuntimeError – If zarr is unavailable.
Return type:: collections.abc.Iterator[None]