"""Atomic writer for BF16 Context prefeatures and four-step Teacher targets.""" 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 .prefeature_schema import ( CONTEXT_BLOCK_IDS, SCHEMA_VERSION_PREFEATURE, validate_case_tensors, validate_context_prefeature, validate_step_tensors, ) from .v2_writer import _convert from .v2_schema import DEFAULT_TEXT_KV_TOKENS class PredictorPrefeatureDatasetWriter: def __init__( self, root: str | os.PathLike[str], *, worker_id: int, text_kv_tokens: int = DEFAULT_TEXT_KV_TOKENS, ) -> None: self.root = Path(root).resolve() self.worker_id = int(worker_id) self.text_kv_tokens = int(text_kv_tokens) self.case_dir = self.root / "cases" self.step_dir = self.root / "steps" self.context_dir = self.root / "context_prefeature" self.video_dir = self.root / "videos" self.manifest_dir = self.root / "manifests" self.manifest_path = self.manifest_dir / f"worker_{self.worker_id:02d}.jsonl" for path in ( self.case_dir, self.step_dir, self.context_dir, self.video_dir, self.manifest_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 self.manifest_path.is_file(): 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["action_id"]), int(item["chunk_id"])) ) return keys def case_path(self, case_id: int) -> Path: return self.case_dir / f"case_{case_id:04d}.safetensors" def step_path(self, case_id: int, action_id: int, chunk_id: int) -> Path: return ( self.step_dir / f"case_{case_id:04d}" / f"action_{action_id:02d}" / f"chunk_{chunk_id:02d}.safetensors" ) def context_path(self, block_id: int, case_id: int, action_id: int, chunk_id: int) -> Path: return ( self.context_dir / f"block_{block_id:02d}" / f"case_{case_id:04d}" / f"action_{action_id:02d}" / f"chunk_{chunk_id:02d}.safetensors" ) def video_path(self, case_id: int, action_id: int) -> Path: return self.video_dir / f"case_{case_id:04d}_action_{action_id:02d}.mp4" def is_chunk_complete(self, case_id: int, action_id: int, chunk_id: int) -> bool: key = (int(case_id), int(action_id), int(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) def save_case(self, case_id: int, tensors: Mapping[str, torch.Tensor]) -> Path: converted = { name: _convert(tensor) for name, tensor in tensors.items() } first_k = converted[f"block_{CONTEXT_BLOCK_IDS[0]:02d}_k_txt"] valid_tokens = int(first_k.shape[2]) if valid_tokens > self.text_kv_tokens: raise ValueError( f"Text KV has {valid_tokens} tokens, exceeding {self.text_kv_tokens}" ) 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}" value = converted[name] if value.shape[2] < self.text_kv_tokens: padded = value.new_zeros( value.shape[0], value.shape[1], self.text_kv_tokens, value.shape[3] ) padded[:, :, : value.shape[2]].copy_(value) converted[name] = padded validate_case_tensors(converted) path = self.case_path(case_id) if not path.is_file(): self._atomic_save(path, converted) return path def save_chunk( self, *, case_id: int, action_id: 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 = (int(case_id), int(action_id), int(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(converted_steps) converted_context: dict[int, dict[str, torch.Tensor]] = {} frame_counts = set() for block_id in CONTEXT_BLOCK_IDS: if block_id not in context_by_block: raise ValueError(f"Missing Context block {block_id}") values = { name: _convert(value) if name not in ( "context_valid_mask", "selected_frame_indices", "rope_temporal_size", "start_rope_start_idx", ) else value.detach().to(device="cpu").contiguous() for name, value in context_by_block[block_id].items() } frame_counts.add(validate_context_prefeature(block_id, values)) converted_context[block_id] = values if len(frame_counts) != 1: raise ValueError("Context frame count differs between selected blocks") context_frames = next(iter(frame_counts)) step_path = self.step_path(*key) self._atomic_save(step_path, converted_steps) context_files: dict[str, str] = {} for block_id, values in converted_context.items(): path = self.context_path(block_id, *key) self._atomic_save(path, values) context_files[str(block_id)] = str(path.relative_to(self.root)) record = { "schema_version": SCHEMA_VERSION_PREFEATURE, "case_id": key[0], "action_id": key[1], "chunk_id": key[2], "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), "padding": "dynamic_at_collate", **dict(metadata), } 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._manifest_keys.add(key) return step_path