o41-training-data

Internal 041 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

uv add o41-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.

from pathlib import Path

import torch

from o41_training_data import (
    BatchStream,
    BatchStreamConfig,
    Checkpoint,
    FileSource,
    ManifestSource,
    ProjectedField,
    Shuffle,
    Stage,
    TrainingData,
)

training_data = TrainingData(
    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, f32, and f64 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. An f64 field keeps double precision through projection and the NumPy view.

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:

from o41_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 TrainingData, 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:

from o41_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

from concurrent.futures import ProcessPoolExecutor

from o41_training_data import TrainingData

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 = TrainingData(
    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:

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

from o41_training_data import TrainingData

training_data = TrainingData(
    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)