jepa-fx-risk/scripts/autoresearch_start.py

"""autoresearch start — scaffold a run dir from an Autoresearch Council backlog leaf.

Usage:
    python scripts/autoresearch_start.py <backlog.json> <rq-id>

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/<rq-id>/ 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/<rq-id>

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 <backlog.json> <rq-id>")
        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()