Cccccz/HY

1"""Atomic safetensors writer and JSONL manifest for Predictor data."""2 3from __future__ import annotations4 5import json6import os7from pathlib import Path8from typing import Any, Mapping9 10import torch11from safetensors.torch import save_file12 13from .schema import SCHEMA_VERSION, validate_case_tensors, validate_chunk_tensors14 15 16def _cpu_contiguous(tensor: torch.Tensor, *, dtype: torch.dtype | None = None) -> torch.Tensor:17    tensor = tensor.detach()18    if dtype is not None and tensor.is_floating_point():19        tensor = tensor.to(dtype=dtype)20    return tensor.to(device="cpu").contiguous()21 22 23class PredictorDatasetWriter:24    def __init__(self, root: str | os.PathLike[str], *, save_dtype: torch.dtype = torch.bfloat16):25        self.root = Path(root).resolve()26        self.case_dir = self.root / "cases"27        self.chunk_dir = self.root / "chunks"28        self.manifest_path = self.root / "manifest.jsonl"29        self.train_eval_manifest_path = self.root / "train_eval_manifest.jsonl"30        self.save_dtype = save_dtype31        self.case_dir.mkdir(parents=True, exist_ok=True)32        self.chunk_dir.mkdir(parents=True, exist_ok=True)33        self._manifest_keys = self._load_existing_keys()34 35    def _load_existing_keys(self) -> set[tuple[int, int, int]]:36        keys: set[tuple[int, int, int]] = set()37        if not self.manifest_path.exists():38            return keys39        with self.manifest_path.open("r", encoding="utf-8") as handle:40            for line in handle:41                if not line.strip():42                    continue43                item = json.loads(line)44                keys.add((int(item["case_id"]), int(item["seed"]), int(item["chunk_id"])))45        return keys46 47    def case_path(self, case_id: int) -> Path:48        return self.case_dir / f"case_{case_id:02d}.safetensors"49 50    def chunk_path(self, case_id: int, seed: int, chunk_id: int) -> Path:51        return self.chunk_dir / f"case_{case_id:02d}_seed_{seed}" / f"chunk_{chunk_id:02d}.safetensors"52 53    def is_chunk_complete(self, case_id: int, seed: int, chunk_id: int) -> bool:54        key = (case_id, seed, chunk_id)55        return key in self._manifest_keys and self.chunk_path(case_id, seed, chunk_id).is_file()56 57    def _atomic_save(self, path: Path, tensors: Mapping[str, torch.Tensor]) -> None:58        path.parent.mkdir(parents=True, exist_ok=True)59        tmp_path = path.with_suffix(path.suffix + f".tmp.{os.getpid()}")60        save_file(dict(tensors), str(tmp_path))61        os.replace(tmp_path, path)62 63    def save_case(self, case_id: int, tensors: Mapping[str, torch.Tensor]) -> Path:64        converted = {65            name: _cpu_contiguous(tensor, dtype=self.save_dtype)66            for name, tensor in tensors.items()67        }68        validate_case_tensors(converted)69        path = self.case_path(case_id)70        self._atomic_save(path, converted)71        return path72 73    def save_chunk(74        self,75        *,76        case_id: int,77        seed: int,78        chunk_id: int,79        tensors: Mapping[str, torch.Tensor],80        metadata: Mapping[str, Any],81    ) -> Path:82        key = (case_id, seed, chunk_id)83        if self.is_chunk_complete(*key):84            return self.chunk_path(*key)85 86        converted = {}87        for name, tensor in tensors.items():88            dtype = self.save_dtype if tensor.is_floating_point() and "timestep" not in name else None89            if "timestep" in name:90                dtype = torch.float3291            converted[name] = _cpu_contiguous(tensor, dtype=dtype)92        validate_chunk_tensors(converted)93 94        path = self.chunk_path(*key)95        self._atomic_save(path, converted)96        record = {97            "schema_version": SCHEMA_VERSION,98            "case_id": case_id,99            "seed": seed,100            "chunk_id": chunk_id,101            "tensor_file": str(path.relative_to(self.root)),102            "case_tensor_file": str(self.case_path(case_id).relative_to(self.root)),103            **dict(metadata),104        }105        line = json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n"106        for manifest in (self.manifest_path, self.train_eval_manifest_path):107            with manifest.open("a", encoding="utf-8") as handle:108                handle.write(line)109                handle.flush()110                os.fsync(handle.fileno())111        self._manifest_keys.add(key)112        return path113