SPB Git

spb/wp29_uqo

Public
4.9 KB · 79 lines python
1"""Command-line interface of the PGAR harness.2  pgar assign  --split pilot --tasks F1-dev-01,... --repeats 1 --seed 13  pgar run     --split pilot [--runs R...] [--workers N]4  pgar status  [--split ...]5  pgar export6  pgar probe   (sentinel isolation probes on all nodes)7  pgar deploy  (install gate + env file on nodes)8"""9from __future__ import annotations10import argparse, json, os, sys, yaml11from .state import StateStore12from .scheduler import Controller13from .exports import export_runs14from . import isolation, launcher1516REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))1718def load_cfg(split):19    base = yaml.safe_load(open(os.path.join(REPO, "configs", "budgets.yaml")))20    cfg = yaml.safe_load(open(os.path.join(REPO, "configs", f"{split}.yaml")))21    out = dict(base); out.update(cfg); return out2223def get_store():24    return StateStore(os.path.join(REPO, "runs", "state.sqlite"))2526def main(argv=None):27    ap = argparse.ArgumentParser(prog="pgar"); sub = ap.add_subparsers(dest="cmd", required=True)28    a = sub.add_parser("assign"); a.add_argument("--split", required=True); a.add_argument("--tasks", required=True); a.add_argument("--repeats", type=int, default=1); a.add_argument("--seed", type=int, required=True); a.add_argument("--arms", default="A,B,C,D"); a.add_argument("--start-seq", type=int, default=None)29    r = sub.add_parser("run"); r.add_argument("--split", required=True); r.add_argument("--runs", default=None); r.add_argument("--workers", type=int, default=None); r.add_argument("--fault", default=None, help="JSON fault spec applied to all selected runs"); r.add_argument("--rerun-errors", action="store_true", help="also select runs in status 'error' (infrastructure rerun policy)")30    s = sub.add_parser("status"); s.add_argument("--split", default=None)31    sub.add_parser("export")32    sub.add_parser("probe")33    sub.add_parser("deploy")34    e = sub.add_parser("evaluate"); e.add_argument("--runs", required=True)35    args = ap.parse_args(argv)36    store = get_store()37    if args.cmd == "assign":38        cfg = load_cfg(args.split); ctl = Controller(REPO, store, cfg, cfg["nodes"])39        existing = store.query("SELECT COUNT(*) AS n FROM runs WHERE split=?", (args.split,))[0]["n"]40        runs = ctl.make_assignment(args.split, args.tasks.split(","), args.repeats, args.seed, arms=tuple(args.arms.split(",")), start_seq=(args.start_seq or existing + 1))41        ctl.register(runs)42        path = os.path.join(REPO, "runs", args.split, f"assignment_seed{args.seed}.json"); os.makedirs(os.path.dirname(path), exist_ok=True)43        json.dump([json.loads(m.to_json()) for m in runs], open(path, "w"), indent=1)44        print(f"registered {len(runs)} runs for split {args.split}: {runs[0].run_id}..{runs[-1].run_id} -> {path}")45    elif args.cmd == "run":46        cfg = load_cfg(args.split); ctl = Controller(REPO, store, cfg, cfg["nodes"])47        ids = args.runs.split(",") if args.runs else [m["run_id"] for m in store.list_runs(split=args.split) if m["status"] == "not_started" or (args.rerun_errors and m["status"] == "error")]48        faults = {rid: json.loads(args.fault) for rid in ids} if args.fault else None49        res = ctl.run_many(ids, max_workers=args.workers, faults=faults)50        print(json.dumps([{k: v for k, v in x.items() if k != "evaluation"} for x in res], indent=1, default=str))51        export_runs(store, REPO)52    elif args.cmd == "status":53        for m in store.list_runs(split=args.split):54            ev = store.query("SELECT scientific_success, n_violations FROM evaluations WHERE run_id=?", (m["run_id"],))55            u = m.get("usage") or {}56            print(f"{m['run_id']:8} {m['task_id']:12} {m['treatment']} r{m['repeat_index']} {m.get('node') or '-':8} {m['status']:16} cost={u.get('spent_usd', 0):.3f} success={ev[0]['scientific_success'] if ev else '-'} viol={ev[0]['n_violations'] if ev else '-'} {m.get('terminal_reason') or ''}")57    elif args.cmd == "export":58        print(export_runs(store, REPO))59    elif args.cmd == "probe":60        cfg = load_cfg("development")61        for n in cfg["nodes"]:62            print(json.dumps(isolation.sentinel_probe(n, "~/wp29")))63    elif args.cmd == "deploy":64        cfg = load_cfg("development")65        for n in cfg["nodes"]:66            launcher.ssh(n, "mkdir -p ~/wp29/gates ~/wp29/runs && chmod 755 ~/wp29/gates", check=True)67            import subprocess68            subprocess.run(["scp", "-q", "-o", "BatchMode=yes", os.path.join(REPO, "src", "pgar", "runtime_gates.py"), f"{n}:~/wp29/gates/pgar_gate.py"], check=True)69            launcher.ssh(n, "chmod 444 ~/wp29/gates/pgar_gate.py && shasum -a 256 ~/wp29/gates/pgar_gate.py", check=True)70            print(n, "gate deployed")71    elif args.cmd == "evaluate":72        from .evaluation import evaluate_run73        for rid in args.runs.split(","):74            print(rid, json.dumps({k: evaluate_run(REPO, store, rid).get(k) for k in ("scientific_success", "quality_score")}))75        export_runs(store, REPO)7677if __name__ == "__main__":78    main()79
Files · main
Menu

Theme