spb/wp46_uqo

1#!/usr/bin/env python32"""Run rolling-origin forecasting jobs (instrument x model x distribution) and write one Parquet per job.34usage: run_benchmark.py --jobs results/manifests/jobs_<name>.json [--shard i/N] [--workers 8] [--out results/forecasts]5The job list is produced by scripts/make_jobs.py from the frozen configs. Each job writes6  <out>/<symbol>/<spec_id>.parquet            (forecasts, one row per origin)7  <out>/<symbol>/<spec_id>.fits.parquet       (refit attempts)8and a line in <out>/_runs.jsonl (git SHA, host, package versions, wall time). Existing outputs are skipped (resumable).9Nothing printed contains return values."""10from __future__ import annotations1112import argparse13import json14import os15import platform16import resource17import socket18import subprocess19import sys20import time21from concurrent.futures import ProcessPoolExecutor, as_completed22from pathlib import Path2324ROOT = Path(__file__).resolve().parents[1]25sys.path.insert(0, str(ROOT / "src"))262728def git_sha() -> str:29    try:30        return subprocess.check_output(["git", "-C", str(ROOT), "rev-parse", "--short", "HEAD"], text=True).strip()31    except Exception:  # noqa: BLE00132        return "unknown"333435def run_one(job: dict, out: str) -> dict:36    import numpy as np37    import pandas as pd38    from wp46.estimation import JobSpec, run_job39    from wp46.preprocessing import load_panel40    js = JobSpec(**{k: v for k, v in job.items() if k in JobSpec.__dataclass_fields__})41    d = Path(out) / js.symbol42    d.mkdir(parents=True, exist_ok=True)43    fp = d / f"{js.spec_id}.parquet"44    if fp.exists():45        return {"spec_id": js.spec_id, "symbol": js.symbol, "skipped": True}46    t0 = time.perf_counter()47    panel = load_panel(js.symbol)48    F, G = run_job(panel, js)49    F.to_parquet(fp, index=False)50    G.to_parquet(d / f"{js.spec_id}.fits.parquet", index=False)51    ru = resource.getrusage(resource.RUSAGE_SELF)52    return {"spec_id": js.spec_id, "symbol": js.symbol, "n_origins": int(len(F)), "n_fits": int(len(G)),53            "status_counts": F["status"].value_counts().to_dict() if len(F) else {},54            "fit_status_counts": G["status"].value_counts().to_dict() if len(G) else {},55            "wall_s": time.perf_counter() - t0, "cpu_s": ru.ru_utime + ru.ru_stime, "maxrss_mb": ru.ru_maxrss / (1024 ** 2 if platform.system() == "Darwin" else 1024),56            "fit_time_total_s": float(G["time_s"].sum()) if len(G) else 0.0, "fc_time_total_s": float(F["time_fc_s"].sum()) if "time_fc_s" in F else 0.0}575859def main() -> None:60    ap = argparse.ArgumentParser()61    ap.add_argument("--jobs", required=True)62    ap.add_argument("--shard", default="0/1")63    ap.add_argument("--workers", type=int, default=max(1, os.cpu_count() - 2))64    ap.add_argument("--out", default=str(ROOT / "results" / "forecasts"))65    a = ap.parse_args()66    jobs = json.loads(Path(a.jobs).read_text())67    i, n = map(int, a.shard.split("/"))68    jobs = [j for k, j in enumerate(jobs) if k % n == i]69    import arch, numpy, scipy, pandas70    meta = {"git_sha": git_sha(), "host": socket.gethostname(), "python": platform.python_version(), "arch": arch.__version__,71            "numpy": numpy.__version__, "scipy": scipy.__version__, "pandas": pandas.__version__, "machine": platform.machine(),72            "cpu_count": os.cpu_count(), "shard": a.shard, "n_jobs": len(jobs), "started": time.strftime("%Y-%m-%dT%H:%M:%S%z")}73    Path(a.out).mkdir(parents=True, exist_ok=True)74    log = Path(a.out) / "_runs.jsonl"75    print(json.dumps(meta), flush=True)76    done = 077    t0 = time.time()78    with ProcessPoolExecutor(a.workers) as ex:79        futs = {ex.submit(run_one, j, a.out): j for j in jobs}80        for f in as_completed(futs):81            j = futs[f]82            try:83                rec = f.result()84            except Exception as e:  # noqa: BLE00185                rec = {"spec_id": j.get("model_id"), "symbol": j.get("symbol"), "error": f"{type(e).__name__}: {str(e)[:200]}"}86            rec.update({"git_sha": meta["git_sha"], "host": meta["host"], "ts": time.strftime("%Y-%m-%dT%H:%M:%S%z")})87            with open(log, "a") as fh:88                fh.write(json.dumps(rec) + "\n")89            done += 190            if done % 10 == 0 or "error" in rec:91                print(f"[{done}/{len(jobs)} {time.time()-t0:7.0f}s] {rec.get('symbol')} {rec.get('spec_id')} {rec.get('error','') or rec.get('wall_s','')}", flush=True)92    print("done", done, "jobs in", round(time.time() - t0), "s")939495if __name__ == "__main__":96    main()97