#!/usr/bin/env python3 """Validate Predictor v4 manifests, file references, and sampled tensors.""" from __future__ import annotations import argparse import json import sys from pathlib import Path from typing import Any from safetensors.torch import load_file ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from predictor_data import ( # noqa: E402 CANDIDATE_BLOCK_IDS, NUM_CHUNKS, SCHEMA_VERSION, validate_case_tensors, validate_clean_prefeature, validate_step_tensors, ) DEFAULT_ROOT = Path( "/mnt/local_nvme/zoubin/cz/self_forcing_predictor_v4_1000_seed0" ) def read_jsonl(path: Path) -> list[dict[str, Any]]: result = [] with path.open("r", encoding="utf-8") as handle: for line_number, line in enumerate(handle, start=1): if not line.strip(): continue try: result.append(json.loads(line)) except json.JSONDecodeError as exc: raise ValueError(f"invalid JSON at {path}:{line_number}") from exc return result def ensure_file(root: Path, relative_path: str) -> Path: path = (root / relative_path).resolve() if root not in path.parents: raise ValueError(f"manifest path escapes dataset root: {relative_path}") if not path.is_file(): raise FileNotFoundError(path) return path def validate_record_paths( root: Path, record: dict[str, Any], records_by_key: dict[tuple[int, int], dict[str, Any]], ) -> None: case_id = int(record["case_id"]) chunk_id = int(record["chunk_id"]) if record["schema_version"] != SCHEMA_VERSION: raise ValueError(f"record {(case_id, chunk_id)} has wrong schema version") if int(record["seed"]) != 0 or not bool(record["seed_reset_per_case"]): raise ValueError(f"record {(case_id, chunk_id)} violates fixed per-case seed 0") if record["schedule"] != "F-F-F-F": raise ValueError(f"record {(case_id, chunk_id)} is not an F-F-F-F trajectory") if tuple(record["candidate_block_ids"]) != CANDIDATE_BLOCK_IDS: raise ValueError(f"record {(case_id, chunk_id)} uses unexpected candidate blocks") ensure_file(root, record["case_tensor_file"]) ensure_file(root, record["step_tensor_file"]) clean = record["clean_prefeature_files"] if set(clean) != {str(value) for value in CANDIDATE_BLOCK_IDS}: raise ValueError(f"record {(case_id, chunk_id)} has incomplete clean files") for path in clean.values(): ensure_file(root, path) expected_previous = ( None if chunk_id == 0 else records_by_key[(case_id, chunk_id - 1)]["step_tensor_file"] ) if record.get("previous_step_tensor_file") != expected_previous: raise ValueError(f"record {(case_id, chunk_id)} has wrong previous step reference") history = record.get("history_clean_prefeature_files") if history is None: raise ValueError(f"record {(case_id, chunk_id)} lacks history clean references") for block_id in CANDIDATE_BLOCK_IDS: paths = history.get(str(block_id), []) if len(paths) != chunk_id: raise ValueError( f"record {(case_id, chunk_id)} block {block_id} history has " f"{len(paths)} chunks, expected {chunk_id}" ) for path in paths: ensure_file(root, path) if int(record["context_frames"]) != chunk_id * 3: raise ValueError(f"record {(case_id, chunk_id)} has wrong context_frames") def validate_tensor_record(root: Path, record: dict[str, Any]) -> None: case_tensors = load_file( str(ensure_file(root, record["case_tensor_file"])), device="cpu", ) validate_case_tensors(case_tensors) step_tensors = load_file( str(ensure_file(root, record["step_tensor_file"])), device="cpu", ) validate_step_tensors(step_tensors) for block_id in CANDIDATE_BLOCK_IDS: values = load_file( str( ensure_file( root, record["clean_prefeature_files"][str(block_id)], ) ), device="cpu", ) validate_clean_prefeature(block_id, values) if int(values["start_frame"].item()) != int(record["chunk_id"]) * 3: raise ValueError("clean prefeature start_frame disagrees with manifest chunk") def validate_kv_metrics(root: Path, minimum_rows: int) -> None: paths = sorted((root / "logs").glob("kv_rebuild_worker_*.jsonl")) rows = [] for path in paths: rows.extend(read_jsonl(path)) if len(rows) < minimum_rows: raise RuntimeError( f"only {len(rows)} KV rebuild metric rows exist, expected at least {minimum_rows}" ) for row in rows: for key, value in row.items(): if key.endswith("_relative_l2") and float(value) > 5e-3: raise ValueError(f"{key} exceeds 5e-3: {value}") if key.endswith("_cosine") and float(value) < 0.9999: raise ValueError(f"{key} is below 0.9999: {value}") def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--dataset_root", type=Path, default=DEFAULT_ROOT) parser.add_argument( "--tensor_records", type=int, default=16, help="number of evenly spaced records to fully load; 0 validates every record", ) parser.add_argument( "--require_kv_metrics", type=int, default=0, help="require at least this many successful live clean-KV reconstruction rows", ) args = parser.parse_args() root = args.dataset_root.resolve() manifest_path = root / "manifest.jsonl" train_path = root / "train_manifest.jsonl" for path in (root / "cases.jsonl", manifest_path, train_path): if not path.is_file(): raise FileNotFoundError(path) cases = read_jsonl(root / "cases.jsonl") records = read_jsonl(manifest_path) train = read_jsonl(train_path) expected_records = len(cases) * NUM_CHUNKS if len(records) != expected_records: raise ValueError(f"manifest has {len(records)} records, expected {expected_records}") if len(train) != len(cases) * (NUM_CHUNKS - 1): raise ValueError("train_manifest does not contain exactly chunks 1..6") records_by_key = { (int(item["case_id"]), int(item["chunk_id"])): item for item in records } if len(records_by_key) != len(records): raise ValueError("manifest contains duplicate case/chunk keys") for record in records: validate_record_paths(root, record, records_by_key) if any(int(record["chunk_id"]) == 0 for record in train): raise ValueError("train_manifest must exclude chunk 0") if args.tensor_records < 0: raise ValueError("--tensor_records must be non-negative") if args.tensor_records == 0 or args.tensor_records >= len(records): tensor_records = records elif args.tensor_records: indices = { round(index * (len(records) - 1) / (args.tensor_records - 1)) if args.tensor_records > 1 else 0 for index in range(args.tensor_records) } tensor_records = [records[index] for index in sorted(indices)] else: tensor_records = [] for record in tensor_records: validate_tensor_record(root, record) if args.require_kv_metrics: validate_kv_metrics(root, args.require_kv_metrics) print( f"Validated {len(records)} manifest records / {len(train)} training chunks; " f"fully loaded {len(tensor_records)} tensor records." ) if __name__ == "__main__": main()