| |
| """Enumerate or fetch pinned public checkpoint resources.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
| import platform |
| import sys |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1]) |
| parser.add_argument("--config", type=Path) |
| parser.add_argument("--registry", type=Path) |
| parser.add_argument("--dry-run", action="store_true") |
| parser.add_argument("--all-public", action="store_true") |
| return parser.parse_args() |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| root = args.root.resolve() |
| project_root = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(project_root / "src")) |
| from dlb.checkpoints import ( |
| fetch_all_resources, |
| load_checkpoint_manifest, |
| require_server_platform, |
| validate_checkpoint_coverage, |
| ) |
| from dlb.registry import load_registry |
|
|
| manifest_path = (args.config or root / "artifacts" / "checkpoints.yaml").resolve() |
| registry_path = (args.registry or root / "configs" / "experiments.yaml").resolve() |
| manifest = load_checkpoint_manifest(manifest_path) |
| registry = load_registry(registry_path) |
| validate_checkpoint_coverage(registry, manifest) |
| if args.dry_run: |
| for resource_id, resource in manifest.resources.items(): |
| print( |
| f"RESOURCE {resource_id} {resource.backend} {resource.provenance} " |
| f"checkpoints/{resource.destination} teacher={resource.teacher_family}" |
| ) |
| for recipe_id, recipe in manifest.recipes.items(): |
| print( |
| f"RECIPE {recipe_id} {recipe.model}/{recipe.dataset} {recipe.provenance} " |
| f"{recipe.output} teacher={recipe.teacher_family}" |
| ) |
| return 0 |
| if not args.all_public: |
| print("ERROR: pass --all-public to acknowledge all public checkpoint downloads", file=sys.stderr) |
| return 2 |
| try: |
| require_server_platform(platform.system()) |
| except RuntimeError as error: |
| print(f"ERROR: {error}", file=sys.stderr) |
| return 2 |
| lock = fetch_all_resources(root, manifest_path, manifest) |
| for resource_id, record in lock["resources"].items(): |
| print(f"{str(record['status']).upper()} {resource_id} {record['destination']}") |
| return 1 if any(record["status"] != "downloaded" for record in lock["resources"].values()) else 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
