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