nvalchemi-zarr-perf

Zarr DataLoader Performance Tuning

Use this skill when optimizing nvalchemi Zarr reads or writing stores that will later be read through the nvalchemi DataLoader.

Overview

The pipeline has clean ownership boundaries:

• Reader: storage I/O only. Returns raw CPU tensor dictionaries plus metadata.
• Dataset: validation, optional validation skipping, device transfer, and async prefetch orchestration. Its canonical explicit batch API is load_batches(batch_index_lists).
• DataLoader: sampler/batch iteration, fused prefetch, stream usage, and batch construction.
• MultiDataset: global index composition over multiple Datasets while routing load_batches requests to child datasets.
• Sampler / batch_sampler: semantic sample order and batch membership. Do not rely on sampler windows to optimize storage I/O.

Reader public methods:

• reader.read(index): one sample.
• reader.read_many(indices): many samples, returned in the request order.

Reader backend hooks:

• _load_sample(index): implement for simple single-sample formats.
• _load_many_samples(indices): implement for batch-optimized formats.
• __len__(): total logical samples.

The base Reader owns metadata finalization and optional pinned memory. Index validity is the concrete reader's responsibility. AtomicDataZarrReader supports negative logical indices, maps through the active sample mask, and implements _load_many_samples as the fast path.

Recommended DataLoader setup

from nvalchemi.data.datapipes import (
    AtomicDataZarrReader,
    Dataset,
    DataLoader,
)

reader = AtomicDataZarrReader("store.zarr")

dataset = Dataset(
    reader,
    device="cuda",
    num_workers=1,          # 1 is enough; concurrent Zarr reads contend
    skip_validation=True,   # safe when store was written by the toolkit
)

loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    prefetch_factor=16,     # up to 64 * 16 = 1024 indices per backend read
    num_streams=2,
    use_streams=True,
    pin_memory=True,        # request pinned CPU tensors from the reader
)