File size: 7,788 Bytes
2847d0b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | #!/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()
|