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 ru0 = resource.getrusage(resource.RUSAGE_SELF)48 panel = load_panel(js.symbol)49 F, G = run_job(panel, js)50 F.to_parquet(fp, index=False)51 G.to_parquet(d / f"{js.spec_id}.fits.parquet", index=False)52 ru = resource.getrusage(resource.RUSAGE_SELF)53 return {"spec_id": js.spec_id, "symbol": js.symbol, "n_origins": int(len(F)), "n_fits": int(len(G)),54 "status_counts": F["status"].value_counts().to_dict() if len(F) else {},55 "fit_status_counts": G["status"].value_counts().to_dict() if len(G) else {},56 "wall_s": time.perf_counter() - t0, "cpu_s": (ru.ru_utime - ru0.ru_utime) + (ru.ru_stime - ru0.ru_stime), "maxrss_mb": ru.ru_maxrss / (1024 ** 2 if platform.system() == "Darwin" else 1024),57 "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}585960def main() -> None:61 ap = argparse.ArgumentParser()62 ap.add_argument("--jobs", required=True)63 ap.add_argument("--shard", default="0/1")64 ap.add_argument("--workers", type=int, default=max(1, os.cpu_count() - 2))65 ap.add_argument("--out", default=str(ROOT / "results" / "forecasts"))66 a = ap.parse_args()67 jobs = json.loads(Path(a.jobs).read_text())68 i, n = map(int, a.shard.split("/"))69 jobs = [j for k, j in enumerate(jobs) if k % n == i]70 import arch, numpy, scipy, pandas71 meta = {"git_sha": git_sha(), "host": socket.gethostname(), "python": platform.python_version(), "arch": arch.__version__,72 "numpy": numpy.__version__, "scipy": scipy.__version__, "pandas": pandas.__version__, "machine": platform.machine(),73 "cpu_count": os.cpu_count(), "shard": a.shard, "n_jobs": len(jobs), "started": time.strftime("%Y-%m-%dT%H:%M:%S%z")}74 Path(a.out).mkdir(parents=True, exist_ok=True)75 log = Path(a.out) / f"_runs_{socket.gethostname()}.jsonl" # per host: rsync-merge safe76 print(json.dumps(meta), flush=True)77 done = 078 t0 = time.time()79 with ProcessPoolExecutor(a.workers) as ex:80 futs = {ex.submit(run_one, j, a.out): j for j in jobs}81 for f in as_completed(futs):82 j = futs[f]83 try:84 rec = f.result()85 except Exception as e: # noqa: BLE00186 rec = {"spec_id": j.get("model_id"), "symbol": j.get("symbol"), "error": f"{type(e).__name__}: {str(e)[:200]}"}87 rec.update({"git_sha": meta["git_sha"], "host": meta["host"], "ts": time.strftime("%Y-%m-%dT%H:%M:%S%z")})88 with open(log, "a") as fh:89 fh.write(json.dumps(rec) + "\n")90 done += 191 if done % 10 == 0 or "error" in rec:92 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)93 print("done", done, "jobs in", round(time.time() - t0), "s")949596if __name__ == "__main__":97 main()98