spb/wpai0002_uqo

1#!/usr/bin/env python32"""Batch-job administration: show status of a rewrite job's batches, and reset the ``collected`` flag of batches3whose results were not all recorded (so that ``scripts/rewrite.py --collect`` fetches them again)."""4from __future__ import annotations56import argparse7import json8import sys9from pathlib import Path1011ROOT = Path(__file__).resolve().parents[1]12sys.path.insert(0, str(ROOT / "src"))13from htr.utils.io import read_json, read_jsonl, write_json  # noqa: E402141516def main() -> None:17    ap = argparse.ArgumentParser()18    ap.add_argument("--job", required=True)19    ap.add_argument("--reset-incomplete", action="store_true")20    a = ap.parse_args()21    job = ROOT / "data/pairs" / a.job22    recorded = {r["request_id"] for r in read_jsonl(job / "rewrites.jsonl")}23    planned = {r["request_id"] for r in read_jsonl(job / "plan.jsonl")}24    print(f"job {a.job}: planned={len(planned)} recorded={len(recorded)}")25    for f in sorted((job / "batches").glob("*.json")):26        m = read_json(f)27        ids = m["custom_ids"]28        n_rec = sum(c in recorded for c in ids)29        n_plan = sum(c in planned for c in ids)30        print(f"  {m['batch_id'][-10:]}: {len(ids)} requests, recorded {n_rec}, in plan {n_plan}, collected={m.get('collected', False)}")31        if a.reset_incomplete and m.get("collected") and n_rec < len(ids):32            m["collected"] = False33            write_json(f, m)34            print("    -> reset collected flag")35    led = read_json(job / "ledger.json")36    print("ledger:", {k: led[k] for k in ("spent_usd", "n_ok", "n_refused", "n_error", "n_truncated", "n_empty") if k in led})373839if __name__ == "__main__":40    main()41