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