Metadata-Version: 2.3
Name: syvain-training-data
Version: 0.0.252
Summary: Syvain training data manifest, loading, and saving utilities
Requires-Dist: msgpack>=1.1.2,<2.0.0
Requires-Dist: obstore>=0.11.0,<0.12.0
Requires-Dist: pydantic>=2.13.4
Requires-Dist: zstandard>=0.25.0,<0.26.0
Requires-Dist: numpy>=2.0.0 ; extra == 'dev'
Requires-Dist: pyarrow>=23.0.1,<24.0.0 ; extra == 'dev'
Requires-Dist: pytest>=8.0.0 ; extra == 'dev'
Requires-Dist: ruff>=0.15.12 ; extra == 'dev'
Requires-Dist: safetensors>=0.7.0,<0.8.0 ; extra == 'dev'
Requires-Dist: torch>=2.13.0 ; extra == 'dev'
Requires-Dist: ty>=0.0.34 ; extra == 'dev'
Requires-Dist: numpy>=2.0.0 ; extra == 'full'
Requires-Dist: pyarrow>=23.0.1,<24.0.0 ; extra == 'full'
Requires-Dist: safetensors>=0.7.0,<0.8.0 ; extra == 'full'
Requires-Dist: torch>=2.13.0 ; extra == 'full'
Requires-Dist: pyarrow>=23.0.1,<24.0.0 ; extra == 'parquet'
Requires-Dist: numpy>=2.0.0 ; extra == 'representation'
Requires-Dist: pyarrow>=23.0.1,<24.0.0 ; extra == 'representation'
Requires-Dist: safetensors>=0.7.0,<0.8.0 ; extra == 'representation'
Requires-Dist: torch>=2.13.0 ; extra == 'representation'
Requires-Dist: torch>=2.13.0 ; extra == 'torch'
Requires-Python: >=3.12, <3.15
Provides-Extra: dev
Provides-Extra: full
Provides-Extra: parquet
Provides-Extra: representation
Provides-Extra: torch
Description-Content-Type: text/markdown

# syvain-training-data

Internal [Syvain](https://syvain.com/) data utility. No secret sauce here, just
a shared helper.

> This is my dataloader. There are many like it, but this one is mine. My
> dataloader is my best friend. It is my life. I must master it as I must master
> my life. My dataloader, without me, is useless. Without my dataloader, I am
> useless.

## Install

```bash
uv add 'syvain-training-data[torch]'
```

This installs the PyTorch data loader used in the primary example below. For
manifest, storage, JSONL gzip, and MessagePack/Zstandard workflows that do not
load training batches, use the lightweight base package:

```bash
uv add syvain-training-data
```

The base install does not install NumPy, PyTorch, PyArrow, or Safetensors.
Parquet users can install `syvain-training-data[parquet]`. Consumers of the
OLMo 3 representation publication can install
`syvain-training-data[representation]`. The `full` extra installs every
optional feature. Requesting an optional feature without its extra raises an
import error that names the required extra.

## Load data

```python
from syvain_training_data import SyvainTrainingData

training_data = SyvainTrainingData(
    s3_base_url="https://t3.storage.dev",
    region="auto",
    access_key_id="...",
    secret_access_key="...",
)


def collate(records):
    ...


loader = training_data.split_data_loader(
    "s3://my-training-bucket/path/to/data-manifest-v1.json",
    collate_fn=collate,
    dataloader_args={"batch_size": 32, "num_workers": 4, ...},
)

train_batches = loader.load("train")
valid_batches = loader.load("valid")
easy_batches = loader.load("train", curriculum_stage="easy")
early_curriculum_batches = loader.load("train", curriculum_stages=["easy", "medium"])
infinite_train_batches = loader.load("train", infinite_iter=True)
```

`curriculum_stages` selects the union of the named stages. It does not guarantee
records are yielded in stage order, especially when `num_workers` is greater
than zero.

Storage reads recover from transient S3/Tigris connection and body failures by
opening a fresh client and resuming immutable shard streams at the last received
byte. Point reads and writes retry the complete operation at the same URI.
Missing objects, authentication failures, and invalid data still fail closed.

## Load joined OLMo 3 representation examples

The representation loader verifies the pinned final manifest, the ETag-bound
indexes, the index-bound artifact manifest, and each selected artifact. It
then joins prepared tokens to the requested teacher projection rows.

```python
examples = training_data.iter_representation_examples(
    manifest_uri,
    manifest_sha256,
    split="analysis",
    sources=["broad_replay", "dolmino"],
    per_source_count=64,
    teacher_layer=31,
    state_kind="post_mlp",
)
```

The loader processes source groups in caller order. Within each source, it
selects and yields the first `per_source_count` examples by ascending
`block_id`. Each `RepresentationExample` owns its returned tensors, so no
Safetensors mapping or full artifact remains live after the loader copies the
selected rows.

The package uses a 10-second connect timeout, a 60-second read-inactivity
timeout, and a 10-minute overall request timeout. Obstore's internal retry
window is deliberately short; the package owns the longer 15-minute
no-progress recovery window so a failed connection pool can be discarded.

When using worker processes, leave PyTorch `DataLoader(timeout=0)` unless the
training runtime has a specific worker watchdog. A positive DataLoader timeout
must be longer than the storage recovery window plus normal shard processing;
a value such as 120 seconds can terminate a healthy worker while it is retrying
a transient object-store outage.

## Derive data

Use the manifest's format to stream source shards when generating a derived
dataset:

```python
from syvain_training_data import iter_shard

for shard in manifest.splits["train"].shards:
    for record in iter_shard(
        manifest.data_format,
        shard,
        storage_config=storage_config,
    ):
        ...
```

## Save data

```python
from concurrent.futures import ProcessPoolExecutor

from syvain_training_data import SyvainTrainingData

def generate_data(split, curriculum_stage, shard_id):
    ...

def save_shard(job):
    saver, split, curriculum_stage, metadata, shard_id = job
    records = generate_data(split, curriculum_stage, shard_id)
    saver.save(
        split,
        curriculum_stage,
        records,
        curriculum_metadata=metadata,
        shard_id=str(shard_id),
    )


training_data = SyvainTrainingData(
    s3_base_url="https://t3.storage.dev",
    region="auto",
    access_key_id="...",
    secret_access_key="...",
)

saver = training_data.dataset_saver(
    "s3://my-training-bucket/path/to/dataset/data-manifest-v1.json",
)

jobs = [
    (saver, "train", stage["name"], stage, shard_id)
    for stage in [
        {"name": "easy", "family": "arithmetic", "weight": 1.0},
        {"name": "medium", "family": "control", "weight": 2.0},
        {"name": "hard", "family": "composition", "weight": 3.0},
    ]
    for shard_id in range(32)
] + [
    (saver, "valid", None, None, shard_id) for shard_id in range(4)
] + [
    (saver, "test", None, None, shard_id) for shard_id in range(4)
]

with ProcessPoolExecutor(max_workers=8) as pool:
    list(pool.map(save_shard, jobs))

manifest = saver.commit_manifest()
```

Each deterministic shard writes a completion descriptor after its data object.
A restarted saver verifies that descriptor and reuses the shard without
consuming `records`. The data manifest is written last as the publication
marker; committing the same completed publication again is idempotent. Call
`saver.recover(split, curriculum_stage, shard_id=...)` to inspect a completed
shard explicitly without constructing a records iterator.

Deterministic Parquet and framed MessagePack shards use an atomic conditional
upload and therefore must each be smaller than 5 GiB. Use more logical shard
IDs when generating a larger derived dataset.

## Copy a manifest

```python
from syvain_training_data import SyvainTrainingData

training_data = SyvainTrainingData(
    s3_base_url="https://t3.storage.dev",
    region="auto",
    access_key_id="...",
    secret_access_key="...",
)

manifest = training_data.load_manifest("s3://my-training-bucket/shared/data-manifest-v1.json")

# Do modifications if needed

training_data.save_manifest("s3://my-training-bucket/new-run/data-manifest-v1.json", manifest)
```
