From 65a58fcca2c09f2f575675428b8cd00d6d847227 Mon Sep 17 00:00:00 2001 From: Mathias Date: Sat, 27 Jun 2026 10:19:17 +0200 Subject: [PATCH] feat(loop): add --run-dir isolation, heartbeat, ntfy-on-crash; scaffold start command Closes jepa-fx-risk#11 Phase A. - scripts/autoresearch_start.py: scaffold runs// from Council backlog leaf; fail-closed on non-autoresearch-ready; strips candidate_metric; writes program.md + run.json (provenance) + train.py copy. 19 TDD tests. - loop.py: --run-dir flag redirects STATUS.md / metrics.json / HEARTBEAT / train.py into the run dir; METRICS_OUT env var passed to train subprocess so it writes metrics.json to the run dir; heartbeat file written each iter phase; ntfy-on-crash via NTFY_URL env var (best-effort). - train.py: METRICS_OUT env var overrides metrics.json path (default unchanged). Launch: LITELLM_KEY=xxx python loop.py --run-dir runs/rq-04 Co-Authored-By: Claude Sonnet 4.6 --- loop.py | 186 ++++++++++++++++++-------- scripts/autoresearch_start.py | 155 ++++++++++++++++++++++ tests/test_autoresearch_start.py | 215 +++++++++++++++++++++++++++++++ train.py | 3 +- 4 files changed, 504 insertions(+), 55 deletions(-) create mode 100644 scripts/autoresearch_start.py create mode 100644 tests/test_autoresearch_start.py diff --git a/loop.py b/loop.py index 0bdf177..659bc3e 100644 --- a/loop.py +++ b/loop.py @@ -4,7 +4,7 @@ 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] + LITELLM_KEY=xxx python loop.py [--iters N] [--model MODEL] [--run-dir runs/rq-04] Env: LITELLM_KEY — LiteLLM master key (required) @@ -12,6 +12,7 @@ Env: 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/) """ import argparse import json @@ -24,14 +25,19 @@ 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")) +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, @@ -64,7 +70,7 @@ def gpu_snapshot() -> str: return "gpu=N/A" -def read_metric() -> float | None: +def read_metric() -> "float | None": if not METRICS_JSON.exists(): return None try: @@ -73,14 +79,15 @@ def read_metric() -> float | None: return None -def run_train() -> tuple[float | None, float, str]: - """Run train.py. Returns (val_vol_r2 or None, wall_secs, stderr_tail).""" +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() - gpu_before = gpu_snapshot() + env = dict(os.environ) + env["METRICS_OUT"] = str(METRICS_JSON.resolve()) try: r = subprocess.run( - [sys.executable, "train.py"], - capture_output=True, text=True, timeout=TRAIN_TIMEOUT, + [sys.executable, str(TRAIN_PY.resolve())], + capture_output=True, text=True, timeout=TRAIN_TIMEOUT, env=env, ) elapsed = time.time() - t0 if r.returncode != 0: @@ -91,10 +98,10 @@ def run_train() -> tuple[float | None, float, str]: return None, TRAIN_TIMEOUT, "TIMEOUT" -def call_agent(iteration: int, best_so_far: float | None) -> str: +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")), + "# 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), @@ -132,73 +139,144 @@ def append_status(line: str): 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") + 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) + 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)) + print("Starting loop | model=%s | iters=%d | baseline=%.4f" % (loop_model, loop_iters, best)) + if args.run_dir: + print(" run-dir:", RUN_DIR) - for i in range(1, LOOP_ITERS + 1): - print("\n--- iter %d/%d ---" % (i, LOOP_ITERS)) - original = TRAIN_PY.read_text() + 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: - 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) + 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:]) + # 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) + TRAIN_PY.write_text(new_code) - gpu = gpu_snapshot() - print(" running train.py [%s]..." % gpu) - metric, secs, err = run_train() + 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)) - continue + 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" + 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)) + 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__": diff --git a/scripts/autoresearch_start.py b/scripts/autoresearch_start.py new file mode 100644 index 0000000..0ad3bf0 --- /dev/null +++ b/scripts/autoresearch_start.py @@ -0,0 +1,155 @@ +"""autoresearch start — scaffold a run dir from an Autoresearch Council backlog leaf. + +Usage: + python scripts/autoresearch_start.py + +Reads the Council backlog JSON (from agentsquad autoresearch_pipe.py Stage-3 output), +finds the node by rq-id, validates it is autoresearch-ready (fail-closed), then +scaffolds runs// with: + + program.md — hypothesis, single metric (stripped), agent search-space seam + run.json — provenance (strategic_question + council_node) + config + train.py — copy of project train.py (the loop edits this, keeps history clean) + +Launch: + LITELLM_KEY=xxx python loop.py --run-dir runs/ + +Refs: jepa-fx-risk#11, agentsquad#44 +""" + +import json +import shutil +import sys +from datetime import datetime, timezone +from pathlib import Path + + +def load_backlog(path: str) -> dict: + try: + with open(path) as f: + return json.load(f) + except FileNotFoundError: + print(f"error: backlog file not found: {path}", file=sys.stderr) + raise + + +def scaffold_run( + backlog_path_or_dict, + rq_id: str, + run_dir: Path, + train_py_src: Path, +) -> None: + """Scaffold a run dir. Raises SystemExit on any validation failure.""" + if isinstance(backlog_path_or_dict, (str, Path)): + backlog = load_backlog(str(backlog_path_or_dict)) + else: + backlog = backlog_path_or_dict + + # Find node + nodes_by_id = {n["id"]: n for n in backlog.get("nodes", [])} + if rq_id not in nodes_by_id: + print(f"error: rq-id {rq_id!r} not found in backlog", file=sys.stderr) + sys.exit(1) + node = nodes_by_id[rq_id] + + # Fail-closed: only autoresearch-ready nodes may be scaffolded + status = node.get("status", "") + if status != "autoresearch-ready": + print( + f"error: {rq_id} has status {status!r}, not 'autoresearch-ready' — refusing to scaffold", + file=sys.stderr, + ) + sys.exit(1) + + # Guard against overwriting an existing run + if run_dir.exists(): + print( + f"error: {run_dir} already exists — remove it first to re-scaffold", + file=sys.stderr, + ) + sys.exit(1) + + metric = (node.get("candidate_metric") or "").strip() + strategic_q = backlog.get("strategic_question", "") + council_node = node["id"] + generated_at = datetime.now(timezone.utc).isoformat() + + run_dir.mkdir(parents=True) + + # --- program.md --- + program_md = f"""# program.md — {council_node}: {node.get("question", "")[:80]} + +## Provenance +- strategic_question: {json.dumps(strategic_q)} +- council_node: {council_node} (autoresearch-ready; Autoresearch Council backlog) +- generated_at: {generated_at} + +## Hypothesis +{node.get("question", "")} + +## Single validation metric (optimise this, nothing else) +`{metric}` — see eval harness for the exact definition. Only this scalar drives +keep/revert decisions. Report alongside but do NOT optimise: + - Kupiec POF p-value (calibration sanity) + - val_vol_r2 (representation quality guard) + +## What the agent MAY modify (the search space) +- Hyperparameters in train.py (model size, LR, window, patch_len, epochs, etc.) +- Conditioning mechanisms (e.g. JEPA_ENABLE_REGIME toggle) +- Loss function weights and architecture depth + +## Frozen (do NOT touch — keeps the ablation clean) +- Data pipeline and splits (train ≤2021, OOS ≥2022, test 2024 held out) +- The metric definition and scoring code +- loop.py, scripts/, tests/ + +## Experiment loop (per Karpathy autoresearch) +Each iter (≤ time-box): apply ONE change to train.py → run → read +`{metric}` → keep if improved (and Kupiec p-value did not collapse), else revert. +Stop on: target reached, max iters, or K consecutive iters with no improvement. +""" + (run_dir / "program.md").write_text(program_md) + + # --- run.json (provenance + config) --- + run_meta = { + "strategic_question": strategic_q, + "council_node": council_node, + "metric": metric, + "generated_at": generated_at, + "model_tier": "homelab", + "max_iters": 10, + "time_box_minutes": 5, + } + (run_dir / "run.json").write_text(json.dumps(run_meta, indent=2) + "\n") + + # --- train.py (loop edits this copy; project root train.py is the template) --- + shutil.copy(train_py_src, run_dir / "train.py") + + +def main() -> None: + if len(sys.argv) != 3: + print("usage: python scripts/autoresearch_start.py ") + sys.exit(1) + + backlog_path, rq_id = sys.argv[1], sys.argv[2] + + project_root = Path(__file__).parent.parent + run_dir = project_root / "runs" / rq_id + train_py_src = project_root / "train.py" + + scaffold_run(backlog_path, rq_id, run_dir, train_py_src) + + backlog = load_backlog(backlog_path) + nodes_by_id = {n["id"]: n for n in backlog.get("nodes", [])} + metric = (nodes_by_id[rq_id].get("candidate_metric") or "").strip() + + print(f"✓ scaffolded {run_dir}") + print(f" node: {rq_id}") + print(f" metric: {metric}") + print() + print("launch:") + print(f" LITELLM_KEY=xxx python loop.py --run-dir runs/{rq_id}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_autoresearch_start.py b/tests/test_autoresearch_start.py new file mode 100644 index 0000000..45c1443 --- /dev/null +++ b/tests/test_autoresearch_start.py @@ -0,0 +1,215 @@ +"""Tests for scripts/autoresearch_start.py — jepa-fx-risk#11 Phase A scaffold. + +Success criterion: `autoresearch start ` scaffolds a +runnable run dir from a ready leaf; refuses non-ready nodes; strips +candidate_metric; records provenance. +""" + +import importlib.util +import json +import sys +from pathlib import Path + +import pytest + +# Load the module without executing main() +_SCRIPT = Path(__file__).parent.parent / "scripts" / "autoresearch_start.py" + + +def _import(): + spec = importlib.util.spec_from_file_location("autoresearch_start", _SCRIPT) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +@pytest.fixture() +def mod(): + return _import() + + +@pytest.fixture() +def backlog(tmp_path): + data = { + "strategic_question": "Test strategic question?", + "generated_at": "2026-06-27T00:00:00Z", + "nodes": [ + { + "id": "rq-01", + "question": "Does X improve Y?", + "case_type": "autoresearch-loop", + "data": "obtainable", + "method": "adjacent", + "falsifiable": "yes", + "candidate_metric": " val_vol_r2", # leading space — bypass test + "depends_on": [], + "status": "autoresearch-ready", + "track": "autoresearch", + "converged": True, + "survived_review": True, + }, + { + "id": "rq-02", + "question": "Not ready yet?", + "case_type": "empirical-study", + "data": "obtainable", + "method": "adjacent", + "falsifiable": "yes", + "candidate_metric": None, + "depends_on": [], + "status": "needs-metric", + "track": "study", + "converged": True, + "survived_review": True, + }, + { + "id": "rq-03", + "question": "A spike.", + "case_type": "spike", + "data": "have", + "method": "yes-named", + "falsifiable": "yes", + "candidate_metric": None, + "depends_on": [], + "status": "spike-ready", + "track": "spike", + "converged": True, + "survived_review": True, + }, + ], + } + p = tmp_path / "backlog.json" + p.write_text(json.dumps(data)) + return p + + +@pytest.fixture() +def fake_train_py(tmp_path): + """Minimal train.py placeholder for scaffold tests.""" + src = tmp_path / "train_template.py" + src.write_text("# train.py placeholder\n") + return src + + +# --------------------------------------------------------------------------- +# fail-closed: refuse non-autoresearch-ready nodes +# --------------------------------------------------------------------------- + +class TestRefuseNonReady: + def test_refuses_needs_metric(self, mod, backlog, fake_train_py, tmp_path): + run_dir = tmp_path / "runs" / "rq-02" + with pytest.raises(SystemExit) as exc: + mod.scaffold_run(backlog, "rq-02", run_dir, fake_train_py) + assert exc.value.code != 0 + + def test_refuses_spike_ready(self, mod, backlog, fake_train_py, tmp_path): + run_dir = tmp_path / "runs" / "rq-03" + with pytest.raises(SystemExit) as exc: + mod.scaffold_run(backlog, "rq-03", run_dir, fake_train_py) + assert exc.value.code != 0 + + def test_refuses_missing_rq_id(self, mod, backlog, fake_train_py, tmp_path): + run_dir = tmp_path / "runs" / "rq-99" + with pytest.raises(SystemExit) as exc: + mod.scaffold_run(backlog, "rq-99", run_dir, fake_train_py) + assert exc.value.code != 0 + + def test_refuses_existing_run_dir(self, mod, backlog, fake_train_py, tmp_path): + run_dir = tmp_path / "runs" / "rq-01" + run_dir.mkdir(parents=True) + with pytest.raises(SystemExit) as exc: + mod.scaffold_run(backlog, "rq-01", run_dir, fake_train_py) + assert exc.value.code != 0 + + +# --------------------------------------------------------------------------- +# scaffold structure: correct files created +# --------------------------------------------------------------------------- + +class TestScaffoldStructure: + @pytest.fixture(autouse=True) + def _scaffold(self, mod, backlog, fake_train_py, tmp_path): + self.run_dir = tmp_path / "runs" / "rq-01" + mod.scaffold_run(backlog, "rq-01", self.run_dir, fake_train_py) + + def test_run_dir_created(self): + assert self.run_dir.is_dir() + + def test_program_md_created(self): + assert (self.run_dir / "program.md").exists() + + def test_run_json_created(self): + assert (self.run_dir / "run.json").exists() + + def test_train_py_copied(self): + assert (self.run_dir / "train.py").exists() + assert (self.run_dir / "train.py").read_text() == "# train.py placeholder\n" + + +# --------------------------------------------------------------------------- +# program.md content +# --------------------------------------------------------------------------- + +class TestProgramMd: + @pytest.fixture(autouse=True) + def _scaffold(self, mod, backlog, fake_train_py, tmp_path): + self.run_dir = tmp_path / "runs" / "rq-01" + mod.scaffold_run(backlog, "rq-01", self.run_dir, fake_train_py) + self.content = (self.run_dir / "program.md").read_text() + + def test_contains_hypothesis(self): + assert "Does X improve Y?" in self.content + + def test_metric_key_stripped(self): + # candidate_metric had leading space " val_vol_r2" — must be stripped + assert "`val_vol_r2`" in self.content + assert "` val_vol_r2`" not in self.content + + def test_contains_strategic_question(self): + assert "Test strategic question?" in self.content + + def test_contains_council_node(self): + assert "rq-01" in self.content + + +# --------------------------------------------------------------------------- +# run.json provenance +# --------------------------------------------------------------------------- + +class TestRunJson: + @pytest.fixture(autouse=True) + def _scaffold(self, mod, backlog, fake_train_py, tmp_path): + self.run_dir = tmp_path / "runs" / "rq-01" + mod.scaffold_run(backlog, "rq-01", self.run_dir, fake_train_py) + self.run = json.loads((self.run_dir / "run.json").read_text()) + + def test_strategic_question_in_provenance(self): + assert self.run["strategic_question"] == "Test strategic question?" + + def test_council_node_in_provenance(self): + assert self.run["council_node"] == "rq-01" + + def test_metric_stripped_in_provenance(self): + assert self.run["metric"] == "val_vol_r2" + assert self.run["metric"] == self.run["metric"].strip() + + def test_generated_at_present(self): + assert "generated_at" in self.run + + def test_max_iters_present(self): + assert "max_iters" in self.run + + +# --------------------------------------------------------------------------- +# load_backlog helper +# --------------------------------------------------------------------------- + +class TestLoadBacklog: + def test_loads_json(self, mod, backlog): + data = mod.load_backlog(str(backlog)) + assert data["strategic_question"] == "Test strategic question?" + assert len(data["nodes"]) == 3 + + def test_missing_file_raises(self, mod, tmp_path): + with pytest.raises((FileNotFoundError, SystemExit)): + mod.load_backlog(str(tmp_path / "nonexistent.json")) diff --git a/train.py b/train.py index 0eea9cc..c276f64 100644 --- a/train.py +++ b/train.py @@ -298,12 +298,13 @@ def main(): phase1_r2 = float(1 - ((yte - pred_h) ** 2).sum() / ss_tot) print("phase1_r2 = %.4f (n_test=%d)" % (phase1_r2, len(yte))) + _metrics_out = os.environ.get("METRICS_OUT", "metrics.json") json.dump({ "val_vol_r2": val_vol_r2, "phase1_r2": phase1_r2, "n_test": len(yte), "knobs": {"WINDOW": WINDOW, "PATCH_LEN": PATCH_LEN, "D_MODEL": D_MODEL, "DEPTH": DEPTH, "ALPHA": ALPHA, "DELTA_T_MAX": DELTA_T_MAX, "EPOCHS": EPOCHS}, - }, open("metrics.json", "w"), indent=2) + }, open(_metrics_out, "w"), indent=2) print("val_vol_r2 = %.4f (n_test=%d, dev=%s)" % (val_vol_r2, len(yte), dev)) # ── EXPORT BLOCK — do NOT edit (agent boundary) ──────────────────────────