SPB Git

spb/wp42_uqo

Public
4.1 KB · 55 lines python
1"""CLI: python -m wp42.cli <command> [...]"""2from __future__ import annotations3import argparse, json, os, signal, sys4from .config import DB_PATH, ALL_ARMS, PRIMARY_MODEL, NODES5from .store import Store6from . import scheduler, watchdog, export as export_mod, bundle78def main(argv=None):9    ap = argparse.ArgumentParser(prog="wp42"); sub = ap.add_subparsers(dest="cmd", required=True)10    a = sub.add_parser("assign"); a.add_argument("--tasks", nargs="+", required=True); a.add_argument("--arms", nargs="+", default=ALL_ARMS); a.add_argument("--reps", type=int, default=2)11    a.add_argument("--seed", type=int, default=4242); a.add_argument("--model", default=PRIMARY_MODEL); a.add_argument("--stage", default="pilot"); a.add_argument("--model-tag", default=""); a.add_argument("--interrupt-s", type=int, default=None)12    r = sub.add_parser("run"); r.add_argument("--where", default="status='assigned'"); r.add_argument("--parallel", type=int, default=12); r.add_argument("--no-judge", action="store_true"); r.add_argument("--envelope", type=float, default=None)13    s = sub.add_parser("status"); s.add_argument("--json", action="store_true")14    e = sub.add_parser("evaluate"); e.add_argument("--run-ids", nargs="*"); e.add_argument("--where", default="status='collected'"); e.add_argument("--no-judge", action="store_true")15    sub.add_parser("export").add_argument("--tag", default="pilot")16    sub.add_parser("kill-all"); sub.add_parser("sync-tools"); sub.add_parser("ledger")17    i = sub.add_parser("intervention"); i.add_argument("--run-id", required=True); i.add_argument("--minutes", type=float, required=True); i.add_argument("--desc", required=True)18    x = sub.add_parser("reset-run"); x.add_argument("--run-ids", nargs="+", required=True)19    args = ap.parse_args(argv); store = Store(DB_PATH)20    if args.cmd == "assign":21        n = scheduler.make_assignment(store, args.tasks, args.arms, args.reps, args.seed, args.model, args.stage, args.model_tag, args.interrupt_s); print(f"assigned {n} new runs")22    elif args.cmd == "run":23        signal.signal(signal.SIGINT, lambda *_: scheduler.STOP.set())24        kw = {"envelope": args.envelope} if args.envelope else {}25        scheduler.run_queue(store, args.where, args.parallel, do_judge=not args.no_judge, **kw)26    elif args.cmd == "status":27        rows = store.q("SELECT status, arm, COUNT(*) AS n, ROUND(SUM(COALESCE(cost_usd,0)+COALESCE(recovery_cost_usd,0)),2) AS cost FROM runs GROUP BY status, arm ORDER BY status, arm")28        if args.json: print(json.dumps(rows, indent=1))29        else:30            for x in rows: print(f"{x['status']:<12} {x['arm']:<3} n={x['n']:<4} cost={x['cost']}")31            print(f"ledger total: {store.ledger_total():.2f} USD")32    elif args.cmd == "evaluate":33        ids = args.run_ids or [x["run_id"] for x in store.runs(args.where)]34        for rid in ids: scheduler.evaluate_and_judge(store, rid, do_judge=not args.no_judge)35    elif args.cmd == "export":36        print(export_mod.export(store, args.tag))37    elif args.cmd == "kill-all":38        scheduler.STOP.set(); watchdog.kill_everything()39    elif args.cmd == "sync-tools":40        for n in NODES: bundle.sync_tools(n); print(f"{n}: tools synced")41    elif args.cmd == "ledger":42        for x in store.q("SELECT kind, COUNT(*) n, ROUND(SUM(cost_usd),3) usd FROM ledger GROUP BY kind"): print(x)43        print("total", round(store.ledger_total(), 3))44    elif args.cmd == "intervention":45        store.add_intervention(args.run_id, args.minutes, args.desc); print("recorded")46    elif args.cmd == "reset-run":47        for rid in args.run_ids:48            store.exec("DELETE FROM events WHERE run_id=?", (rid,)); store.exec("DELETE FROM receipts WHERE run_id=?", (rid,)); store.exec("DELETE FROM judgments WHERE run_id=?", (rid,))49            store.exec("DELETE FROM metrics WHERE run_id=?", (rid,)); store.exec("DELETE FROM artifacts WHERE run_id=?", (rid,))50            store.update_run(rid, status="assigned", node=None, session_id=None, cost_usd=None, recovery_cost_usd=None, recovery_attempted=0, recovery_success=None, selfcheck_pass_first=None, selfcheck_pass_final=None, notes="reset")51            print("reset", rid)5253if __name__ == "__main__":54    main()55
Files · main
Menu

Theme