scripts/validate_predictor_v4_dataset.py

#!/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()