| """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 |
| |
| |
| |
| 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 |
|
|
