syvain-training-data 0.0.118__tar.gz

syvain_training_data-0.0.118/PKG-INFO

@@ -0,0 +1,121 @@
+Metadata-Version: 2.3
+Name: syvain-training-data
+Version: 0.0.118
+Summary: Syvain training data manifest, loading, and saving utilities
+Requires-Dist: obstore>=0.10.1
+Requires-Dist: pydantic>=2.13.3
+Requires-Dist: torch>=2.0.0,<2.11
+Requires-Dist: typing-extensions>=4.15.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'
+Requires-Python: >=3.10
+Provides-Extra: dev
+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
+```
+
+## 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")
+infinite_train_batches = loader.load("train", infinite_iter=True)
+```
+
+## 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, shard_id = job
+    records = generate_data(split, curriculum_stage, shard_id)
+    saver.save(split, curriculum_stage, records)
+
+
+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, shard_id)
+    for stage in ["easy", "medium", "hard"]
+    for shard_id in range(32)
+] + [
+    (saver, "valid", None, shard_id) for shard_id in range(4)
+] + [
+    (saver, "test", 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()
+```
+
+## 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)
+```

syvain_training_data-0.0.118/README.md

@@ -0,0 +1,106 @@
+# 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
+```
+
+## 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")
+infinite_train_batches = loader.load("train", infinite_iter=True)
+```
+
+## 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, shard_id = job
+    records = generate_data(split, curriculum_stage, shard_id)
+    saver.save(split, curriculum_stage, records)
+
+
+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, shard_id)
+    for stage in ["easy", "medium", "hard"]
+    for shard_id in range(32)
+] + [
+    (saver, "valid", None, shard_id) for shard_id in range(4)
+] + [
+    (saver, "test", 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()
+```
+
+## 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)
+```

syvain_training_data-0.0.118/pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+name = "syvain-training-data"
+version = "0.0.118"
+description = "Syvain training data manifest, loading, and saving utilities"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "obstore>=0.10.1",
+    "pydantic>=2.13.3",
+    "torch>=2.0.0,<2.11",
+    "typing-extensions>=4.15.0",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=8.0.0", "ruff>=0.15.12", "ty>=0.0.34"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[build-system]
+requires = ["uv_build>=0.11.9,<0.12"]
+build-backend = "uv_build"

syvain_training_data-0.0.118/src/syvain_training_data/__init__.py

@@ -0,0 +1,23 @@
+from syvain_training_data.client import SyvainTrainingData
+from syvain_training_data.loader import SplitDataLoader
+from syvain_training_data.manifest import (
+    DataManifest,
+    JsonObject,
+    JsonValue,
+    ManifestCurriculumStage,
+    ManifestDatasetEntry,
+    ManifestShard,
+)
+from syvain_training_data.saver import DatasetSaver
+
+__all__ = [
+    "DataManifest",
+    "DatasetSaver",
+    "JsonObject",
+    "JsonValue",
+    "ManifestCurriculumStage",
+    "ManifestDatasetEntry",
+    "ManifestShard",
+    "SplitDataLoader",
+    "SyvainTrainingData",
+]

syvain_training_data-0.0.118/src/syvain_training_data/client.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any, TypeVar
+
+from syvain_training_data.loader import SplitDataLoader
+from syvain_training_data.manifest import (
+    DataManifest,
+    manifest_json_bytes,
+    parse_manifest_json,
+)
+from syvain_training_data.saver import DatasetSaver
+from syvain_training_data.storage import (
+    StorageConfig,
+    put_object_bytes,
+    read_object_bytes,
+)
+
+BatchT = TypeVar("BatchT")
+
+
+class SyvainTrainingData:
+    def __init__(
+        self,
+        *,
+        s3_base_url: str,
+        region: str,
+        access_key_id: str,
+        secret_access_key: str,
+        virtual_hosted_style_request: bool = False,
+    ) -> None:
+        self._storage_config = StorageConfig(
+            s3_base_url=s3_base_url,
+            region=region,
+            access_key_id=access_key_id,
+            secret_access_key=secret_access_key,
+            virtual_hosted_style_request=virtual_hosted_style_request,
+        )
+
+    def load_manifest(self, manifest_uri: str) -> DataManifest:
+        body = read_object_bytes(
+            manifest_uri,
+            storage_config=self._storage_config,
+        )
+        return parse_manifest_json(body, source=manifest_uri)
+
+    def save_manifest(
+        self,
+        manifest_uri: str,
+        manifest: DataManifest | Mapping[str, object],
+    ) -> None:
+        put_object_bytes(
+            manifest_uri,
+            manifest_json_bytes(manifest),
+            storage_config=self._storage_config,
+            attributes={"Content-Type": "application/json"},
+        )
+
+    def split_data_loader(
+        self,
+        manifest_uri: str,
+        collate_fn: Callable[[Sequence[Mapping[str, Any]]], BatchT],
+        dataloader_args: Mapping[str, Any] | None = None,
+    ) -> SplitDataLoader[BatchT]:
+        return SplitDataLoader(
+            manifest=self.load_manifest(manifest_uri),
+            storage_config=self._storage_config,
+            collate_fn=collate_fn,
+            dataloader_args={} if dataloader_args is None else dict(dataloader_args),
+        )
+
+    def dataset_saver(
+        self,
+        manifest_uri: str,
+    ) -> DatasetSaver:
+        return DatasetSaver(
+            storage_config=self._storage_config,
+            manifest_uri=manifest_uri,
+        )
+
+
+__all__ = ["SyvainTrainingData"]

syvain_training_data-0.0.118/src/syvain_training_data/gzjsonl.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+import codecs
+import gzip
+import json
+import logging
+import zlib
+from collections.abc import Iterable, Iterator, Mapping
+from io import BytesIO
+from typing import Any
+
+from torch.utils.data import get_worker_info
+
+from syvain_training_data.manifest import ManifestShard
+from syvain_training_data.storage import (
+    StorageConfig,
+    parse_s3_uri,
+    store_for_uri,
+)
+
+LOGGER = logging.getLogger(__name__)
+
+def records_to_body(records: Iterable[Mapping[str, Any]]) -> tuple[bytes, int]:
+    output = BytesIO()
+    record_count = 0
+    with gzip.GzipFile(
+        fileobj=output, mode="wb", compresslevel=6, mtime=0
+    ) as gzip_file:
+        for record in records:
+            line = json.dumps(record, separators=(",", ":")).encode("utf-8")
+            gzip_file.write(line)
+            gzip_file.write(b"\n")
+            record_count += 1
+    return output.getvalue(), record_count
+
+
+def iter_shard(
+    shard: ManifestShard,
+    *,
+    storage_config: StorageConfig,
+) -> Iterator[Mapping[str, Any]]:
+    source = shard.uri
+    parsed_uri = parse_s3_uri(source)
+    worker = get_worker_info()
+    worker_id = "main" if worker is None else str(worker.id)
+    yielded_records = 0
+    LOGGER.info(
+        "Worker %s reading shard %s with %s expected records",
+        worker_id,
+        source,
+        shard.records,
+    )
+
+    try:
+        store, _ = store_for_uri(source, storage_config=storage_config)
+        response = store.get(parsed_uri.key)
+    except Exception as exc:
+        raise RuntimeError(f"Unreadable shard {source}") from exc
+
+    decompressor = zlib.decompressobj(wbits=16 + zlib.MAX_WBITS)
+    decoder = codecs.getincrementaldecoder("utf-8")()
+    pending_line = ""
+    line_number = 0
+
+    try:
+        for chunk in response.stream(min_chunk_size=1024 * 1024):
+            decoded = decoder.decode(decompressor.decompress(chunk), final=False)
+            if not decoded:
+                continue
+
+            pending_line += decoded
+            lines = pending_line.split("\n")
+            pending_line = lines.pop()
+            for line in lines:
+                line_number += 1
+                if not line.strip():
+                    continue
+                record = _parse_record_line(
+                    line,
+                    source=source,
+                    line_number=line_number,
+                )
+                if record is not None:
+                    yielded_records += 1
+                    yield record
+
+        pending_line += decoder.decode(decompressor.flush(), final=True)
+    except (zlib.error, UnicodeDecodeError) as exc:
+        raise RuntimeError(
+            f"Unreadable gzip-compressed UTF-8 JSONL shard {source}"
+        ) from exc
+
+    if not decompressor.eof or decompressor.unused_data:
+        raise RuntimeError(
+            "Invalid gzip-compressed UTF-8 JSONL shard "
+            f"{source}: eof={decompressor.eof} "
+            f"unused_data={bool(decompressor.unused_data)}"
+        )
+
+    if pending_line.strip():
+        line_number += 1
+        record = _parse_record_line(
+            pending_line,
+            source=source,
+            line_number=line_number,
+        )
+        if record is not None:
+            yielded_records += 1
+            yield record
+
+    LOGGER.info(
+        "Worker %s finished shard %s with %s/%s yielded records",
+        worker_id,
+        source,
+        yielded_records,
+        shard.records,
+    )
+
+
+def _parse_record_line(
+    line: str,
+    *,
+    source: str,
+    line_number: int,
+) -> Mapping[str, Any] | None:
+    try:
+        raw_record = json.loads(line)
+    except json.JSONDecodeError as exc:
+        LOGGER.error(
+            "Skipping invalid JSON record in %s:%s: %s", source, line_number, exc
+        )
+        return None
+
+    if not isinstance(raw_record, Mapping):
+        LOGGER.error("Skipping non-object JSON record in %s:%s", source, line_number)
+        return None
+
+    return raw_record
+
+
+__all__ = ["iter_shard", "records_to_body"]

syvain_training_data-0.0.118/src/syvain_training_data/loader.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any, Generic, TypeVar, cast, overload
+
+from torch.utils.data import DataLoader, IterableDataset, get_worker_info
+from typing_extensions import Literal
+
+from syvain_training_data import gzjsonl
+from syvain_training_data.manifest import (
+    DataManifest,
+    DataManifestDataFormat,
+    ManifestShard,
+    curriculum_stage_entry,
+    manifest_split_entry,
+)
+from syvain_training_data.storage import StorageConfig
+
+LOGGER = logging.getLogger(__name__)
+
+BatchT = TypeVar("BatchT")
+
+
+@dataclass(frozen=True)
+class SplitDataLoader(Generic[BatchT]):
+    manifest: DataManifest
+    storage_config: StorageConfig
+    collate_fn: Callable[[Sequence[Mapping[str, Any]]], BatchT]
+    dataloader_args: Mapping[str, Any]
+
+    @overload
+    def load(
+        self,
+        split: str,
+        curriculum_stage: str | None = None,
+        infinite_iter: Literal[False] = False,
+    ) -> DataLoader[BatchT]: ...
+
+    @overload
+    def load(
+        self,
+        split: str,
+        curriculum_stage: str | None = None,
+        infinite_iter: Literal[True] = True,
+    ) -> Iterator[BatchT]: ...
+
+    def load(
+        self,
+        split: str,
+        curriculum_stage: str | None = None,
+        infinite_iter: bool = False,
+    ) -> DataLoader[BatchT] | Iterator[BatchT]:
+        split_entry = manifest_split_entry(self.manifest, split)
+        if curriculum_stage is not None:
+            split_entry = curriculum_stage_entry(split_entry, curriculum_stage)
+
+        dataset = _ManifestShardDataset(
+            shards=split_entry.shards,
+            data_format=self.manifest.data_format,
+            storage_config=self.storage_config,
+        )
+        loader_kwargs = dict(self.dataloader_args)
+        if "collate_fn" in loader_kwargs:
+            raise ValueError(
+                "Pass collate_fn to split_data_loader, not dataloader_args"
+            )
+        loader_kwargs["collate_fn"] = self.collate_fn
+        loader_kwargs = _cap_worker_args(
+            loader_kwargs,
+            shard_count=dataset.shard_count,
+            split=split,
+        )
+
+        loader = cast(DataLoader[BatchT], DataLoader(dataset, **loader_kwargs))
+        if not infinite_iter:
+            return loader
+        return _iter_dataloader_forever(loader, split=split)
+
+
+class _ManifestShardDataset(IterableDataset[Mapping[str, Any]]):
+    def __init__(
+        self,
+        *,
+        shards: Sequence[ManifestShard],
+        data_format: DataManifestDataFormat,
+        storage_config: StorageConfig,
+    ) -> None:
+        self._shards = tuple(shards)
+        self._data_format = data_format
+        self._storage_config = storage_config
+
+    @property
+    def shard_count(self) -> int:
+        return len(self._shards)
+
+    def __iter__(self) -> Iterator[Mapping[str, Any]]:
+        for shard in _worker_shards(self._shards):
+            if self._data_format == "jsonl.gz":
+                yield from gzjsonl.iter_shard(
+                    shard,
+                    storage_config=self._storage_config,
+                )
+            else:
+                raise ValueError(f"Unknown data format: {self._data_format}")
+
+
+def _cap_worker_args(
+    loader_kwargs: dict[str, Any],
+    *,
+    shard_count: int,
+    split: str,
+) -> dict[str, Any]:
+    requested_workers = int(loader_kwargs.get("num_workers", 0))
+    if requested_workers < 0:
+        raise ValueError("dataloader_args['num_workers'] must be non-negative")
+
+    num_workers = min(requested_workers, shard_count)
+    loader_kwargs["num_workers"] = num_workers
+    if num_workers < requested_workers:
+        LOGGER.debug(
+            "Capping %s dataloader workers from %s to %s because the split has %s shards",
+            split,
+            requested_workers,
+            num_workers,
+            shard_count,
+        )
+
+    if num_workers == 0:
+        loader_kwargs.pop("persistent_workers", None)
+        loader_kwargs.pop("prefetch_factor", None)
+        loader_kwargs.pop("multiprocessing_context", None)
+        if int(loader_kwargs.get("timeout", 0)) > 0:
+            loader_kwargs.pop("timeout", None)
+    return loader_kwargs
+
+
+def _worker_shards(shards: Sequence[ManifestShard]) -> Sequence[ManifestShard]:
+    worker = get_worker_info()
+    if worker is None:
+        return shards
+    return shards[worker.id :: worker.num_workers]
+
+
+def _iter_dataloader_forever(
+    loader: DataLoader[BatchT],
+    *,
+    split: str,
+) -> Iterator[BatchT]:
+    while True:
+        yielded_batch = False
+        for batch in loader:
+            yielded_batch = True
+            yield batch
+        if not yielded_batch:
+            raise RuntimeError(f"{split} split produced no batches")
+
+
+__all__ = ["SplitDataLoader"]

syvain_training_data-0.0.118/src/syvain_training_data/manifest.py

@@ -0,0 +1,198 @@
+from __future__ import annotations
+
+import json
+from typing import Literal, Mapping
+from urllib.parse import urlparse
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator
+from typing_extensions import TypeAliasType
+
+JsonValue = TypeAliasType(
+    "JsonValue",
+    str | int | float | bool | None | list["JsonValue"] | dict[str, "JsonValue"],
+)
+JsonObject = TypeAliasType("JsonObject", dict[str, JsonValue])
+DataManifestDataFormat = TypeAliasType("DataManifestDataFormat", Literal["jsonl.gz"])
+
+MANIFEST_VERSION = 1
+
+MODEL_CONFIG = ConfigDict(extra="forbid", frozen=True, strict=True)
+
+
+class ManifestShard(BaseModel):
+    model_config = MODEL_CONFIG
+
+    uri: str = Field(min_length=1)
+    records: int = Field(gt=0)
+    size_bytes: int = Field(gt=0)
+    etag: str = ""
+
+    @field_validator("uri")
+    @classmethod
+    def validate_uri(cls, value: str) -> str:
+        parsed = urlparse(value)
+        if parsed.scheme != "s3" or not parsed.netloc or not parsed.path.lstrip("/"):
+            raise ValueError(f"Expected non-empty s3://bucket/key URI, got: {value}")
+        return value
+
+
+class ManifestCurriculumStage(BaseModel):
+    model_config = MODEL_CONFIG
+
+    records: int = Field(gt=0)
+    shards: list[ManifestShard] = Field(min_length=1)
+
+    @field_validator("shards")
+    @classmethod
+    def validate_shard_records(cls, value: list[ManifestShard]) -> list[ManifestShard]:
+        if sum(shard.records for shard in value) <= 0:
+            raise ValueError("stage shards must contain records")
+        return value
+
+    def model_post_init(self, __context: object) -> None:
+        shard_records = sum(shard.records for shard in self.shards)
+        if self.records != shard_records:
+            raise ValueError(
+                f"stage records {self.records} do not match shard records "
+                f"{shard_records}"
+            )
+
+
+class ManifestDatasetEntry(BaseModel):
+    model_config = MODEL_CONFIG
+
+    records: int = Field(gt=0)
+    shards: list[ManifestShard] = Field(min_length=1)
+    curriculum: dict[str, ManifestCurriculumStage] = Field(default_factory=dict)
+
+    @field_validator("curriculum")
+    @classmethod
+    def validate_curriculum(
+        cls, value: dict[str, ManifestCurriculumStage]
+    ) -> dict[str, ManifestCurriculumStage]:
+        for stage_name in value:
+            if not stage_name:
+                raise ValueError("curriculum stage names must not be empty")
+        return value
+
+    def model_post_init(self, __context: object) -> None:
+        shard_records = sum(shard.records for shard in self.shards)
+        if self.records != shard_records:
+            raise ValueError(
+                f"split records {self.records} do not match shard records "
+                f"{shard_records}"
+            )
+
+
+class DataManifest(BaseModel):
+    model_config = MODEL_CONFIG
+
+    version: Literal[1]
+    data_format: DataManifestDataFormat
+    splits: dict[str, ManifestDatasetEntry] = Field(min_length=1)
+
+    @field_validator("splits")
+    @classmethod
+    def validate_splits(
+        cls, value: dict[str, ManifestDatasetEntry]
+    ) -> dict[str, ManifestDatasetEntry]:
+        for split in value:
+            if not split:
+                raise ValueError("split names must not be empty")
+        return value
+
+
+def manifest_split_entry(manifest: DataManifest, split: str) -> ManifestDatasetEntry:
+    try:
+        return manifest.splits[split]
+    except KeyError as exc:
+        available = ", ".join(sorted(manifest.splits))
+        raise ValueError(
+            f"Unknown split {split!r}; available splits: {available}"
+        ) from exc
+
+
+def curriculum_stage_entry(
+    split_entry: ManifestDatasetEntry,
+    stage_name: str,
+) -> ManifestCurriculumStage:
+    try:
+        return split_entry.curriculum[stage_name]
+    except KeyError as exc:
+        available = ", ".join(sorted(split_entry.curriculum))
+        raise ValueError(
+            f"Unknown curriculum stage {stage_name!r}; available stages: {available}"
+        ) from exc
+
+
+def parse_manifest_json(body: bytes, *, source: str) -> DataManifest:
+    try:
+        raw_manifest = json.loads(body.decode("utf-8"))
+        return DataManifest.model_validate(raw_manifest)
+    except UnicodeDecodeError as exc:
+        raise ValueError(f"Data manifest is not valid UTF-8: {source}") from exc
+    except json.JSONDecodeError as exc:
+        raise ValueError(f"Data manifest is not valid JSON: {source}") from exc
+    except ValidationError as exc:
+        raise ValueError(f"Data manifest has invalid shape: {source}") from exc
+
+
+def manifest_json_bytes(manifest: DataManifest | Mapping[str, object]) -> bytes:
+    typed_manifest = (
+        manifest
+        if isinstance(manifest, DataManifest)
+        else DataManifest.model_validate(manifest)
+    )
+    return (typed_manifest.model_dump_json(indent=2, exclude_none=True) + "\n").encode(
+        "utf-8"
+    )
+
+
+def dataset_entry_from_shards(
+    shards: list[ManifestShard],
+    *,
+    curriculum: Mapping[str, ManifestCurriculumStage] | None = None,
+) -> ManifestDatasetEntry:
+    return ManifestDatasetEntry(
+        records=sum(shard.records for shard in shards),
+        shards=shards,
+        curriculum={} if curriculum is None else dict(curriculum),
+    )
+
+
+def curriculum_stage_from_shards(
+    shards: list[ManifestShard],
+) -> ManifestCurriculumStage:
+    return ManifestCurriculumStage(
+        records=sum(shard.records for shard in shards),
+        shards=shards,
+    )
+
+
+def manifest_from_splits(
+    splits: Mapping[str, ManifestDatasetEntry],
+    *,
+    data_format: DataManifestDataFormat = "jsonl.gz",
+) -> DataManifest:
+    return DataManifest(
+        version=MANIFEST_VERSION, data_format=data_format, splits=dict(splits)
+    )
+
+
+__all__ = [
+    "DataManifest",
+    "DataManifestDataFormat",
+    "JsonObject",
+    "JsonValue",
+    "MANIFEST_VERSION",
+    "ManifestCurriculumStage",
+    "ManifestDatasetEntry",
+    "ManifestShard",
+    "curriculum_stage_entry",
+    "curriculum_stage_from_shards",
+    "dataset_entry_from_shards",
+    "manifest_from_splits",
+    "manifest_json_bytes",
+    "manifest_split_entry",
+    "parse_manifest_json",
+]

syvain_training_data-0.0.118/src/syvain_training_data/py.typed

@@ -0,0 +1 @@
+

syvain_training_data-0.0.118/src/syvain_training_data/saver.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+import multiprocessing
+import uuid
+from collections.abc import Callable, Iterable, Mapping
+from typing import Any, TypedDict, cast
+
+from syvain_training_data import gzjsonl
+from syvain_training_data.manifest import (
+    DataManifest,
+    DataManifestDataFormat,
+    ManifestCurriculumStage,
+    ManifestDatasetEntry,
+    ManifestShard,
+    curriculum_stage_from_shards,
+    manifest_from_splits,
+)
+from syvain_training_data.storage import (
+    StorageConfig,
+    join_s3_uri,
+    parse_s3_uri,
+    put_object_bytes,
+    safe_path_segment,
+)
+
+
+class _SavedShard(TypedDict):
+    split: str
+    curriculum_stage: str | None
+    shard: dict[str, object]
+
+
+_RecordsToBody = Callable[[Iterable[Mapping[str, Any]]], tuple[bytes, int]]
+
+
+class DatasetSaver:
+    def __init__(
+        self,
+        *,
+        storage_config: StorageConfig,
+        manifest_uri: str,
+        data_format: DataManifestDataFormat = "jsonl.gz",
+    ) -> None:
+        manifest_uri = manifest_uri.strip()
+        if not manifest_uri:
+            raise ValueError("manifest_uri must be provided")
+
+        self._storage_config = storage_config
+        self._manifest_uri = manifest_uri
+        self._shards_base_uri = _manifest_shards_base_uri(manifest_uri)
+        self._data_format = data_format
+        if data_format == "jsonl.gz":
+            self._records_to_body: _RecordsToBody = gzjsonl.records_to_body
+        else:
+            raise ValueError(f"Unknown data format: {data_format}")
+        self._manager = multiprocessing.Manager()
+        self._saved_shards = self._manager.list()
+        self._lock = self._manager.RLock()
+
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        state["_manager"] = None
+        return state
+
+    @property
+    def manifest_uri(self) -> str:
+        return self._manifest_uri
+
+    def save(
+        self,
+        split: str,
+        curriculum_stage: str | None,
+        records: Iterable[Mapping[str, Any]],
+    ) -> ManifestShard:
+        if not split:
+            raise ValueError("split must not be empty")
+
+        body, record_count = self._records_to_body(records)
+        if record_count == 0:
+            raise ValueError("Cannot save an empty record shard")
+
+        shard_uri = self._shard_uri(split=split, curriculum_stage=curriculum_stage)
+        response = put_object_bytes(
+            shard_uri,
+            body,
+            storage_config=self._storage_config,
+            attributes={
+                "Content-Type": "application/gzip",
+                "content-format": "jsonl",
+                "compression": "gzip",
+            },
+        )
+        etag = response.get("e_tag") or response.get("etag") or ""
+        shard = ManifestShard(
+            uri=shard_uri,
+            records=record_count,
+            size_bytes=len(body),
+            etag=str(etag),
+        )
+        saved: _SavedShard = {
+            "split": split,
+            "curriculum_stage": curriculum_stage,
+            "shard": shard.model_dump(mode="json"),
+        }
+        with self._lock:
+            self._saved_shards.append(saved)
+        return shard
+
+    def commit_manifest(self) -> DataManifest:
+        manifest = self.build_manifest()
+        from syvain_training_data.manifest import manifest_json_bytes
+
+        put_object_bytes(
+            self.manifest_uri,
+            manifest_json_bytes(manifest),
+            storage_config=self._storage_config,
+            attributes={"Content-Type": "application/json"},
+        )
+        return manifest
+
+    def build_manifest(self) -> DataManifest:
+        with self._lock:
+            saved_shards = [
+                cast(_SavedShard, item) for item in list(self._saved_shards)
+            ]
+        if not saved_shards:
+            raise ValueError("Cannot commit a manifest with no saved shards")
+
+        split_names = sorted({item["split"] for item in saved_shards})
+        split_entries: dict[str, ManifestDatasetEntry] = {}
+        for split in split_names:
+            split_items = [item for item in saved_shards if item["split"] == split]
+            curriculum_entries = _curriculum_entries(split_items)
+            root_shards = _sorted_shards(split_items)
+            split_entries[split] = ManifestDatasetEntry(
+                records=sum(shard.records for shard in root_shards),
+                shards=root_shards,
+                curriculum=curriculum_entries,
+            )
+
+        return manifest_from_splits(split_entries, data_format=self._data_format)
+
+    def _shard_uri(self, *, split: str, curriculum_stage: str | None) -> str:
+        split_segment = safe_path_segment(split)
+        shard_name = f"part-{uuid.uuid4().hex}.jsonl.gz"
+        if curriculum_stage is None:
+            return join_s3_uri(self._shards_base_uri, split_segment, shard_name)
+        stage_segment = safe_path_segment(curriculum_stage)
+        return join_s3_uri(
+            self._shards_base_uri,
+            split_segment,
+            "curriculum",
+            stage_segment,
+            shard_name,
+        )
+
+
+def _manifest_shards_base_uri(manifest_uri: str) -> str:
+    parsed = parse_s3_uri(manifest_uri)
+    key = parsed.key.rstrip("/")
+    if key != parsed.key:
+        raise ValueError(f"Expected manifest file URI, got: {manifest_uri}")
+    parent_key = key.rpartition("/")[0]
+    shards_key = "shards" if not parent_key else f"{parent_key}/shards"
+    return f"s3://{parsed.bucket}/{shards_key}"
+
+
+def _curriculum_entries(items: list[_SavedShard]) -> dict[str, ManifestCurriculumStage]:
+    stage_names = sorted(
+        {
+            item["curriculum_stage"]
+            for item in items
+            if item["curriculum_stage"] is not None
+        }
+    )
+    return {
+        stage_name: curriculum_stage_from_shards(
+            _sorted_shards(
+                item for item in items if item["curriculum_stage"] == stage_name
+            )
+        )
+        for stage_name in stage_names
+    }
+
+
+def _sorted_shards(items: Iterable[_SavedShard]) -> list[ManifestShard]:
+    shards = [ManifestShard.model_validate(item["shard"]) for item in items]
+    return sorted(shards, key=lambda shard: shard.uri)
+
+
+__all__ = ["DatasetSaver"]

syvain_training_data-0.0.118/src/syvain_training_data/storage.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import builtins
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any
+from urllib.parse import quote, urlparse
+
+from obstore.store import S3Store
+
+
+@dataclass(frozen=True)
+class S3Uri:
+    bucket: str
+    key: str
+
+    @property
+    def uri(self) -> str:
+        return f"s3://{self.bucket}/{self.key}"
+
+
+@dataclass(frozen=True)
+class StorageConfig:
+    s3_base_url: str
+    region: str
+    access_key_id: str
+    secret_access_key: str
+    virtual_hosted_style_request: bool = False
+
+    def __post_init__(self) -> None:
+        _require_non_empty("s3_base_url", self.s3_base_url)
+        _require_non_empty("region", self.region)
+        _require_non_empty("access_key_id", self.access_key_id)
+        _require_non_empty("secret_access_key", self.secret_access_key)
+
+    def open_store(self, bucket: str) -> S3Store:
+        kwargs: dict[str, Any] = {
+            "endpoint": self.s3_base_url,
+            "region": self.region,
+            "access_key_id": self.access_key_id,
+            "secret_access_key": self.secret_access_key,
+            "virtual_hosted_style_request": self.virtual_hosted_style_request,
+        }
+        return S3Store(bucket, **kwargs)
+
+
+def parse_s3_uri(uri: str) -> S3Uri:
+    parsed = urlparse(uri)
+    key = parsed.path.lstrip("/")
+    if parsed.scheme != "s3" or not parsed.netloc or not key:
+        raise ValueError(f"Expected non-empty s3://bucket/key URI, got: {uri}")
+    return S3Uri(bucket=parsed.netloc, key=key)
+
+
+def join_s3_uri(base_uri: str, *parts: str) -> str:
+    parsed = parse_s3_uri(_ensure_base_key(base_uri))
+    normalized_key = parsed.key.rstrip("/")
+    suffix = "/".join(_clean_key_part(part) for part in parts if part)
+    if suffix:
+        normalized_key = f"{normalized_key}/{suffix}"
+    return f"s3://{parsed.bucket}/{normalized_key}"
+
+
+def safe_path_segment(value: str) -> str:
+    stripped = value.strip()
+    if not stripped:
+        raise ValueError("S3 path segment must not be empty")
+    return quote(stripped, safe="-_.=+")
+
+
+def store_for_uri(
+    uri: str,
+    *,
+    storage_config: StorageConfig,
+) -> tuple[S3Store, S3Uri]:
+    parsed = parse_s3_uri(uri)
+    store = storage_config.open_store(parsed.bucket)
+    return store, parsed
+
+
+def read_object_bytes(
+    uri: str,
+    *,
+    storage_config: StorageConfig,
+) -> builtins.bytes:
+    store, parsed = store_for_uri(uri, storage_config=storage_config)
+    return bytes(store.get(parsed.key).bytes())
+
+
+def put_object_bytes(
+    uri: str,
+    body: builtins.bytes,
+    *,
+    storage_config: StorageConfig,
+    attributes: dict[str, str] | None = None,
+) -> Mapping[str, Any]:
+    store, parsed = store_for_uri(uri, storage_config=storage_config)
+    return store.put(parsed.key, body, attributes=attributes)
+
+
+def _ensure_base_key(uri: str) -> str:
+    parsed = urlparse(uri)
+    if parsed.scheme == "s3" and parsed.netloc and parsed.path.lstrip("/"):
+        return uri
+    raise ValueError(f"Expected non-empty s3://bucket/key URI, got: {uri}")
+
+
+def _clean_key_part(value: str) -> str:
+    return value.strip().strip("/")
+
+
+def _require_non_empty(field_name: str, value: str) -> None:
+    if not isinstance(value, str) or not value.strip():
+        raise ValueError(f"{field_name} must be provided")
+
+
+__all__ = [
+    "S3Uri",
+    "StorageConfig",
+    "join_s3_uri",
+    "parse_s3_uri",
+    "put_object_bytes",
+    "read_object_bytes",
+    "safe_path_segment",
+    "store_for_uri",
+]