dataio

A Rust-backed data loader for ML training. You declare a pipeline once; the native runtime fetches, decodes, transforms, and batches across a shared worker pool and hands PyTorch zero-copy tensors via DLPack. One URI surface covers local files, S3 / R2, GCS, Azure, and Hugging Face — the same plan runs unchanged against any of them.

import dataio

pipe = dataio.pipeline(
    dataio.slot("image")
        .decode_image(mode="rgb")
        .resize_short(256)
        .center_crop(h=224, w=224)
        .normalize("imagenet"),
    dataio.batch(64, prefetch=4),
)

with pipe.bind(records).run("torch", concurrency=64) as loader:
    for batch in loader:
        train_step(batch["image"])

Why dataio

• Native fast path. Decode, resize, normalize, and collate are Rust kernels on a shared pool — no per-sample Python on the hot path. The GIL is only held to hand the finished batch back.
• Zero-copy by default. Decoded tensors reach PyTorch through DLPack without a copy; reads slice memory-mapped buffers in place; outputs that a consumer buffer-reads are exported, not duplicated.
• One surface, every backend. s3://, gs://, az://, hf://, and local paths share the same pipeline ops and the same flat IO verbs.
• Built for training. Deterministic sharding, seeded shuffle, and checkpoint resume that replays a partial batch so transform RNG stays aligned across a restart.

Install

pip install dataio-rs

Python 3.11 or newer. Wheels are CPU-safe by default; CUDA, nvJPEG, and cuFile / GDS are discovered through dlopen at import time, so the same wheel runs on a laptop or an 8-GPU node. Print the capability report for the current host:

python -m dataio

Pipeline

A pipeline is a static plan: optional record maps, one or more slot chains, exactly one shape stage (batch / collate / group_by), and optional post-batch maps or native batch transforms. Building it does no work — binding a record source and calling .run() starts the runtime.

pipe = dataio.pipeline(
    dataio.map(lambda r: {"image": r["path"], "label": r["class_id"]}),

    dataio.slot("image")
        .decode_image(mode="rgb")
        .resize_short(256)
        .center_crop(h=224, w=224)
        .normalize("imagenet"),

    dataio.slot("label").decode_bytes(),
    dataio.batch(64, drop_last=True, prefetch=4),
)

Op factories are grouped into namespaces; the slot methods (.decode_image(), .resize(), .normalize(), …) are shorthand for the same native ops.

Slot-local conditionals

dataio.cond builds native predicates that compile into the plan signature. Field predicates run before decode, so they can gate optional slots; shape predicates run after decode and live inside transform branches.

from dataio import cond as c

pipe = dataio.pipeline(
    dataio.slot("image").decode_image(mode="rgb"),
    dataio.slot("mask", optional=True)
        .active_when(c.col("mask_id").is_not_null())
        .decode_image(mode="gray"),
    dataio.batch(32),
)

# Shape-driven branching after decode
dataio.slot("image") \
    .decode_image(mode="rgb") \
    .when(c.shape(axis=0).gt(512)) \
    .then(dataio.transforms.center_crop(h=512, w=512)) \
    .otherwise(dataio.transforms.resize(h=512, w=512))

Grouped batches

When an upstream sampler owns batch boundaries, group_by emits one batch per distinct key in arrival order. Pair it with from_sampler to adapt a PyTorch-style BatchSampler (or any iterable yielding lists of rows); the optional context= callback surfaces per-record sidecar data on batch.context, aligned with batch.keys:

pipe = dataio.pipeline(
    dataio.slot("latent").decode_safetensors(tensor_key=dataio.col("tensor_key")),
    dataio.group_by("batch_id", prefetch=4),
)

records = dataio.from_sampler(
    batch_sampler,
    lambda row: {
        "key": row["key"],
        "latent": row["cache_path"],
        "tensor_key": row["tensor_key"],
        "_row": row,
    },
)

with pipe.bind(records).run("torch", concurrency=128,
                            context=lambda r: r["_row"]) as loader:
    for batch in loader:
        latents = batch["latent"]
        rows = batch.context  # list[dict] aligned with batch.keys

from_sampler injects the batch_id field for you. The record_fn argument is optional — omit it when the sampler already yields valid record dicts.

Runtime

pipe.bind(records).run(...) returns a context-managed loader. Worker pools are process-wide and set once via environment variables; per-loader knobs control admission and prefetch:

concurrency="auto" (the default) is enough for local data and moderate batches; raise it for remote object reads or high-latency sources. Output preserves submit order by default — pass order="completion" to yield each batch as soon as it finishes.

Error handling

loader = pipe.bind(records).run("torch", on_failure="skip")

Per-row errors land on batch.errors. For automatic logging, pass log_errors=True (routes through the dataio.rows logger) or a callable for custom handling:

loader = pipe.bind(records).run("torch", log_errors=True)
loader = pipe.bind(records).run("torch", log_errors=my_logger.warning)

Distributed training and resume

loader = pipe.bind(records).run(
    "torch",
    shard={"id": rank, "of": world_size, "pad": True},
    seed=42,
    concurrency=64,
)

pad=True keeps every rank at the same batch count; seed drives shuffling and native batch-transform RNG. state_dict() / load_state_dict() checkpoint and resume the loader — mid-batch resume replays the partial batch from its first sample, so transform RNG stays aligned across the restart.

Direct IO

Flat verbs for scripts, evaluation, and one-shot reads:

dataio.head("s3://bucket/file.bin")
chunk = dataio.read("s3://bucket/file.bin", offset=0, size=1 << 20)
etag  = dataio.write("s3://bucket/out.bin", chunk)

dataio.exists("s3://bucket/out.bin")
dataio.ls("s3://bucket/prefix/")
dataio.glob("s3://bucket/**/*.safetensors")

For large eager reads, fill a caller-owned writable buffer instead of allocating a fresh Python bytes object:

buf = bytearray(1 << 30)
n = dataio.fetch_into("s3://bucket/checkpoint.bin", buf)
view = memoryview(buf)[:n]

Composable handles for repeated operations on one source:

src = dataio.reader("s3://bucket/shard.tar.gz")
payload = src.gunzip().tar_entry("0001.jpg").read_all()

with dataio.writer("s3://bucket/checkpoint.bin") as w:
    w.write_chunks(huge_generator())

Credentials and endpoints are configured once per scheme:

dataio.configure_credentials("s3", method="default")  # platform chain
dataio.configure_credentials("r2", access_key_id="...", secret_access_key="...")
dataio.configure_endpoint("r2", "https://<account>.r2.cloudflarestorage.com")

method="default" resolves through env vars, a shared profile, instance metadata, or workload identity — whichever the platform exposes.

Format helpers

arr   = dataio.read_numpy("/data/x.npy")
image = dataio.read_image("s3://bucket/image.png", mode="rgb")
cols  = dataio.read_parquet("/data/shard.parquet", columns=["key", "height"])
rows  = (dataio.scan_parquet("/data/shard.parquet")
            .filter(lambda r: r["height"] > 512)
            .collect())

for example in dataio.iter_tfrecord("/data/train.tfrecord"):
    label = example["label"]

safetensors

Zero-copy mmap reads for both single files and sharded checkpoints:

dataio.write_safetensors(
    "/fsx/ckpt-step-12000/",
    state_dict,
    max_shard_size="5GB",
    metadata={"format": "pt"},
)

state_dict = dataio.read_safetensors("/fsx/ckpt-step-12000/", framework="torch")

with dataio.open_safetensors("/fsx/model.safetensors") as st:
    tensor = st.read(key="layer.weight", framework="torch")

On Linux hosts with CUDA and a GDS-ready filesystem, reads land directly in device memory:

state_dict = dataio.read_safetensors(
    "/fsx/ckpt-step-12000/",
    framework="torch",
    device="cuda:0",
)

dataio.gds_available() and dataio.gds_info() report the active state.

RDMA transport

dataio.rdma is a native one-sided RDMA client + server over libfabric (OFI), built for fast peer-to-peer checkpoint transport over EFA / verbs / RoCE. It is CPU-safe: libfabric.so.1 is discovered via dlopen at first use, and on a host without an RDMA provider available() returns False while tcp;ofi_rxm / sockets providers still run as slow functional fallbacks. Handles are context managers.

import dataio.rdma as rdma

with rdma.serve(bind="0.0.0.0:0", buffer_size=1 << 20) as server:
    addr = server.local_addr
    with rdma.connect(addr) as client:
        rdma.write(client, offset=0, data=b"hello rdma")
        assert rdma.read(client, offset=0, size=10) == b"hello rdma"

The data path is pipelined and zero-copy: a transfer is split into chunk_size segments with up to max_inflight posted concurrently (EFA is latency-bound per op, so a single in-flight op leaves the link idle), and a single endpoint reaches ~8 GB/s. read_into lands bytes straight into a caller-owned buffer (bytearray / numpy / pinned torch tensor / mmap) with no intermediate copy; write registers the caller's buffer directly as the RMA source.

Serving a checkpoint file needs no extra protocol — a self-describing format like safetensors works as-is, and a CUDA-connected client can pull each tensor's byte range straight into VRAM via GPUDirect RDMA:

# producer: serve weights read-only
server = rdma.serve_file("/fsx/model.safetensors", bind="0.0.0.0:0")

# consumer: GPUDirect-RDMA tensor ranges into VRAM
client = rdma.connect(server.local_addr, cuda=True)
import torch
dst = torch.empty(shape, dtype=torch.bfloat16, device="cuda:0")
rdma.read_into(client, dst, offset=tensor_byte_offset)  # NIC → VRAM, host untouched

GPUDirect needs an HMEM-capable provider and libfabric API ≥ 1.9. Rather than write a C fi_getinfo probe, run the self-diagnostic — it runs the real plain-RMA and HMEM opens and reports exactly why each would or wouldn't match:

for line in rdma.diagnose()["summary"]:
    print(line)

rdma.available() / rdma.info() give the boolean and per-provider capability snapshot.

API at a glance

Documentation

• docs/arch.md — architecture, runtime IR, zero-copy contracts, DLPack ownership, the RDMA crate, security.
• docs/benchmark.md — benchmark commands, regression gates, the full cross-framework matrix.
• examples/ — runnable end-to-end scripts. Good starting points: examples/multi_slot_pipeline.py, examples/safetensors_latents.py, examples/conditional_slots.py.

Development

The Rust toolchain is pinned in rust-toolchain.toml. Editable build:

uv run --no-sync --with maturin maturin develop --profile staging

Tests:

cargo test
uv run --no-sync python -m pytest python/tests -q

Wheel:

uv run --no-sync --with maturin maturin build --profile staging --out dist

Performance

Decode pipeline (variable-size JPEG → per-record bucket → symmetric CHW normalize), 32-core AMD EPYC 9R14, CPU only. Median samples/s; higher is better.

Full matrix, scaling charts, and reproduce commands in docs/benchmark.md.

License

Apache-2.0. Source: https://github.com/Mikubill/dataio-rs. PyPI: https://pypi.org/project/dataio-rs/.