"""Crash-safe writer for Predictor v4 Full-DiT teacher trajectories.""" from __future__ import annotations import json import os from collections.abc import Mapping from pathlib import Path from typing import Any import torch from safetensors.torch import save_file from .schema import ( CANDIDATE_BLOCK_IDS, SCHEMA_VERSION, validate_case_tensors, validate_clean_prefeature, validate_step_tensors, ) def _bf16_cpu(tensor: torch.Tensor) -> torch.Tensor: return tensor.detach().to(device="cpu", dtype=torch.bfloat16).contiguous() def atomic_write_text(path: Path, text: str) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}") with temporary.open("w", encoding="utf-8") as handle: handle.write(text) handle.flush() os.fsync(handle.fileno()) os.replace(temporary, path) def atomic_save_safetensors( path: Path, tensors: Mapping[str, torch.Tensor], ) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}") try: save_file(dict(tensors), str(temporary)) os.replace(temporary, path) finally: if temporary.exists(): temporary.unlink() class PredictorV4DatasetWriter: """Write one worker shard and fsync every committed manifest record.""" def __init__( self, root: str | os.PathLike[str], *, worker_id: int, block_ids: tuple[int, ...] = CANDIDATE_BLOCK_IDS, ) -> None: self.root = Path(root).resolve() self.worker_id = int(worker_id) self.block_ids = tuple(int(value) for value in block_ids) self.case_dir = self.root / "cases" self.step_dir = self.root / "steps" self.clean_dir = self.root / "clean_prefeature" self.manifest_dir = self.root / "manifests" self.log_dir = self.root / "logs" self.manifest_path = self.manifest_dir / f"worker_{self.worker_id:02d}.jsonl" for path in ( self.case_dir, self.step_dir, self.clean_dir, self.manifest_dir, self.log_dir, ): path.mkdir(parents=True, exist_ok=True) self._records = self._read_committed_records() def _read_committed_records(self) -> dict[tuple[int, int], dict[str, Any]]: records: dict[tuple[int, int], dict[str, Any]] = {} if not self.manifest_path.is_file(): return records text = self.manifest_path.read_text(encoding="utf-8") lines = text.splitlines(keepends=True) committed_lines: list[str] = [] for line_number, line in enumerate(lines, start=1): if not line.strip(): committed_lines.append(line) continue try: item = json.loads(line) except json.JSONDecodeError as exc: is_truncated_tail = ( line_number == len(lines) and not line.endswith(("\n", "\r")) ) if not is_truncated_tail: raise ValueError( f"invalid JSON at {self.manifest_path}:{line_number}" ) from exc # A process may be killed between write(2) and fsync(2). Repair # only an unterminated final record; corruption elsewhere is # never silently ignored. atomic_write_text(self.manifest_path, "".join(committed_lines)) break key = (int(item["case_id"]), int(item["chunk_id"])) if key in records and records[key] != item: raise ValueError(f"conflicting duplicate worker manifest record {key}") records[key] = item committed_lines.append(line) return records def case_path(self, case_id: int) -> Path: return self.case_dir / f"case_{int(case_id):06d}.safetensors" def step_path(self, case_id: int, chunk_id: int) -> Path: return ( self.step_dir / f"case_{int(case_id):06d}" / f"chunk_{int(chunk_id):02d}.safetensors" ) def clean_path(self, block_id: int, case_id: int, chunk_id: int) -> Path: return ( self.clean_dir / f"block_{int(block_id):02d}" / f"case_{int(case_id):06d}" / f"chunk_{int(chunk_id):02d}.safetensors" ) def _record_files_exist(self, record: Mapping[str, Any]) -> bool: paths = [ self.root / str(record["case_tensor_file"]), self.root / str(record["step_tensor_file"]), ] paths.extend( self.root / str(value) for value in record["clean_prefeature_files"].values() ) return all(path.is_file() for path in paths) def is_chunk_complete(self, case_id: int, chunk_id: int) -> bool: record = self._records.get((int(case_id), int(chunk_id))) return record is not None and self._record_files_exist(record) def save_case( self, case_id: int, tensors: Mapping[str, torch.Tensor], ) -> Path: converted = {name: _bf16_cpu(value) for name, value in tensors.items()} validate_case_tensors(converted, self.block_ids) path = self.case_path(case_id) if not path.is_file(): atomic_save_safetensors(path, converted) return path def save_chunk( self, *, case_id: int, chunk_id: int, step_tensors: Mapping[str, torch.Tensor], clean_features: Mapping[int, torch.Tensor], start_frame: int, metadata: Mapping[str, Any], ) -> Path: key = (int(case_id), int(chunk_id)) if self.is_chunk_complete(*key): return self.step_path(*key) if not self.case_path(case_id).is_file(): raise RuntimeError(f"case tensor file must be saved before chunk {key}") converted_steps: dict[str, torch.Tensor] = {} for name, value in step_tensors.items(): if name.endswith("_timestep"): converted_steps[name] = value.detach().to( device="cpu", dtype=torch.int64 ).contiguous() else: converted_steps[name] = _bf16_cpu(value) validate_step_tensors(converted_steps) converted_clean: dict[int, dict[str, torch.Tensor]] = {} for block_id in self.block_ids: if block_id not in clean_features: raise ValueError(f"missing clean prefeature for block {block_id}") values = { "self_attn_input": _bf16_cpu(clean_features[block_id]), "start_frame": torch.tensor([int(start_frame)], dtype=torch.int64), "num_frames": torch.tensor([3], dtype=torch.int64), } validate_clean_prefeature(block_id, values) converted_clean[block_id] = values step_path = self.step_path(*key) atomic_save_safetensors(step_path, converted_steps) clean_files: dict[str, str] = {} for block_id, values in converted_clean.items(): path = self.clean_path(block_id, *key) atomic_save_safetensors(path, values) clean_files[str(block_id)] = str(path.relative_to(self.root)) record = { "schema_version": SCHEMA_VERSION, "case_id": key[0], "chunk_id": key[1], "case_tensor_file": str(self.case_path(case_id).relative_to(self.root)), "step_tensor_file": str(step_path.relative_to(self.root)), "clean_prefeature_files": clean_files, "candidate_block_ids": list(self.block_ids), "context_frames": int(start_frame), **dict(metadata), } existing = self._records.get(key) if existing is None: line = json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n" with self.manifest_path.open("a", encoding="utf-8") as handle: handle.write(line) handle.flush() os.fsync(handle.fileno()) self._records[key] = record elif existing != record: raise ValueError(f"recomputed metadata differs for manifest record {key}") return step_path