DatasetRT

DatasetRT is a correctness-first dataset cache for ML training loops.

It gives you a deterministic, immutable cache on disk, backed by a Rust runtime and exposed through a small Python API. You keep your model code in PyTorch, JAX, TensorFlow, NumPy, or plain Python; DatasetRT handles cache integrity, metadata, sampling weights, and repeatable iteration without becoming another framework.

Authorship

Created by Vadym Stupakov vadim.stupakov@gmail.com.

Why ML Users Need This

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 built around one rule:

If it affects correctness, Rust owns it.

Rust owns:

• immutable cache publication
• manifest and checksum validation
• metadata schema validation
• shard offsets and index generation
• deterministic weighted sampling
• iterator state
• bounded reader/writer prefetch and worker pools
• weight table validation

Python stays thin and ergonomic. It describes your source data and receives bytes plus metadata back.

Quickstart

from pathlib import Path

import polars as pl

from dataset_rt import (
    CacheInput,
    CacheSourcesDatasetError,
    CacheSourcesDatasetSuccess,
    CachedDataset,
    ReaderConfig,
    ShardCompression,
    WriterConfig,
)


class Images:
    name = "train_images"

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


result = CachedDataset.from_cache_sources(
    Images(),
    Path("cache"),
    reader_config=ReaderConfig(seed=42, prefetch_size=64, num_workers=4, shuffle=True),
    writer_config=WriterConfig(
        prefetch_size=64,
        num_threads=4,
        shard_compression=ShardCompression(algo="none", ratio=1.0),
        show_progress=True,
    ),
)

match result:
    case CacheSourcesDatasetSuccess(dataset, results):
        pass
    case CacheSourcesDatasetError(results, message):
        raise RuntimeError(message)

for sample in dataset:
    image = decode_image(sample.data)  # domain decoding stays in Python
    label = sample.metadata["label"]

CachedDataset.from_cache_sources creates missing caches, reuses valid existing caches, and returns a result containing the loaded dataset plus per-source write outcomes. The cache directory argument is always a base cache directory; Rust writes each source under base_cache_dir / name. Cache writing shows committed samples/s and MB/s for the active source and source-count ETA for multi-source writes by default; pass WriterConfig(show_progress=False) for quiet jobs.

PyTorch

When PyTorch is installed, turn the same DatasetRT object into a sized IterableDataset:

torch_dataset = dataset.to_torch_iterable_dataset()
loader = torch.utils.data.DataLoader(torch_dataset, batch_size=None)

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

The adapter does not decode payloads or add a Torch dependency to DatasetRT. It yields CachedSample values and reports len(torch_dataset).

Metadata-Aware Weights

Weights are not a loose list that can drift out of alignment. DatasetRT exposes them as a Polars table with sample identity and metadata:

weights = dataset.weight_table()

rare = weights.with_columns(
    pl.when(pl.col("label") == "rare_class")
    .then(5.0)
    .otherwise(1.0)
    .alias("weight")
)

dataset.set_weight_table(rare)

The table contains:

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

Rust validates that every physical (cache_id, sample_id) appears exactly once and that every weight is positive and finite.

Multiple Sources

result = CachedDataset.from_cache_sources(
    [TrainImages(), SyntheticImages(), HardNegatives()],
    Path("cache"),
    reader_config=ReaderConfig(seed=123),
    writer_config=WriterConfig(prefetch_size=128, num_threads=8, show_progress=False),
)

If one source fails during a multi-source write, DatasetRT reports that source as CacheWriteError and keeps going. CacheSourcesDatasetSuccess.results tells you which sources were loaded and which were missing or malformed. A loaded dataset means every successful cache was validated from its manifest.

Storage Layout

cache/
    train_images/
        manifest.json
        metadata.arrow
        index.bin
        shards/
            000000.bin
            000001.bin

Metadata is stored separately from payload bytes. This keeps sampling, filtering, auditing, and weight editing independent of domain payload decoding. Each shard record also embeds the same metadata redundantly so raw record inspection and visualization can show sample context without joining back through Arrow.

What DatasetRT Does Not Do

DatasetRT does not decode JPEGs, PNGs, tensors, or framework-specific objects in the Rust core.

The core returns payload bytes. Your Python code or optional adapters can decode those bytes into tensors, arrays, images, token sequences, or any other domain object.

DatasetRT v0.1 also intentionally supports only:

• bytes-like payloads: bytes, bytearray, memoryview
• primitive metadata: bool, int, float, str
• shard compression: ShardCompression(algo="none", ratio=1.0) or ShardCompression(algo="lz4", ratio=...)

LZ4 compression is applied per payload record so random access stays direct. The common Rust zstd crate uses C bindings, so zstd compression is not enabled for the first stable version.

Documentation

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

Artifact Builds

Distribution artifacts are built locally with just build-all. Linux wheels use Zig cross-compilation; macOS arm64 and x86_64 wheels build on the local host. GitHub Actions stays checks-only.

Wheels use Python's stable ABI (cp310-abi3) and support Python 3.10 through 3.13.

Status

DatasetRT is at foundational v0.1 architecture. The core cache lifecycle, immutable storage, metadata, deterministic weighted sampling, Rust-owned reader/writer prefetching, and Polars weight table are in place.