generated from mathias/template-go-web
feat(loop): autoresearch keep/revert loop + first iteration (val_vol_r2 0.2821→0.3749, +9.3%)
loop.py: Karpathy-style keep/revert loop. Agent (berget/gemma4-31b, iguana model, NOT koala GPU) proposes one change to train.py per iter → train.py runs on koala GPU (<2s) → read val_vol_r2 from metrics.json → keep if improved, else restore original content. STATUS.md tracks per-iter metric + delta + GPU snap. Iter 1 kept: improved EMBED_DIM/capacity, +9.3% on OOS R². Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Claude Sonnet 4.6
@@ -0,0 +1,205 @@
|
||||
"""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]
|
||||
|
||||
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
|
||||
"""
|
||||
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"))
|
||||
STATUS_MD = Path("STATUS.md")
|
||||
METRICS_JSON = Path("metrics.json")
|
||||
TRAIN_PY = Path("train.py")
|
||||
|
||||
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. Returns (val_vol_r2 or None, wall_secs, stderr_tail)."""
|
||||
t0 = time.time()
|
||||
gpu_before = gpu_snapshot()
|
||||
try:
|
||||
r = subprocess.run(
|
||||
[sys.executable, "train.py"],
|
||||
capture_output=True, text=True, timeout=TRAIN_TIMEOUT,
|
||||
)
|
||||
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(Path("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 main():
|
||||
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")
|
||||
|
||||
# establish baseline
|
||||
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:
|
||||
print("Baseline run failed:", err); 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))
|
||||
|
||||
for i in range(1, LOOP_ITERS + 1):
|
||||
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:
|
||||
print(" agent call failed:", e)
|
||||
append_status("| %d | ERR | — | agent-fail | — | — | %s |" % (i, str(e)[:60]))
|
||||
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)
|
||||
|
||||
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))
|
||||
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)
|
||||
print(" val_vol_r2=%.4f delta=%+.4f action=%s [%.0fs]" % (metric, delta, action, secs))
|
||||
|
||||
print("\nDone. Best val_vol_r2 = %.4f (baseline was %.4f, delta %+.4f)" % (best, baseline, best - baseline))
|
||||
print("STATUS.md updated.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user