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