File size: 2,569 Bytes
840c3a4 | 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 | #!/usr/bin/env python3
"""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())
|