Training with Grain

Grain’s random-access source contract requires two methods: __len__() and integer __getitem__(). SDS datasets provide both. Use an application adapter to convert generic records into the array tree expected by a model.

Install Grain in the training application; it is intentionally not an SDS dependency:

uv add 'grain>=0.2.17,<0.3'

import threading

import grain
import numpy as np
import signal_dataset as sds


class SDSSource(grain.sources.RandomAccessDataSource):
    def __init__(self, uri: str) -> None:
        self._uri = uri
        self._dataset = None
        self._lock = threading.Lock()

    def _open(self) -> sds.Dataset:
        # Double-checked: the lock guards the one-time open, not the read.
        if self._dataset is None:
            with self._lock:
                if self._dataset is None:
                    self._dataset = sds.open(self._uri)
        return self._dataset

    def __len__(self) -> int:
        return len(self._open())

    def __getitem__(self, index: int) -> dict[str, np.ndarray]:
        record = self._open()[index]
        return {
            "inputs": record["samples"].data,
            "target": np.asarray(record.metadata["label"], dtype=np.int32),
        }

    def __getstate__(self) -> dict[str, object]:
        return {"_uri": self._uri}

    def __setstate__(self, state: dict[str, object]) -> None:
        self._uri = str(state["_uri"])
        self._dataset = None
        self._lock = threading.Lock()

    def __repr__(self) -> str:
        return f"SDSSource(uri={self._uri!r})"

Build the training pipeline with Grain-owned transforms:

def train_step(inputs: np.ndarray, target: np.ndarray) -> None:
    """Stand-in for the real optimizer step."""
    assert inputs.shape[0] == target.shape[0]


source = SDSSource("captures.sds")
pipeline = (
    grain.MapDataset.source(source)
    .shuffle(seed=20260831)
    .batch(2, drop_remainder=True)
    .repeat(num_epochs=1)
    .to_iter_dataset(
        read_options=grain.ReadOptions(num_threads=16, prefetch_buffer_size=2)
    )
)

for batch in pipeline:
    train_step(batch["inputs"], batch["target"])

Grain provides shuffle, sharding, mapping, random transforms, batching, repetition, prefetch, and iterator checkpointing. SDS remains responsible only for immutable indexed records. Keep model conversion in the adapter; do not pass an SDS Record directly to Grain batching.

The lock guards the one-time open, and nothing else. Holding it across the read instead would serialize every one of Grain’s reader threads – ReadOptions.num_threads defaults to 16 – and make that setting inert. Grain asks that __getitem__ be thread-safe and deterministic, meaning a given index returns the same value; concurrent reads of distinct indices from an immutable dataset already satisfy that.

Pickling drops the live SDS reader and lock; each worker process reopens the same immutable dataset URI. The stable __repr__ supports Grain DataLoader checkpoint validation.

For annotations, pin the set once in __init__ or _open, then read the signal record and annotation at the same ordinal. Grain slicing can shard a source without copying SDS data:

worker_index, worker_count = 0, 2
worker_source = grain.MapDataset.source(source)[worker_index::worker_count]