Cccccz/HY

1"""Atomic writer for BF16 Context prefeatures and four-step Teacher targets."""2 3from __future__ import annotations4 5import json6import os7from pathlib import Path8from typing import Any, Mapping9 10import torch11from safetensors.torch import save_file12 13from .prefeature_schema import (14    CONTEXT_BLOCK_IDS,15    SCHEMA_VERSION_PREFEATURE,16    validate_case_tensors,17    validate_context_prefeature,18    validate_step_tensors,19)20from .v2_writer import _convert21from .v2_schema import DEFAULT_TEXT_KV_TOKENS22 23 24class PredictorPrefeatureDatasetWriter:25    def __init__(26        self,27        root: str | os.PathLike[str],28        *,29        worker_id: int,30        text_kv_tokens: int = DEFAULT_TEXT_KV_TOKENS,31    ) -> None:32        self.root = Path(root).resolve()33        self.worker_id = int(worker_id)34        self.text_kv_tokens = int(text_kv_tokens)35        self.case_dir = self.root / "cases"36        self.step_dir = self.root / "steps"37        self.context_dir = self.root / "context_prefeature"38        self.video_dir = self.root / "videos"39        self.manifest_dir = self.root / "manifests"40        self.manifest_path = self.manifest_dir / f"worker_{self.worker_id:02d}.jsonl"41        for path in (42            self.case_dir,43            self.step_dir,44            self.context_dir,45            self.video_dir,46            self.manifest_dir,47        ):48            path.mkdir(parents=True, exist_ok=True)49        self._manifest_keys = self._load_existing_keys()50 51    def _load_existing_keys(self) -> set[tuple[int, int, int]]:52        keys: set[tuple[int, int, int]] = set()53        if self.manifest_path.is_file():54            with self.manifest_path.open("r", encoding="utf-8") as handle:55                for line in handle:56                    if line.strip():57                        item = json.loads(line)58                        keys.add(59                            (int(item["case_id"]), int(item["action_id"]), int(item["chunk_id"]))60                        )61        return keys62 63    def case_path(self, case_id: int) -> Path:64        return self.case_dir / f"case_{case_id:04d}.safetensors"65 66    def step_path(self, case_id: int, action_id: int, chunk_id: int) -> Path:67        return (68            self.step_dir69            / f"case_{case_id:04d}"70            / f"action_{action_id:02d}"71            / f"chunk_{chunk_id:02d}.safetensors"72        )73 74    def context_path(self, block_id: int, case_id: int, action_id: int, chunk_id: int) -> Path:75        return (76            self.context_dir77            / f"block_{block_id:02d}"78            / f"case_{case_id:04d}"79            / f"action_{action_id:02d}"80            / f"chunk_{chunk_id:02d}.safetensors"81        )82 83    def video_path(self, case_id: int, action_id: int) -> Path:84        return self.video_dir / f"case_{case_id:04d}_action_{action_id:02d}.mp4"85 86    def is_chunk_complete(self, case_id: int, action_id: int, chunk_id: int) -> bool:87        key = (int(case_id), int(action_id), int(chunk_id))88        return (89            key in self._manifest_keys90            and self.step_path(*key).is_file()91            and all(self.context_path(block_id, *key).is_file() for block_id in CONTEXT_BLOCK_IDS)92        )93 94    @staticmethod95    def _atomic_save(path: Path, tensors: Mapping[str, torch.Tensor]) -> None:96        path.parent.mkdir(parents=True, exist_ok=True)97        tmp = path.with_suffix(path.suffix + f".tmp.{os.getpid()}")98        save_file(dict(tensors), str(tmp))99        os.replace(tmp, path)100 101    def save_case(self, case_id: int, tensors: Mapping[str, torch.Tensor]) -> Path:102        converted = {103            name: _convert(tensor)104            for name, tensor in tensors.items()105        }106        first_k = converted[f"block_{CONTEXT_BLOCK_IDS[0]:02d}_k_txt"]107        valid_tokens = int(first_k.shape[2])108        if valid_tokens > self.text_kv_tokens:109            raise ValueError(110                f"Text KV has {valid_tokens} tokens, exceeding {self.text_kv_tokens}"111            )112        converted["text_valid_mask"] = (113            torch.arange(self.text_kv_tokens)[None] < valid_tokens114        )115        for block_id in CONTEXT_BLOCK_IDS:116            for field in ("k_txt", "v_txt"):117                name = f"block_{block_id:02d}_{field}"118                value = converted[name]119                if value.shape[2] < self.text_kv_tokens:120                    padded = value.new_zeros(121                        value.shape[0], value.shape[1], self.text_kv_tokens, value.shape[3]122                    )123                    padded[:, :, : value.shape[2]].copy_(value)124                    converted[name] = padded125        validate_case_tensors(converted)126        path = self.case_path(case_id)127        if not path.is_file():128            self._atomic_save(path, converted)129        return path130 131    def save_chunk(132        self,133        *,134        case_id: int,135        action_id: int,136        chunk_id: int,137        step_tensors: Mapping[str, torch.Tensor],138        context_by_block: Mapping[int, Mapping[str, torch.Tensor]],139        metadata: Mapping[str, Any],140    ) -> Path:141        key = (int(case_id), int(action_id), int(chunk_id))142        if self.is_chunk_complete(*key):143            return self.step_path(*key)144 145        converted_steps = {146            name: _convert(tensor, timestep="timestep" in name)147            for name, tensor in step_tensors.items()148        }149        validate_step_tensors(converted_steps)150        converted_context: dict[int, dict[str, torch.Tensor]] = {}151        frame_counts = set()152        for block_id in CONTEXT_BLOCK_IDS:153            if block_id not in context_by_block:154                raise ValueError(f"Missing Context block {block_id}")155            values = {156                name: _convert(value) if name not in (157                    "context_valid_mask",158                    "selected_frame_indices",159                    "rope_temporal_size",160                    "start_rope_start_idx",161                ) else value.detach().to(device="cpu").contiguous()162                for name, value in context_by_block[block_id].items()163            }164            frame_counts.add(validate_context_prefeature(block_id, values))165            converted_context[block_id] = values166        if len(frame_counts) != 1:167            raise ValueError("Context frame count differs between selected blocks")168        context_frames = next(iter(frame_counts))169 170        step_path = self.step_path(*key)171        self._atomic_save(step_path, converted_steps)172        context_files: dict[str, str] = {}173        for block_id, values in converted_context.items():174            path = self.context_path(block_id, *key)175            self._atomic_save(path, values)176            context_files[str(block_id)] = str(path.relative_to(self.root))177 178        record = {179            "schema_version": SCHEMA_VERSION_PREFEATURE,180            "case_id": key[0],181            "action_id": key[1],182            "chunk_id": key[2],183            "step_tensor_file": str(step_path.relative_to(self.root)),184            "case_tensor_file": str(self.case_path(case_id).relative_to(self.root)),185            "context_tensor_files": context_files,186            "context_block_ids": list(CONTEXT_BLOCK_IDS),187            "context_frames": int(context_frames),188            "padding": "dynamic_at_collate",189            **dict(metadata),190        }191        line = json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n"192        with self.manifest_path.open("a", encoding="utf-8") as handle:193            handle.write(line)194            handle.flush()195            os.fsync(handle.fileno())196        self._manifest_keys.add(key)197        return step_path198