| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
|
|
| REPO_ROOT = Path(__file__).resolve().parents[1] |
| DEFAULT_CHECKPOINT = REPO_ROOT / "checkpoints" / "29999" |
| ASSET_ID = "real_piper_x_6tasks" |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser( |
| description=( |
| "Validate the released StreamPI Piper-X 20k " |
| "deployment checkpoint." |
| ) |
| ) |
| parser.add_argument( |
| "--checkpoint", |
| type=Path, |
| default=DEFAULT_CHECKPOINT, |
| help="Checkpoint directory. Default: checkpoints/29999", |
| ) |
| return parser.parse_args() |
|
|
|
|
| def directory_size(path: Path) -> int: |
| total = 0 |
|
|
| for item in path.rglob("*"): |
| if item.is_file(): |
| total += item.stat().st_size |
|
|
| return total |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| checkpoint = args.checkpoint.expanduser().resolve() |
|
|
| required_files = ( |
| checkpoint / "params" / "_METADATA", |
| checkpoint / "params" / "_sharding", |
| checkpoint / "params" / "manifest.ocdbt", |
| checkpoint |
| / "assets" |
| / ASSET_ID |
| / "norm_stats.json", |
| ) |
|
|
| missing = [ |
| path for path in required_files |
| if not path.is_file() |
| ] |
|
|
| if missing: |
| details = "\n".join( |
| f" - {path}" for path in missing |
| ) |
| raise FileNotFoundError( |
| "Missing deployment checkpoint files:\n" |
| + details |
| ) |
|
|
| forbidden_paths = ( |
| checkpoint / "train_state", |
| checkpoint / "optimizer", |
| ) |
|
|
| unexpected = [ |
| path for path in forbidden_paths |
| if path.exists() |
| ] |
|
|
| if unexpected: |
| details = "\n".join( |
| f" - {path}" for path in unexpected |
| ) |
| raise RuntimeError( |
| "Training-only artifacts remain in the " |
| "deployment checkpoint:\n" |
| + details |
| ) |
|
|
| norm_path = ( |
| checkpoint |
| / "assets" |
| / ASSET_ID |
| / "norm_stats.json" |
| ) |
|
|
| with norm_path.open("r", encoding="utf-8") as stream: |
| norm_stats = json.load(stream) |
|
|
| if not isinstance(norm_stats, dict) or not norm_stats: |
| raise ValueError( |
| f"Invalid or empty norm stats: {norm_path}" |
| ) |
|
|
| parameter_files = [ |
| path |
| for path in (checkpoint / "params").rglob("*") |
| if path.is_file() |
| ] |
|
|
| if not parameter_files: |
| raise RuntimeError( |
| f"No parameter files found under {checkpoint / 'params'}" |
| ) |
|
|
| symlinks = [ |
| path |
| for path in checkpoint.rglob("*") |
| if path.is_symlink() |
| ] |
|
|
| if symlinks: |
| details = "\n".join( |
| f" - {path}" for path in symlinks |
| ) |
| raise RuntimeError( |
| "Deployment checkpoint contains symbolic links:\n" |
| + details |
| ) |
|
|
| for required in required_files: |
| print( |
| "OK:", |
| required.relative_to(checkpoint), |
| ) |
|
|
| print(f"checkpoint={checkpoint}") |
| print(f"parameter_files={len(parameter_files)}") |
| print(f"checkpoint_bytes={directory_size(checkpoint)}") |
| print(f"norm_stat_keys={sorted(norm_stats)}") |
| print("train_state_present=false") |
| print("symlink_count=0") |
| print("CHECKPOINT_STRUCTURE_OK") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
