Metadata-Version: 2.4
Name: syvain-training-data
Version: 0.0.354
Requires-Dist: msgpack>=1.1.2,<2.0.0
Requires-Dist: numpy>=2.0.0
Requires-Dist: obstore>=0.11.0,<0.12.0
Requires-Dist: pyarrow>=23.0.1,<24.0.0
Requires-Dist: pydantic>=2.13.4
Requires-Dist: zstandard>=0.25.0,<0.26.0
Requires-Dist: pytest>=8.0.0 ; extra == 'dev'
Requires-Dist: ruff>=0.15.12 ; extra == 'dev'
Requires-Dist: ty>=0.0.34 ; extra == 'dev'
Provides-Extra: dev
Summary: Syvain training data manifest, streaming, and saving utilities
Requires-Python: >=3.14, <3.15
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM

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

The package ships a compiled Rust core with wheels for Linux x86-64, Linux
aarch64 and macOS arm64. It depends on PyArrow and NumPy and stays independent
of any training framework.

## Stream batches

A batch stream is a deterministic, resumable iterator over ordered curriculum
stages. Each stage is a weighted mix of sources, and a source is either a
selection from a data manifest or a plain list of shard files. Records are
projected to the declared fixed-shape columns; everything else is skipped
before decoding finishes.

```python
from pathlib import Path

import torch

from syvain_training_data import (
    BatchStream,
    BatchStreamConfig,
    Checkpoint,
    FileSource,
    ManifestSource,
    ProjectedField,
    Shuffle,
    Stage,
    SyvainTrainingData,
)

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

config = BatchStreamConfig(
    fields=[
        ProjectedField(name="tokens", source="input_ids", dtype="i32", length=2048, pad=0),
        ProjectedField(name="loss_mask", source="loss_mask", dtype="bool", length=2048, pad=False),
    ],
    stages=[
        Stage(
            sources=[
                ManifestSource(
                    manifest_uri="s3://my-training-bucket/path/to/data-manifest-v1.json",
                    split="train",
                    curriculum_stages=["easy", "medium"],
                    weight=3.0,
                ),
                FileSource(
                    files=["s3://my-training-bucket/extra/part-000.jsonl.gz"],
                    encoding="jsonl.gz",
                    weight=1.0,
                ),
            ],
            batch_size=32,
            shuffle=Shuffle(buffer_rows=4096),
            infinite=True,
        ),
        Stage(
            sources=[
                ManifestSource(
                    manifest_uri="s3://my-training-bucket/path/to/data-manifest-v1.json",
                    split="train",
                    curriculum_stages=["hard"],
                )
            ],
            batch_size=32,
            shuffle=Shuffle(buffer_rows=4096),
            infinite=True,
        ),
    ],
    seed=1,
    threads=8,
    prefetch_batches=4,
)

STREAM_CHECKPOINT = Path(f"checkpoints/stream-rank{config.rank}.json")


def open_stream() -> tuple[BatchStream, int]:
    """Continue from the last checkpoint when there is one, otherwise start fresh."""
    if STREAM_CHECKPOINT.exists():
        checkpoint = Checkpoint.from_json(STREAM_CHECKPOINT.read_text())
        return training_data.resume_batch_stream(checkpoint), checkpoint.batches_yielded
    return training_data.batch_stream(config), 0


stream, step = open_stream()
with stream:
    for batch in stream:
        step += 1
        arrays = batch.arrays()  # {"tokens": int32 (32, 2048), "loss_mask": uint8 (32, 2048)}
        tokens = torch.from_numpy(arrays["tokens"])
        loss_mask = torch.from_numpy(arrays["loss_mask"]).bool()
        ...  # forward, backward, optimizer step
        if step % 1000 == 0:
            save_model_checkpoint(...)
            STREAM_CHECKPOINT.write_text(stream.checkpoint().to_json())
        if batch.stage == 0 and ready_for_hard_examples:
            stream.advance(1)
```

Both `batch_stream` and `resume_batch_stream` open eagerly and return the same
`BatchStream`, so the choice belongs in one small function and the loop does
not care which path produced the stream. Open-time failures, such as a
missing shard or a projected field that no record carries, raise at that call.

Write the stream checkpoint together with the model checkpoint, after the
optimizer step that consumed the batch. A checkpoint taken after a batch was
delivered treats that batch as consumed, so the resumed stream yields the
next one. Drive the cadence from a counter that every rank advances in step,
such as the local batch count above, which `checkpoint.batches_yielded`
restores. `batch.batch_index` is the global batch number: it is
rank-independent and each rank only ever sees its own residue class, so it is
not a per-rank cadence. A checkpoint belongs to the rank that wrote it, since
the resolved config carries `rank` and `world_size`; write one file per rank.

Order is a pure function of the config, the seed and the checkpoint. Ranks
share one global batch sequence and take interleaved slices of it, so set
`rank` and `world_size` per process. A checkpoint holds the resolved config
without credentials, the loop cursor with shuffle coordinates, and the
counters `stage`, `batches_yielded`, `rows_yielded` and per-source `epochs`.
Because the checkpoint stores the resolved file lists, a manifest republished
after the checkpoint does not change what the resumed stream reads.

`curriculum_stages` in one `ManifestSource` selects the union of those stages in
the listed order. Ordered curricula are separate `Stage` entries; the run
decides when to call `advance`, which takes effect at the next batch boundary.
Rows already in the shuffle buffer drain naturally, so the first batches after
an advance can still carry rows from the previous stage. `batch.stage` reports
the stage that was active when the batch was produced. A same-stage advance is
a no-op, and a finite stage ends the stream instead of advancing on its own.
Every rank yields the same number of batches: a rank completes the round of
`world_size` global batches before yielding its own, and a finite stage drops
a final round that cannot be filled. With several ranks, call `advance` after
the same number of local batches on every rank; the ranks then switch stages
from one identical cursor and the new stage stays partitioned between them.
`close` and `advance` wait for the shard read in flight to finish, including
its retry budget during a storage outage; they do not interrupt it.

Every column is a contiguous numeric buffer: `i32`, `i64` and `f32` as
declared, `bool` as `uint8` values `0` and `1`. `batch.arrays()` returns NumPy
views shaped `(rows, *dims)` over the Arrow buffers without copying, and
`batch.columns` is the underlying `pyarrow.RecordBatch`. Fields declare one
`length` or a nested `shape`; each axis truncates or pads with `pad`.

The core validates the projected fields against one record of every source at
open, bounds its working set from the manifest row counts and the declared
shapes, and rejects a config whose bound exceeds `memory_limit_bytes`, which
defaults to half of the host's physical memory. Reads recover from transient
object-store failures by resuming at the last received byte for up to 64
consecutive failures or 15 minutes without progress; missing objects,
authentication failures and malformed shards fail closed.

Sources can also be local paths, for example when a training job copies its
shards to local disk before the run. A config whose sources are all local needs
no storage settings:

```python
from syvain_training_data import open_batch_stream, resume_batch_stream

config = BatchStreamConfig(
    fields=[ProjectedField(name="tokens", source="input_ids", dtype="i32", length=2048)],
    stages=[
        Stage(
            sources=[
                ManifestSource(manifest_uri="/data/run/data-manifest-v1.json", split="train"),
                FileSource(files=["/data/extra/part-000.jsonl.gz"], encoding="jsonl.gz"),
            ],
            batch_size=32,
            infinite=True,
        )
    ],
    seed=1,
)

with open_batch_stream(config) as stream:
    ...

with resume_batch_stream(checkpoint) as stream:
    ...
```

A local manifest lists its shards by local path. Mixing local and `s3://`
sources in one stream goes through `SyvainTrainingData`, which supplies the
storage settings; opening an `s3://` source without them fails at open.

`training_data.resolve_batch_stream_config(config)` returns the resolved
loader config a stream would open. It is the same document a checkpoint
stores and is convenient for inspecting which shards a run will read.

## 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.

To publish cumulative curricula without duplicating records, save the disjoint
shards once and register each ordered selection:

```python
first = saver.save("train", None, first_records, shard_id="delta-001")
second = saver.save("train", None, second_records, shard_id="delta-002")
saver.register_curriculum("train", "small", [first])
saver.register_curriculum("train", "large", [first, second])
candidate = saver.build_manifest()
manifest = saver.commit_manifest()
```

`register_curriculum` accepts only shards already saved or recovered in that
split. It preserves the supplied order and keeps every physical shard once in
the root split. Use `curriculum_metadata` for selector metadata. Repeating the
same registration is idempotent; conflicting definitions and repeated shards
are rejected. After a process restart, recover the shards and register the
same selections before building or committing the manifest. A registration
also cannot replace a different stage inferred from `save` calls.

Root splits use URI order by default. To prescribe their traversal order, call
`saver.register_split_order("train", [second, first])` after saving or recovering
all shards. The list must contain every shard in that split exactly once with
its original metadata. This changes only root order; curriculum selections keep
their own order. `build_manifest()` and `commit_manifest()` preserve the registered
order. Repeating the same registration is idempotent; a different order conflicts.
Adding a shard afterward makes the registration incomplete and prevents manifest
publication. After restart, recover all shards and register the order again.

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

