/cupynumeric-parallel-data-load

name: cupynumeric-parallel-data-load
description: Load a sharded, on-disk dataset (sharded .npy, Parquet/Arrow, raw binary, sharded HDF5, custom layouts) into a distributed cuPyNumeric ndarray via a manual partition + leaf @task launch with CPU/OMP/GPU variants. Use when no single-call loader fits, including when per-shard row counts differ across files. Prefer cupynumeric.load or legate.io.hdf5.from_file when they apply.
license: CC-BY-4.0 OR Apache-2.0
compatibility: linux-x86_64, linux-aarch64, wsl-x86_64
metadata:
  version: "1.0.0"
  author: "NVIDIA Corporation <legate@nvidia.com>"
  upstream: https://github.com/nv-legate/cupynumeric
  docs: https://docs.nvidia.com/cupynumeric/latest/
  tags:
    - cupynumeric
    - legate
    - data-loading
    - io
    - distributed
    - parallel
    - gpu
    - sharded-data

Parallel sharded data -> cupynumeric load

**Why this skill exists.** cupynumeric mirrors NumPy's array API, including `cupynumeric.load` for a single `.npy` file. Beyond that, file *loading* lives in Legate, not cupynumeric:

| Format | Built-in loader | |---|---| | Single `.npy` | `cupynumeric.load(path)` (NumPy-API parity) | | HDF5 (single file) | `legate.io.hdf5.from_file` / `from_file_batched` | | Sharded multi-file (any format), Parquet/Arrow, raw binary, custom layouts | **No built-in loader — this skill.** |

This skill shows the canonical way to fill the gap in the last row: write a Legate Python task that calls the third-party reader the format needs (`h5py`, `pyarrow`, `np.memmap`, ...) inside the task body, and let Legate distribute the reads across GPUs / nodes. For the formats with a built-in loader, prefer it unless you need a custom in-task body (mmap-based loader, format-specific decoder, sidecar metadata, partial / sharded reads).

Canonical pattern: **manual partition + manual task launch, sized to the machine, not the files.** Only axis 0 is sharded; trailing axes ride along inside each tile. Per-shard row counts may differ across files (only `dtype` and trailing axes must match); the launch fills every available processor regardless of how many files there are.

`.npy` is the worked example because the header carries shape and dtype on disk, but the skeleton applies to any format with cheap range/slice reads (raw binary, HDF5, Parquet/Arrow — see "Other formats" below). Reference implementation: [`assets/examples/parallel_npy_load.py`](assets/examples/parallel_npy_load.py).

Data layout assumption

This skill is purely about **loading** — it assumes the data is already laid out on a shared filesystem in some predictable, indexable way. Producing those files is out of scope (the example ships a `write` subcommand for convenience, but real users bring their own).

The worked example assumes one specific layout:

• A directory containing files named `shard_0000.npy`, `shard_0001.npy`,

... in a contiguous integer sequence (zero-padded width 4).

• All shards share the same `dtype` and the same trailing axes

(`shape[1:]`); **axis 0 (rows per shard) may differ across files** — the recipe builds a cumulative row-offset table and reads each file's overlapping slice from inside the leaf task.

• The directory is visible to every rank (shared filesystem for

multi-node runs).

The example's `discover_layout()` prints what it found and hard-fails with a descriptive error when the layout is wrong (missing directory, no shards, mismatched `dtype` / trailing axes, or a hole in the contiguous `shard_NNNN.npy` sequence).

If your data lives in a different layout — fixed-stride raw binary, an HDF5 file with one dataset per shard, a directory tree, ... — only the glob pattern, the per-file reader (step 4 below), and the metadata discovery (step 1 below) change. The partitioning and launch machinery is layout-agnostic.

When to use

See the format table above for the routing decision (built-in loader vs. this skill). Beyond that, two additional cues that this skill is the right fit:

• Replacing sequential `np.concatenate([read(f) for f in files])` with

parallel per-GPU reads.

• Demonstrating how a user-defined Legate Python task writes into a

cupynumeric output array via a manual launch.

Examples

Paths below are written relative to this skill's directory (the script ships at `assets/examples/parallel_npy_load.py`). Adjust the prefix to match wherever your skill is installed (e.g. `skills/cupynumeric-parallel-data-load/assets/...` if the skill lives under a top-level `skills/` directory).

# Single-node, 4 GPUs.
legate --gpus 4 --fbmem 4000 --min-gpu-chunk 1 \
    assets/examples/parallel_npy_load.py \
    read --shard-dir /shared/scratch/demo

# Multi-node, 2 nodes x 4 GPUs (slurm), shared filesystem at --shard-dir.
# Generate the shards once on rank 0, then re-run `read` at any scale.
legate --launcher srun --nodes 2 --cpus 1 \
    assets/examples/parallel_npy_load.py \
    write --shard-dir /shared/scratch/demo

legate --launcher srun --nodes 2 --ranks-per-node 4 \
    --gpus 4 --fbmem 4000 --min-gpu-chunk 1 \
    assets/examples/parallel_npy_load.py \
    read --shard-dir /shared/scratch/demo

No layout flags — the read driver walks every `.npy` header to recover per-file row counts, the trailing shape, and the dtype, then derives `tile_rows` from the available processor count.

`--min-gpu-chunk 1` is only needed when the per-tile element count is below Legate's default minimum chunk size for GPU launches (e.g. the worked example's defaults — total rows split across 4 GPUs at `~1M` per tile — fall below the threshold and would otherwise be folded onto a single GPU). For production-sized datasets (tens of millions of elements per tile or larger) you can drop the flag and let Legate use its default. Bumping it to a moderate value (e.g. `--min-gpu-chunk 1024`) is fine when each tile is large enough that per-task overhead matters more than getting *every* GPU a tile.

Instructions

Five steps from a `.npy` worked e