EventBridge Scheduler (nightly, optional)        EventBridge Scheduler (07:30 safety stop)
                              │                                                 │
                              ▼                                                 ▼
 Operator CLI   ──► Kickoff Lambda ──► AgentCore Runtime (LangGraph agent) ──SSM Run Command──► EC2 GPU box
                                          │      │        │                                   (uv, git,
                                          │      │        └──► Amazon Bedrock (Claude)         train.py)
                                          │      └──► AgentCore Memory (insights)                  │
                                          └──► DynamoDB (results) · S3 (train.py, logs, results.tsv) ◄┘

def parse_run_output(text: str) -> RunResult:
    """Parse the output of the remote experiment script.

    The script prints `===EXIT=== <code>`, then the summary metrics grepped from
    run.log, then (only when val_bpb is missing) the tail of run.log after
    `===TAIL===`. A run without a numeric val_bpb counts as a crash.
    """
    exit_match = re.search(r"===EXIT===\s+(-?\d+)", text)
    exit_code = int(exit_match.group(1)) if exit_match else None
    metrics = {name: float(value) for name, value in METRIC_RE.findall(text)}
    tail = text.split("===TAIL===", 1)[1].strip() if "===TAIL===" in text else ""
    if "val_bpb" not in metrics:
        return RunResult(crashed=True, exit_code=exit_code, log_tail=tail[-4000:])
    return RunResult(
        crashed=False,
        val_bpb=metrics["val_bpb"],
        peak_vram_mb=metrics.get("peak_vram_mb", 0.0),
        exit_code=exit_code,
    )

def decide(result: RunResult, best_bpb: float | None) -> str:
    """Return 'keep', 'discard' or 'crash'.

    The first successful run becomes the baseline. After that, a change is kept
    only if it strictly lowers val_bpb.
    """
    if result.crashed:
        return "crash"
    if best_bpb is None:
        return "keep"
    return "keep" if result.val_bpb < best_bpb else "discard"

def parse_proposal(text: str) -> dict:
    """Parse the model's response into {'description', 'hypothesis', 'edits'}."""
    descriptions = _tag("description", text)
    if not descriptions or not descriptions[0].strip():
        raise ProposalError("missing <description>")
    hypotheses = _tag("hypothesis", text)
    edits = []
    for i, block in enumerate(_tag("edit", text)):
        old, new = _tag("old", block), _tag("new", block)
        if len(old) != 1 or len(new) != 1:
            raise ProposalError(f"edit {i}: needs exactly one <old> and one <new>")
        edits.append({"old": old[0], "new": new[0]})
    if not edits:
        raise ProposalError("no <edit> blocks found")
    return {
        "description": _one_line(descriptions[0], 100),
        "hypothesis": _one_line(hypotheses[0], 600) if hypotheses else "",
        "edits": edits,
    }

def apply_edits(source: str, edits: list[dict]) -> str:
    """Apply exact search/replace edits. Each `old` must match exactly once."""
    if not edits:
        raise EditError("no edits to apply")
    updated = source
    for i, edit in enumerate(edits):
        old, new = edit.get("old"), edit.get("new")
        if not isinstance(old, str) or not old.strip() or not isinstance(new, str):
            raise EditError(f"edit {i}: 'old' must be non-empty text and 'new' must be text")
        count = updated.count(old)
        if count == 0:
            raise EditError(f"edit {i}: the <old> text was not found in train.py")
        if count > 1:
            raise EditError(f"edit {i}: the <old> text matches {count} places; include more context")
        updated = updated.replace(old, new, 1)
    if updated == source:
        raise EditError("the edits did not change train.py")
    return updated

def build_graph(deps: Deps):
    def propose(state: ResearchState) -> dict:
        if not state["history"]:
            log.info("experiment 1: baseline")
            baseline = {"description": "baseline", "hypothesis": "unmodified train.py", "edits": []}
            return {"proposal": baseline, "candidate": state["head"]}

        source = deps.box.read_train_py(state["run_tag"])
        insights = deps.memory.recall(RECALL_QUERY)
        feedback = ""
        for attempt in range(1, MAX_PROPOSAL_ATTEMPTS + 1):
            prompt = build_user_prompt(
                gpu=deps.gpu,
                best_bpb=state.get("best_bpb"),
                history=state["history"],
                insights=insights,
                last_crash=state.get("last_crash", ""),
                train_py=source,
                feedback=feedback,
            )
            text = deps.researcher.propose(deps.program, prompt)
            try:
                proposal = parse_proposal(text)
                updated = apply_edits(source, proposal["edits"])
            except (ProposalError, EditError) as err:
                log.warning("proposal attempt %d rejected: %s", attempt, err)
                feedback = (
                    f"Your previous response was rejected: {err}. Follow the response format "
                    "exactly and copy each <old> block verbatim from the current train.py."
                )
                continue
            seq = state["seq"] + 1
            commit = deps.box.commit_train_py(state["run_tag"], seq, updated, proposal["description"])
            log.info("experiment %d: %s (%s)", seq, proposal["description"], commit)
            return {"proposal": proposal, "candidate": commit}
        return {"proposal": None, "candidate": None}

    def execute(state: ResearchState) -> dict:
        if not state.get("proposal"):
            return {"result": None}
        output = deps.box.run_experiment(state["run_tag"], state["seq"] + 1, state["candidate"])
        return {"result": parse_run_output(output).to_dict()}

    def evaluate(state: ResearchState) -> dict:
        seq = state["seq"] + 1
        prev_best = state.get("best_bpb")
        head, best = state["head"], prev_best
        proposal = state.get("proposal")
        if proposal is None:  # the model never produced an applicable edit
            result = RunResult(crashed=True)
            status, commit = "invalid", head
            proposal = {"description": "(no valid edit produced)", "hypothesis": ""}
        else:
            result = RunResult(**state["result"])
            status, commit = decide(result, prev_best), state["candidate"]
            if status == "keep":
                head, best = commit, result.val_bpb
            elif commit != head:
                deps.box.reset_to(head)  # discard: rewind the branch

        row = {
            "seq": seq,
            "commit": commit,
            "status": status,
            "val_bpb": result.val_bpb,
            "memory_gb": result.memory_gb,
            "description": proposal["description"],
            "hypothesis": proposal.get("hypothesis", ""),
            "ts": int(time.time()),
        }
        history = [*state["history"], row]
        _persist(deps, state["run_tag"], row, history, head, best, prev_best, result)
        log.info("experiment %d -> %s val_bpb=%.6f best=%s", seq, status, result.val_bpb, best)
        return {
            "seq": seq,
            "head": head,
            "best_bpb": best,
            "history": history,
            "last_crash": result.log_tail if status == "crash" else "",
            "proposal": None,
            "candidate": None,
            "result": None,
        }

    def route(state: ResearchState) -> str:
        reason = stop_reason(deps, state)
        if reason:
            log.info("stopping: %s", reason)
            return END
        deps.store.refresh_lock(deps.lock_owner)
        return "propose"

    graph = StateGraph(ResearchState)
    graph.add_node("propose", propose)
    graph.add_node("execute", execute)
    graph.add_node("evaluate", evaluate)
    graph.set_entry_point("propose")
    graph.add_edge("propose", "execute")
    graph.add_edge("execute", "evaluate")
    graph.add_conditional_edges("evaluate", route, {"propose": "propose", END: END})
    return graph.compile()

    def _as_user(self, body: str) -> str:
        script = (
            "set -euo pipefail\n"
            'export PATH="$HOME/.local/bin:$PATH"\n'
            f"cd {shlex.quote(self.cfg.repo_dir)}\n" + body
        )
        encoded = base64.b64encode(script.encode()).decode()
        return f'sudo -u {self.cfg.gpu_user} -H bash -lc "$(echo {encoded} | base64 -d)"'

    def run(self, body: str, timeout_s: int = 300, as_root: bool = False) -> str:
        """Run a script on the box and return stdout. Raises on non-zero exit."""
        resp = self.ssm.send_command(
            InstanceIds=[self.cfg.instance_id],
            DocumentName="AWS-RunShellScript",
            Parameters={
                "commands": [body if as_root else self._as_user(body)],
                "executionTimeout": [str(timeout_s)],
            },
            TimeoutSeconds=300,
            Comment="autoresearch",
        )
        command_id = resp["Command"]["CommandId"]
        deadline = time.time() + timeout_s + 300
        time.sleep(2)
        while time.time() < deadline:
            try:
                inv = self.ssm.get_command_invocation(
                    CommandId=command_id, InstanceId=self.cfg.instance_id
                )
            except self.ssm.exceptions.InvocationDoesNotExist:
                time.sleep(3)
                continue
            if inv["Status"] in _PENDING:
                time.sleep(10)
                continue
            if inv["Status"] != "Success":
                raise RemoteCommandError(
                    f"{inv['Status']} (exit {inv.get('ResponseCode')}): "
                    f"{inv.get('StandardErrorContent', '')[-2000:]}"
                )
            return inv.get("StandardOutputContent", "")
        raise RemoteCommandError(f"command {command_id} did not finish in time")

    def run_experiment(self, run_tag: str, seq: int, commit: str) -> str:
        """Run one fixed-budget training job; return the marker-delimited summary."""
        timeout_s = self.cfg.experiment_timeout_s
        log_key = f"runs/{run_tag}/logs/{seq:04d}-{commit}.log"
        body = (
            # Kill leftovers from an interrupted run. The [p] trick stops pkill
            # from matching this script's own command line.
            "pkill -f '[p]ython.*train\\.py' || true\n"
            "rm -f run.log\n"
            "set +e\n"
            f"timeout --kill-after=30 {timeout_s} uv run train.py > run.log 2>&1\n"
            "code=$?\n"
            "set -e\n"
            'echo "===EXIT=== $code"\n'
            'echo "===METRICS==="\n'
            "grep -E '^(val_bpb|peak_vram_mb|training_seconds|num_params_M):' run.log || true\n"
            'echo "===TAIL==="\n'
            "grep -q '^val_bpb:' run.log || tail -n 40 run.log\n"
            + self._cp("run.log", self._s3_uri(log_key)).rstrip("\n")
            + " || true\n"
        )
        return self.run(body, timeout_s=timeout_s + 300)

class ExperimentMemory:
    """Long-term memory of what worked, backed by AgentCore Memory.

    Each finished experiment is written as a short conversational event. The
    memory's semantic strategy distils those events into insight records that
    later sessions (on any night) retrieve before proposing a change.
    """

    def __init__(self, cfg: Config):
        self.client = boto3.client("bedrock-agentcore", region_name=cfg.region)
        self.memory_id = cfg.memory_id
        self.actor_id = cfg.actor_id
        self.namespace = f"/autoresearch/{cfg.actor_id}/insights/"

    def record(self, run_tag: str, experiment: str, outcome: str) -> None:
        try:
            self.client.create_event(
                memoryId=self.memory_id,
                actorId=self.actor_id,
                sessionId=run_tag,
                eventTimestamp=datetime.now(timezone.utc),
                payload=[
                    {"conversational": {"role": "USER", "content": {"text": experiment}}},
                    {"conversational": {"role": "ASSISTANT", "content": {"text": outcome}}},
                ],
            )
        except ClientError as err:
            log.warning("memory write failed: %s", err)

    def recall(self, query: str, top_k: int = 8) -> list[str]:
        try:
            resp = self.client.retrieve_memory_records(
                memoryId=self.memory_id,
                namespace=self.namespace,
                searchCriteria={"searchQuery": query, "topK": top_k},
            )
        except ClientError as err:
            log.warning("memory read failed: %s", err)
            return []
        texts = []
        for record in resp.get("memoryRecordSummaries", []):
            text = (record.get("content") or {}).get("text", "").strip()
            if text:
                texts.append(text)
        return texts

DEFAULT_PROGRAM = """\
You are an autonomous ML researcher. You improve a small single-GPU GPT
training script, train.py, one experiment at a time.

Objective: the lowest val_bpb (validation bits per byte; lower is better).
Training always stops after a fixed 5-minute budget, so you are optimising
what this GPU can learn in 5 minutes.

You may change anything in train.py: architecture, optimizer, learning rates
and schedules, batch size, model size, the training loop. You may not change
prepare.py or the evaluation, and you may not add dependencies.

Constraints and taste:
- The script must run without crashing and finish within the time budget.
- Some extra VRAM is fine for a real gain; don't exceed the GPU's memory.
- Prefer simplicity. A tiny gain that adds a lot of hacky code is not worth
  it; the same result with less code is a win.
- Learn from the history and the insights. Don't repeat ideas that already
  failed unless your variant is meaningfully different. Combine near-misses
  and try bolder ideas when progress stalls.
- If there is no successful run yet, your only goal is to make train.py run
  on this GPU (for example: lower DEPTH or batch size, shorten sequences, or
  replace a kernel the GPU does not support with a PyTorch-native one).
- If the last experiment crashed because of a simple bug, you may fix it.

Respond with ONE experiment in exactly this format:
<description>one line, at most 80 characters</description>
<hypothesis>why this should lower val_bpb</hypothesis>
<edit>
<old>
text copied verbatim from train.py; it must match exactly one place
</old>
<new>
the replacement text
</new>
</edit>
You may use several <edit> blocks. Keep each <old> short but unique.
"""

@app.entrypoint
def invoke(payload: dict) -> dict:
    payload = payload or {}
    action = payload.get("action", "start")
    try:
        params = RunParams.from_payload(payload)
    except ValueError as err:
        return {"error": str(err)}
    cfg = Config.from_env()
    store = ResultsStore(cfg)

    if action == "status":
        state, rows = store.load(params.run_tag)
        return {"run_tag": params.run_tag, "state": state, "recent": rows[-10:]}
    if action != "start":
        return {"error": f"unknown action: {action}"}

    owner = f"{params.run_tag}-{uuid.uuid4().hex[:8]}"
    if not store.acquire_lock(owner):
        return {"status": "busy", "message": "another research session is using the GPU"}

    task_id = app.add_async_task("autoresearch_session", {"run_tag": params.run_tag})

    def worker() -> None:
        try:
            run_session(cfg, params, owner, store)
        except Exception:
            log.exception("research session %s failed", params.run_tag)
        finally:
            app.complete_async_task(task_id)

    threading.Thread(target=worker, name=f"autoresearch-{params.run_tag}", daemon=True).start()
    return {
        "status": "started",
        "run_tag": params.run_tag,
        "max_hours": params.max_hours,
        "max_experiments": params.max_experiments,
    }

def run_session(cfg: Config, params: RunParams, lock_owner: str, store: ResultsStore) -> dict:
    box = GpuBox(cfg)
    try:
        box.ensure_ready()
        box.clear_stop(params.run_tag)
        saved, rows = store.load(params.run_tag)
        head = box.checkout(params.run_tag, reset_to=saved["head"] if saved else None)
        done = saved["experiments"] if saved else 0
        researcher = Researcher(cfg.model_id, cfg.region)
        deps = Deps(
            box=box,
            store=store,
            memory=ExperimentMemory(cfg),
            researcher=researcher,
            metrics=Metrics(cfg),
            program=box.load_program(DEFAULT_PROGRAM),
            gpu=cfg.gpu_description,
            lock_owner=lock_owner,
        )
        state = {
            "run_tag": params.run_tag,
            "head": head,
            "best_bpb": saved.get("best_bpb") if saved else None,
            "seq": done,
            "history": rows,
            "last_crash": "",
            "deadline": time.time() + params.max_hours * 3600 - SAFETY_MARGIN_S,
            "max_experiments": done + params.max_experiments,
        }
        log.info("session %s: resuming after %d experiments at %s", params.run_tag, done, head)
        final = build_graph(deps).invoke(state, config={"recursion_limit": 100_000})
        log.info(
            "session %s finished: %d experiments, best val_bpb=%s, tokens in/out=%d/%d",
            params.run_tag,
            final["seq"],
            final.get("best_bpb"),
            researcher.input_tokens,
            researcher.output_tokens,
        )
        return final
    finally:
        store.release_lock(lock_owner)
        if cfg.stop_instance_on_finish:
            box.stop_instance()

        memory = CfnResource(
            self,
            "Memory",
            type="AWS::BedrockAgentCore::Memory",
            properties={
                "Name": "autoresearch_memory",
                "Description": "What the autoresearch agent learned across nights",
                "EventExpiryDuration": 90,
                "MemoryStrategies": [
                    {
                        "SemanticMemoryStrategy": {
                            "Name": "experiment_insights",
                            "Description": "Facts about which train.py changes helped, hurt or crashed",
                            "NamespaceTemplates": ["/autoresearch/{actorId}/insights/"],
                        }
                    }
                ],
            },
        )

        # ...

        runtime = CfnResource(
            self,
            "Runtime",
            type="AWS::BedrockAgentCore::Runtime",
            properties={
                "AgentRuntimeName": AGENT_NAME,
                "Description": "Autonomous train.py research loop (autoresearch)",
                "AgentRuntimeArtifact": {"ContainerConfiguration": {"ContainerUri": image.image_uri}},
                "NetworkConfiguration": {"NetworkMode": "PUBLIC"},
                "ProtocolConfiguration": "HTTP",
                "RoleArn": agent_role.role_arn,
                "LifecycleConfiguration": {"IdleRuntimeSessionTimeout": 900, "MaxLifetime": 28800},
                "EnvironmentVariables": {
                    "AUTORESEARCH_REGION": r,
                    "GPU_INSTANCE_ID": instance.instance_id,
                    "GPU_DESCRIPTION": gpu_description,
                    "ARTIFACT_BUCKET": bucket.bucket_name,
                    "RESULTS_TABLE": table.table_name,
                    "MEMORY_ID": memory.get_att("MemoryId").to_string(),
                    "MODEL_ID": model_id,
                    "REPO_DIR": REPO_DIR,
                    "BASE_REF": repo_ref,
                    "EXPERIMENT_TIMEOUT_S": "600",
                    "STOP_INSTANCE_ON_FINISH": "true",
                },
            },
        )