"""Atomic writer for fresh Predictor-v2 step targets and four-layer context KV.""" from __future__ import annotations import json import os from pathlib import Path from typing import Any, Mapping import torch from safetensors.torch import save_file from .v2_schema import ( CONTEXT_BLOCK_IDS, DEFAULT_CONTEXT_KV_FRAMES, DEFAULT_TEXT_KV_TOKENS, SCHEMA_VERSION_V2, TOKENS_PER_FRAME, validate_case_tensors_v2, validate_step_tensors_v2, validate_vision_context_v2, ) def _convert(tensor: torch.Tensor, *, timestep: bool = False) -> torch.Tensor: tensor = tensor.detach() if timestep: tensor = tensor.to(dtype=torch.float32) elif tensor.is_floating_point(): tensor = tensor.to(dtype=torch.bfloat16) return tensor.to(device="cpu").contiguous() class PredictorV2DatasetWriter: def __init__( self, root: str | os.PathLike[str], *, text_kv_tokens: int = DEFAULT_TEXT_KV_TOKENS, context_kv_frames: int = DEFAULT_CONTEXT_KV_FRAMES, ) -> None: self.root = Path(root).resolve() self.text_kv_tokens = int(text_kv_tokens) self.context_kv_frames = int(context_kv_frames) if self.text_kv_tokens <= 0 or self.context_kv_frames < 0: raise ValueError("Invalid padded KV sizes") self.case_dir = self.root / "cases" self.step_dir = self.root / "steps" self.context_dir = self.root / "context" self.manifest_path = self.root / "manifest.jsonl" self.train_eval_manifest_path = self.root / "train_eval_manifest.jsonl" for path in (self.case_dir, self.step_dir, self.context_dir): path.mkdir(parents=True, exist_ok=True) self._manifest_keys = self._load_existing_keys() def _load_existing_keys(self) -> set[tuple[int, int, int]]: keys: set[tuple[int, int, int]] = set() if not self.manifest_path.exists(): return keys with self.manifest_path.open("r", encoding="utf-8") as handle: for line in handle: if line.strip(): item = json.loads(line) keys.add((int(item["case_id"]), int(item["seed"]), int(item["chunk_id"]))) return keys def case_path(self, case_id: int) -> Path: return self.case_dir / f"case_{case_id:02d}.safetensors" def step_path(self, case_id: int, seed: int, chunk_id: int) -> Path: return self.step_dir / f"case_{case_id:02d}_seed_{seed}" / f"chunk_{chunk_id:02d}.safetensors" def context_path(self, block_id: int, case_id: int, seed: int, chunk_id: int) -> Path: return ( self.context_dir / f"block_{block_id:02d}" / f"case_{case_id:02d}_seed_{seed}" / f"chunk_{chunk_id:02d}.safetensors" ) def is_chunk_complete(self, case_id: int, seed: int, chunk_id: int) -> bool: key = (case_id, seed, chunk_id) return ( key in self._manifest_keys and self.step_path(*key).is_file() and all(self.context_path(block_id, *key).is_file() for block_id in CONTEXT_BLOCK_IDS) ) @staticmethod def _atomic_save(path: Path, tensors: Mapping[str, torch.Tensor]) -> None: path.parent.mkdir(parents=True, exist_ok=True) tmp = path.with_suffix(path.suffix + f".tmp.{os.getpid()}") save_file(dict(tensors), str(tmp)) os.replace(tmp, path) @staticmethod def _pad_sequence( tensor: torch.Tensor, target_tokens: int, *, name: str, ) -> torch.Tensor: valid_tokens = int(tensor.shape[2]) if valid_tokens > target_tokens: raise ValueError( f"{name} has {valid_tokens} tokens, exceeding padded size {target_tokens}" ) if valid_tokens == target_tokens: return tensor padded = tensor.new_zeros((*tensor.shape[:2], target_tokens, tensor.shape[3])) padded[:, :, :valid_tokens].copy_(tensor) return padded def save_case(self, case_id: int, tensors: Mapping[str, torch.Tensor]) -> Path: converted = {name: _convert(tensor) for name, tensor in tensors.items()} first_key = f"block_{CONTEXT_BLOCK_IDS[0]:02d}_k_txt" valid_tokens = int(converted[first_key].shape[2]) converted["text_valid_mask"] = ( torch.arange(self.text_kv_tokens)[None] < valid_tokens ) for block_id in CONTEXT_BLOCK_IDS: for field in ("k_txt", "v_txt"): name = f"block_{block_id:02d}_{field}" converted[name] = self._pad_sequence( converted[name], self.text_kv_tokens, name=name ) validate_case_tensors_v2(converted) path = self.case_path(case_id) self._atomic_save(path, converted) return path def save_chunk( self, *, case_id: int, seed: int, chunk_id: int, step_tensors: Mapping[str, torch.Tensor], context_by_block: Mapping[int, Mapping[str, torch.Tensor]], metadata: Mapping[str, Any], ) -> Path: key = (case_id, seed, chunk_id) if self.is_chunk_complete(*key): return self.step_path(*key) converted_steps = { name: _convert(tensor, timestep="timestep" in name) for name, tensor in step_tensors.items() } validate_step_tensors_v2(converted_steps) context_frames = None converted_context: dict[int, dict[str, torch.Tensor]] = {} for block_id in CONTEXT_BLOCK_IDS: if block_id not in context_by_block: raise ValueError(f"Missing context block {block_id}") cache = {name: _convert(tensor) for name, tensor in context_by_block[block_id].items()} valid_tokens = int(cache["k_vision"].shape[2]) target_tokens = self.context_kv_frames * TOKENS_PER_FRAME cache["context_valid_mask"] = ( torch.arange(target_tokens)[None] < valid_tokens ) for field in ("k_vision", "v_vision"): cache[field] = self._pad_sequence( cache[field], target_tokens, name=f"block_{block_id}_{field}" ) frames = validate_vision_context_v2(block_id, cache) if context_frames is None: context_frames = frames elif frames != context_frames: raise ValueError("Context frame count differs between blocks") converted_context[block_id] = cache step_path = self.step_path(*key) self._atomic_save(step_path, converted_steps) context_files: dict[str, str] = {} for block_id, cache in converted_context.items(): path = self.context_path(block_id, *key) self._atomic_save(path, cache) context_files[str(block_id)] = str(path.relative_to(self.root)) record = { "schema_version": SCHEMA_VERSION_V2, "case_id": case_id, "seed": seed, "chunk_id": chunk_id, "step_tensor_file": str(step_path.relative_to(self.root)), "case_tensor_file": str(self.case_path(case_id).relative_to(self.root)), "context_tensor_files": context_files, "context_block_ids": list(CONTEXT_BLOCK_IDS), "context_frames": int(context_frames or 0), "text_kv_tokens": self.text_kv_tokens, "context_kv_tokens": self.context_kv_frames * TOKENS_PER_FRAME, "padding_masks": True, **dict(metadata), } line = json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n" for manifest in (self.manifest_path, self.train_eval_manifest_path): with manifest.open("a", encoding="utf-8") as handle: handle.write(line) handle.flush() os.fsync(handle.fileno()) self._manifest_keys.add(key) return step_path