generated from mathias/template-go-web
mathias
8eb0358c01
train.py is copied into runs/<rq-id>/ by autoresearch_start.py, so when loop.py executes it, sys.path[0] is the run dir — which has no scripts/. train.py's frozen VaR-eval block does `from scripts.var_breach import ...`, which then fails with ModuleNotFoundError: No module named 'scripts' on every loop run (CI and the documented manual launch alike). Fix in the harness, not the frozen train.py/scripts boundary: prepend the project root (loop.py's own dir, where scripts/ lives) to the subprocess PYTHONPATH. Verified red→green locally: the import fails without it and resolves with it (scripts/ is an implicit namespace package, no __init__.py).
288 lines
10 KiB
Python
288 lines
10 KiB
Python
"""loop.py — Karpathy-style autoresearch loop for jepa-fx-risk.
|
|
|
|
Agent (on iguana/berget — NOT koala, whose GPU is reserved for train.py) reads
|
|
program.md + train.py + STATUS.md, proposes ONE change to train.py, we run it,
|
|
keep if val_vol_r2 improved else git-revert. Appends per-iter record to STATUS.md.
|
|
|
|
LITELLM_KEY=xxx python loop.py [--iters N] [--model MODEL] [--run-dir runs/rq-04]
|
|
|
|
Env:
|
|
LITELLM_KEY — LiteLLM master key (required)
|
|
LITELLM_BASE — default http://localhost:30401/v1
|
|
LOOP_MODEL — default berget/gemma4-31b (non-thinking; iguana/berget only)
|
|
LOOP_ITERS — default 3
|
|
TRAIN_TIMEOUT — seconds per train.py run, default 120
|
|
NTFY_URL — optional: POST crash/stall alerts here (e.g. ntfy.sh/<topic>)
|
|
"""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import textwrap
|
|
from pathlib import Path
|
|
|
|
import urllib.request
|
|
|
|
LITELLM_BASE = os.environ.get("LITELLM_BASE", "http://localhost:30401/v1")
|
|
LITELLM_KEY = os.environ.get("LITELLM_KEY", "")
|
|
LOOP_MODEL = os.environ.get("LOOP_MODEL", "berget/gemma4-31b")
|
|
LOOP_ITERS = int(os.environ.get("LOOP_ITERS", "3"))
|
|
TRAIN_TIMEOUT = int(os.environ.get("TRAIN_TIMEOUT", "120"))
|
|
NTFY_URL = os.environ.get("NTFY_URL", "")
|
|
|
|
# Resolved by main() once --run-dir is parsed.
|
|
RUN_DIR = Path(".")
|
|
STATUS_MD = Path("STATUS.md")
|
|
METRICS_JSON = Path("metrics.json")
|
|
TRAIN_PY = Path("train.py")
|
|
HEARTBEAT = Path("HEARTBEAT")
|
|
|
|
AGENT_SYSTEM = textwrap.dedent("""\
|
|
You are the autoresearch agent for jepa-fx-risk. Your job: propose ONE small,
|
|
targeted change to train.py to improve val_vol_r2 (OOS R² predicting 1-day
|
|
realized vol from frozen embeddings). Higher is better.
|
|
|
|
Rules:
|
|
- Return ONLY the full new content of train.py — nothing else, no explanation,
|
|
no markdown fence. Raw Python only.
|
|
- Change ONE thing at a time (one knob, one structural idea).
|
|
- Do NOT touch prepare_data.py, loop.py, or the data pipeline — only train.py.
|
|
- Do NOT add new data sources or new files.
|
|
- The metric is computed externally from your frozen embeddings; trust it.
|
|
""")
|
|
|
|
|
|
def read_file(p: Path) -> str:
|
|
return p.read_text() if p.exists() else ""
|
|
|
|
|
|
def gpu_snapshot() -> str:
|
|
try:
|
|
out = subprocess.check_output(
|
|
["nvidia-smi", "--query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu",
|
|
"--format=csv,noheader,nounits"], timeout=5, text=True
|
|
).strip()
|
|
util, mem_used, mem_total, temp = [x.strip() for x in out.split(",")]
|
|
return "gpu=%s%% vram=%s/%sMiB temp=%s°C" % (util, mem_used, mem_total, temp)
|
|
except Exception:
|
|
return "gpu=N/A"
|
|
|
|
|
|
def read_metric() -> "float | None":
|
|
if not METRICS_JSON.exists():
|
|
return None
|
|
try:
|
|
return float(json.loads(METRICS_JSON.read_text())["val_vol_r2"])
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def run_train() -> "tuple[float | None, float, str]":
|
|
"""Run train.py from project root with METRICS_OUT pointing into the run dir."""
|
|
t0 = time.time()
|
|
env = dict(os.environ)
|
|
env["METRICS_OUT"] = str(METRICS_JSON.resolve())
|
|
# train.py is copied into the run dir, so sys.path[0] is that run dir — which
|
|
# has no scripts/. Put the project root (where loop.py + scripts/ live) on
|
|
# PYTHONPATH so train.py's `from scripts.var_breach import ...` resolves.
|
|
env["PYTHONPATH"] = str(Path(__file__).resolve().parent) + os.pathsep + env.get("PYTHONPATH", "")
|
|
try:
|
|
r = subprocess.run(
|
|
[sys.executable, str(TRAIN_PY.resolve())],
|
|
capture_output=True, text=True, timeout=TRAIN_TIMEOUT, env=env,
|
|
)
|
|
elapsed = time.time() - t0
|
|
if r.returncode != 0:
|
|
return None, elapsed, (r.stderr or r.stdout)[-300:]
|
|
metric = read_metric()
|
|
return metric, elapsed, ""
|
|
except subprocess.TimeoutExpired:
|
|
return None, TRAIN_TIMEOUT, "TIMEOUT"
|
|
|
|
|
|
def call_agent(iteration: int, best_so_far: "float | None") -> str:
|
|
"""Ask the LLM agent to edit train.py. Returns new train.py content."""
|
|
context = "\n\n".join([
|
|
"# program.md\n" + read_file(RUN_DIR / "program.md"),
|
|
"# train.py (current)\n" + read_file(TRAIN_PY),
|
|
"# STATUS.md (history)\n" + read_file(STATUS_MD)[-2000:],
|
|
"# metrics.json (last run)\n" + read_file(METRICS_JSON),
|
|
"Iteration %d. Best val_vol_r2 so far: %s. Improve it." % (
|
|
iteration, "%.4f" % best_so_far if best_so_far is not None else "none yet"
|
|
),
|
|
])
|
|
payload = json.dumps({
|
|
"model": LOOP_MODEL,
|
|
"messages": [
|
|
{"role": "system", "content": AGENT_SYSTEM},
|
|
{"role": "user", "content": context},
|
|
],
|
|
"temperature": 0.7,
|
|
"max_tokens": 4096,
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
LITELLM_BASE + "/chat/completions",
|
|
data=payload,
|
|
headers={"Authorization": "Bearer " + LITELLM_KEY,
|
|
"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
resp = urllib.request.urlopen(req, timeout=60)
|
|
data = json.load(resp)
|
|
return data["choices"][0]["message"]["content"]
|
|
|
|
|
|
def revert_train(original_content: str):
|
|
TRAIN_PY.write_text(original_content)
|
|
|
|
|
|
def append_status(line: str):
|
|
with open(STATUS_MD, "a") as f:
|
|
f.write(line + "\n")
|
|
|
|
|
|
def write_heartbeat(iteration: int, status: str = "alive"):
|
|
"""Update HEARTBEAT so watchdogs can detect stalls."""
|
|
HEARTBEAT.write_text("%s iter=%d ts=%.0f\n" % (status, iteration, time.time()))
|
|
|
|
|
|
def ntfy(msg: str):
|
|
"""POST an alert to NTFY_URL (best-effort; silently ignored on any error)."""
|
|
if not NTFY_URL:
|
|
return
|
|
try:
|
|
req = urllib.request.Request(
|
|
NTFY_URL, data=msg.encode(), method="POST",
|
|
headers={"Content-Type": "text/plain"},
|
|
)
|
|
urllib.request.urlopen(req, timeout=5)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def main():
|
|
global RUN_DIR, STATUS_MD, METRICS_JSON, TRAIN_PY, HEARTBEAT
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--iters", type=int, default=LOOP_ITERS)
|
|
parser.add_argument("--model", default=LOOP_MODEL)
|
|
parser.add_argument(
|
|
"--run-dir", default=None,
|
|
help="run dir scaffolded by autoresearch_start.py; "
|
|
"STATUS.md, metrics.json, HEARTBEAT, and train.py live here",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
loop_iters = args.iters
|
|
loop_model = args.model
|
|
|
|
if args.run_dir:
|
|
RUN_DIR = Path(args.run_dir)
|
|
if not RUN_DIR.is_dir():
|
|
print("ERROR: run dir not found:", RUN_DIR); sys.exit(1)
|
|
|
|
STATUS_MD = RUN_DIR / "STATUS.md"
|
|
METRICS_JSON = RUN_DIR / "metrics.json"
|
|
TRAIN_PY = RUN_DIR / "train.py"
|
|
HEARTBEAT = RUN_DIR / "HEARTBEAT"
|
|
|
|
if not LITELLM_KEY:
|
|
print("ERROR: set LITELLM_KEY"); sys.exit(1)
|
|
|
|
if not STATUS_MD.exists():
|
|
STATUS_MD.write_text(
|
|
"# Autoresearch STATUS\n\n"
|
|
"| iter | val_vol_r2 | delta | action | secs | gpu | change |\n"
|
|
"|------|-----------|-------|--------|------|-----|--------|\n"
|
|
)
|
|
|
|
baseline = read_metric()
|
|
if baseline is None:
|
|
print("No metrics.json — running train.py for baseline...")
|
|
m, secs, err = run_train()
|
|
if m is None:
|
|
msg = "Baseline run failed: " + err
|
|
print(msg)
|
|
ntfy("[jepa-fx-risk] loop CRASH — " + msg)
|
|
sys.exit(1)
|
|
baseline = m
|
|
print("Baseline: val_vol_r2 = %.4f (%.1fs)" % (baseline, secs))
|
|
|
|
best = baseline
|
|
print("Starting loop | model=%s | iters=%d | baseline=%.4f" % (loop_model, loop_iters, best))
|
|
if args.run_dir:
|
|
print(" run-dir:", RUN_DIR)
|
|
|
|
iter_index = 0
|
|
try:
|
|
for i in range(1, loop_iters + 1):
|
|
iter_index = i
|
|
write_heartbeat(i, "agent-call")
|
|
print("\n--- iter %d/%d ---" % (i, loop_iters))
|
|
original = TRAIN_PY.read_text()
|
|
|
|
print(" calling agent (%s)..." % loop_model)
|
|
t_agent = time.time()
|
|
try:
|
|
new_code = call_agent(i, best)
|
|
except Exception as e:
|
|
msg = str(e)
|
|
print(" agent call failed:", msg)
|
|
append_status("| %d | ERR | — | agent-fail | — | — | %s |" % (i, msg[:60]))
|
|
write_heartbeat(i, "agent-fail")
|
|
ntfy("[jepa-fx-risk] iter %d agent FAIL — %s" % (i, msg[:80]))
|
|
continue
|
|
agent_secs = time.time() - t_agent
|
|
print(" agent replied in %.1fs" % agent_secs)
|
|
|
|
# strip accidental markdown fences
|
|
if new_code.strip().startswith("```"):
|
|
lines = new_code.strip().splitlines()
|
|
new_code = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])
|
|
|
|
TRAIN_PY.write_text(new_code)
|
|
|
|
write_heartbeat(i, "training")
|
|
gpu = gpu_snapshot()
|
|
print(" running train.py [%s]..." % gpu)
|
|
metric, secs, err = run_train()
|
|
|
|
if metric is None:
|
|
print(" train.py FAILED — reverting. err:", err[:100])
|
|
revert_train(original)
|
|
append_status("| %d | FAIL | — | revert | %.0fs | %s | run error |" % (i, secs, gpu))
|
|
write_heartbeat(i, "train-fail")
|
|
ntfy("[jepa-fx-risk] iter %d train FAIL — %s" % (i, err[:80]))
|
|
continue
|
|
|
|
delta = metric - best
|
|
if metric > best:
|
|
best = metric
|
|
action = "KEEP"
|
|
else:
|
|
revert_train(original)
|
|
action = "revert"
|
|
|
|
summary = "| %d | %.4f | %+.4f | %s | %.0fs | %s | iter%d |" % (
|
|
i, metric, delta, action, secs, gpu, i)
|
|
append_status(summary)
|
|
write_heartbeat(i, "done")
|
|
print(" val_vol_r2=%.4f delta=%+.4f action=%s [%.0fs]" % (metric, delta, action, secs))
|
|
|
|
except Exception as e:
|
|
msg = "loop CRASH at iter %d: %s" % (iter_index, e)
|
|
print("FATAL:", msg)
|
|
ntfy("[jepa-fx-risk] " + msg)
|
|
raise
|
|
|
|
print("\nDone. Best val_vol_r2 = %.4f (baseline was %.4f, delta %+.4f)" % (best, baseline, best - baseline))
|
|
print("STATUS.md updated.")
|
|
write_heartbeat(loop_iters, "done")
|
|
ntfy("[jepa-fx-risk] loop done. best val_vol_r2=%.4f (delta %+.4f)" % (best, best - baseline))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|