DatasetRT

DatasetRT is a correctness-first dataset cache for ML training loops that need fast restarts, deterministic sampling, and metadata-aware control over what gets sampled.

It keeps payloads immutable on disk, keeps sampling state in Rust, and lets your Python training code steer the dataset through one Polars metadata table.

metadata = dataset.get_metadata()

dev_run = metadata.head(100)
dataset.update_metadata(dev_run)

for sample in dataset:
    train(sample.data, sample.metadata)

Rows in the active metadata table are the rows DatasetRT samples from. Remove rows for a fast development run, duplicate rows or update weight for class balancing and OHEM, add extra columns for training annotations, and let Rust validate that every row still maps to an immutable cache sample.

Why It Exists

Dataset bugs are expensive. A silent shuffle change, corrupt shard, mismatched metadata row, or weight vector applied to the wrong sample can waste training runs and make experiments impossible to reproduce.

DatasetRT is for the boring, high-stakes part of ML infrastructure:

• cache expensive preprocessing once and restart quickly
• keep payload bytes immutable while metadata stays editable at runtime
• filter or rebalance samples without rewriting cache files
• sample deterministically from the same cache contents, seed, epoch, and active metadata table
• validate metadata and weights against stable (cache_id, sample_id) identities
• keep reader and writer queues bounded in Rust

The design rule is simple:

If it affects correctness, Rust owns it.

Python remains the ergonomic edge: it describes sources, edits metadata with Polars, and decodes payload bytes into tensors, images, arrays, or domain objects.

The Core Idea

DatasetRT exposes the active dataset as a Polars table:

cache_id | sample_id | <stored metadata columns...> | weight | <extra columns...>

That table is the runtime control plane.

import polars as pl

metadata = dataset.get_metadata()

balanced = (
    metadata.join(
        metadata.group_by("label").agg(pl.len().alias("class_count")),
        on="label",
    )
    .with_columns((1.0 / pl.col("class_count")).alias("weight"))
    .drop("class_count")
)

dataset.update_metadata(balanced)

update_metadata() is runtime-only. It does not rewrite metadata.arrow, index.bin, shards, or manifests. Future iterators use the updated active table; already-created iterators keep their snapshot.

With drt.ReaderConfig(shuffle=True), DatasetRT performs deterministic weighted multinomial sampling with replacement over the active table. With shuffle=False, it emits active rows in cyclic table order, including duplicates. Use dataset.set_epoch_len(n) to make future iterators emit a finite window of n samples. Ordered reads continue through the active table cyclically across iterator boundaries, while shuffled reads continue a deterministic multinomial draw stream. Updating metadata resets the stream and sets the epoch length to the new table row count. Call set_epoch_len afterward to override it.

Installation

For GPU and distributed loader measurements, see the portable benchmark commands: python -m dataset_rt.run_benchmark --ranks 4 --num-workers 8 --device cuda.

DatasetRT supports Python 3.10 through 3.13.

uv add dataset-rt

Or with pip:

pip install dataset-rt

Quickstart

Wrap your existing data source as a Python iterable that yields payload bytes and primitive metadata.

from pathlib import Path

import dataset_rt as drt


class Images:
    name = "train_images"

    def __iter__(self):
        for image_id, image_bytes, label in load_my_images():
            yield drt.CacheInput(
                data=image_bytes,
                metadata={"image_id": image_id, "label": label},
            )


runtime = drt.DatasetRuntime(num_workers=4)
result = runtime.from_cache_sources(
    Images(),
    Path("cache"),
    reader_config=drt.ReaderConfig(seed=42, shuffle=True),
)

match result:
    case drt.CacheSourcesDatasetSuccess(dataset, results):
        pass
    case error:
        raise RuntimeError(error)

for sample in dataset:
    image = decode_image(sample.data)
    label = sample.metadata["label"]
    train(image, label)

The Rust cache stores bytes. DatasetRT does not decode JPEGs, tensors, tokens, or framework objects in the core; your Python code owns domain decoding.

When To Use It

Use DatasetRT when you care about:

• restart speed after expensive preprocessing
• deterministic training-data order
• class balancing, OHEM, or metadata-driven sampling
• safe development subsets without separate cache copies
• cache integrity and bounded native prefetching

It is not trying to be a training framework, image decoder, tensor format, or model-specific data pipeline. It is the cache and sampling layer underneath those pieces.

Documentation

• Python API
• Runtime Model
• Determinism
• Storage Format
• Architecture
• Serialization Boundary
• Build Artifacts
• Development

Status

DatasetRT is early production infrastructure. The core cache lifecycle, immutable storage format, runtime metadata table, deterministic weighted sampling, and Rust-owned reader/writer prefetching are in place.

Citation

@software{stupakov_datasetrt_2026,
  author = {Stupakov, Vadym},
  title = {DatasetRT: A Correctness-First Dataset Cache Runtime},
  year = {2026},
  url = {https://github.com/Red-Eyed/DatasetRT}
}