Metadata-Version: 2.4
Name: dataset-rt
Version: 0.2.4
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Rust
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Dist: pydantic>=2
Requires-Dist: polars>=1
Requires-Dist: maturin[zig]>=1.7,<2 ; extra == 'dev'
Requires-Dist: pytest>=8 ; extra == 'dev'
Requires-Dist: ruff>=0.5 ; extra == 'dev'
Requires-Dist: pyrefly>=0.21 ; extra == 'dev'
Requires-Dist: pydantic-settings>=2 ; extra == 'dev'
Provides-Extra: dev
License-File: LICENSE
Summary: Framework-independent dataset cache runtime
Home-Page: https://github.com/Red-Eyed/DatasetRT
Author-email: Vadym Stupakov <vadim.stupakov@gmail.com>
License-Expression: MIT
Requires-Python: >=3.10
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM

# 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 over one reused Rust worker pool
- samples metadata validation

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

## Quickstart

```python
from pathlib import Path

import polars as pl

from dataset_rt import (
    CacheInput,
    CacheSourcesDatasetError,
    CacheSourcesDatasetSuccess,
    DatasetRuntime,
    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},
            )


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

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"]
```

`DatasetRuntime` creates exactly the requested number of Rust worker threads once and reuses them for cache loading, reading, and writing. `runtime.from_cache_sources` creates missing caches, reuses existing cache directories, 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. Existing cache checksum validation is opt-in with `validate_cache=True`; by default DatasetRT avoids hashing every payload shard during restart.

## PyTorch

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

```python
torch_dataset = dataset.to_torch_iterable_dataset()
loader = torch.utils.data.DataLoader(torch_dataset, batch_size=None, num_workers=0)

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)`. Keep PyTorch `DataLoader(num_workers=0)`; use `DatasetRuntime(num_workers=N)` for parallel cache reads.

## Samples Metadata

Weights are not a loose list that can drift out of alignment. DatasetRT exposes dataset-level samples metadata as a Polars table with stable identity columns, stored metadata, and editable weights:

```python
metadata = dataset.samples_metadata()

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

dataset.set_samples_metadata(rare)
```

The table contains:

```text
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

```python
large_runtime = DatasetRuntime(num_workers=8)
result = large_runtime.from_cache_sources(
    [TrainImages(), SyntheticImages(), HardNegatives()],
    Path("cache"),
    reader_config=ReaderConfig(seed=123),
    writer_config=WriterConfig(prefetch_size=128, 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

```text
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](docs/architecture.md)
- [Python API](docs/python-api.md)
- [Storage Format](docs/storage-format.md)
- [Runtime Model](docs/runtime.md)
- [Determinism](docs/determinism.md)
- [Serialization Boundary](docs/serialization.md)
- [Development](docs/development.md)
- [Build Artifacts](docs/build.md)
- [Changelog](CHANGELOG.md)

## 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 samples metadata table are in place.

