File size: 2,693 Bytes
5cd1287 | 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 | #!/usr/bin/env python3
"""Validate transferred RoboTwin assets before a smoke or Slurm launch."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from train_frameskip_robotwin_continuation import load_cache
def manifest_support(path: Path) -> set[tuple[str, str]]:
payload = json.loads(path.read_text(encoding="utf-8"))
support = set()
for dataset in payload["datasets"]:
name = dataset.get("name", dataset.get("dataset_name"))
for episode in dataset.get("episodes", dataset.get("trajectories", [])):
trajectory_id = episode.get("trajectory_id", episode.get("id"))
support.add((name, str(trajectory_id)))
if not support:
raise RuntimeError("valid-anchor manifest has no trajectory support")
return support
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--checkpoint", required=True, type=Path)
parser.add_argument("--manifest", required=True, type=Path)
parser.add_argument("--cache", required=True, type=Path)
args = parser.parse_args()
commit_path = args.checkpoint / "checkpoint.json"
commit = json.loads(commit_path.read_text(encoding="utf-8"))
if commit.get("schema_version") != 1 or int(commit.get("completed_steps", -1)) != 10000:
raise RuntimeError("checkpoint commit is not a complete step-10000 state")
if int(commit.get("world_size", -1)) != 2:
raise RuntimeError("10K checkpoint world size must be 2")
inventory = commit.get("files", [])
missing = [
entry["name"]
for entry in inventory
if not (args.checkpoint / entry["name"]).is_file()
or (args.checkpoint / entry["name"]).stat().st_size != int(entry["bytes"])
]
if not inventory or missing:
raise RuntimeError(f"incomplete 10K checkpoint inventory: missing={missing}")
parent_config = args.checkpoint.parent.parent / "config.full.yaml"
if not parent_config.is_file():
raise RuntimeError(f"missing parent config required by exact resume: {parent_config}")
cache, payload = load_cache(str(args.cache))
support = manifest_support(args.manifest)
cache_support = set(cache)
if support != cache_support:
raise RuntimeError(
"FrameSkip cache support differs from manifest support: "
f"missing={len(support - cache_support)}, extra={len(cache_support - support)}"
)
print("FRAMESKIP_ROBOTWIN_ASSET_PREFLIGHT=PASS")
print(f"checkpoint_files={len(inventory)}")
print(f"trajectory_support={len(support)}")
print(f"cache_records={len(payload['records'])}")
if __name__ == "__main__":
main()
|