opal-optimizer 0.3.5__py3-none-any.whl

opal/backends/unavailable.py

@@ -0,0 +1,161 @@
+"""Waiting for an agent that is unavailable for reasons that are not its own doing.
+
+An agent is UNAVAILABLE when its CLI cannot run a turn at all: Claude Code is not logged in (or its key is invalid), a
+usage, spend or rate limit is reached, the API is overloaded or the network is down, the CLI is missing; with a mailbox,
+no broker is answering. Such a decision is not failed: it is kept and waited for, so that a run whose agent is down for
+a while (a subscription switched, a laptop with the broker asleep) makes the same decisions as one whose agent never
+was. The wait has no deadline unless `unavailable_timeout` (seconds) is given; it does not count against the decision's
+own `timeout`, and the agent is told nothing about it. One line says when the wait starts (what is wrong and what to
+do), a reminder follows every REMIND seconds, and one line says when it ends.
+
+Failures of the agent's own turn (a crash inside the call, a time limit, a budget it spent) are not this: they keep the
+handling of backends/claude.py (calls made again while the decision's time allows, then a failed decision).
+"""
+
+from __future__ import annotations
+
+import datetime
+import re
+import time
+
+RECHECK, RECHECK_MAX = 60.0, 300.0  # seconds before the first new try, and the longest gap between tries
+REMIND = 600.0  # seconds between reminders that the run is still waiting
+RESET_MARGIN = 30.0  # seconds after a reset time the message names before trying again
+
+# Claude Code's words for the kinds of unavailability (2.1.30x; "You've hit your org's monthly spend limit · run
+# /usage-credits ... · your session limit resets 1:10am (America/New_York)", "Not logged in · Please run /login",
+# "API Error: 529 overloaded")
+LOGIN = re.compile(
+    r"not logged in|run /login|invalid api key|authentication[_ ]error|oauth token (?:has )?expired", re.I
+)
+# Claude Code's own words for a usage or spend limit: a call whose answer is one of these did not run, also when the CLI
+# exits 0 with it as its answer
+LIMIT = re.compile(
+    r"hit your (?:[\w']+ ){0,3}limit|usage limit reached|spend limit reached|credit balance (?:is )?too low", re.I
+)
+RATE = re.compile(r"rate[_ ]limit|too many requests|\b429\b", re.I)  # in the error of a failed call only
+DOWN = re.compile(
+    r"overloaded|API Error: 5\d\d|\b529\b|ECONNREFUSED|ECONNRESET|ENOTFOUND|ETIMEDOUT|EAI_AGAIN|fetch failed"
+    r"|socket hang up|connection (?:error|refused|reset|timed out)|network (?:error|is unreachable)|unable to connect",
+    re.I,
+)
+HINTS = dict(
+    login="log in on this machine (run `claude`, then /login); the decision goes on by itself",
+    limit="wait for the limit to reset, raise it, or log in to another account; the decision goes on by itself",
+    down="the API or the network is down; the decision goes on by itself once it is back",
+    missing="install Claude Code or put it on PATH; the decision goes on by itself",
+    broker="start a broker for this run; the decision goes on by itself once one answers",
+)
+
+
+def classify(texts):
+    """('login' | 'limit' | 'down', the CLI's words) for the first short text that says the agent is unavailable, else
+    None. Only short texts count, so that an agent's own long answer that mentions a rate limit is never taken for
+    one."""
+    texts = [str(t).strip() for t in texts if t and len(str(t).strip()) < 500]
+    for kind, pattern in (("login", LOGIN), ("limit", LIMIT), ("limit", RATE), ("down", DOWN)):
+        for t in texts:
+            if pattern.search(t):
+                return kind, " ".join(t.split())[:300]
+    return None
+
+
+_RESET = re.compile(r"resets\s+(\d{1,2})(?::(\d{2}))?\s*([ap]m)?(?:\s*\(([A-Za-z_]+/[A-Za-z_]+)\))?", re.I)
+
+
+def seconds_to_reset(text, now=None):
+    """Seconds until just after the reset time a limit message names ("resets 1:10am (America/New_York)", "resets
+    3pm"), the next such time from now; None if it names none."""
+    m = _RESET.search(str(text or ""))
+    if not m:
+        return None
+    hour, minute, ampm, zone = int(m.group(1)), int(m.group(2) or 0), (m.group(3) or "").lower(), m.group(4)
+    if ampm == "pm" and hour != 12:
+        hour += 12
+    elif ampm == "am" and hour == 12:
+        hour = 0
+    if not (0 <= hour < 24 and 0 <= minute < 60):
+        return None
+    tz = None
+    if zone:
+        try:
+            from zoneinfo import ZoneInfo
+
+            tz = ZoneInfo(zone)
+        except Exception:
+            tz = None
+    now_dt = datetime.datetime.fromtimestamp(time.time() if now is None else now, tz=tz)
+    at = now_dt.replace(hour=hour, minute=minute, second=0, microsecond=0)
+    if at <= now_dt:
+        at += datetime.timedelta(days=1)
+    return (at - now_dt).total_seconds() + RESET_MARGIN
+
+
+def _minutes(seconds):
+    m = seconds / 60
+    return f"{seconds:.0f} s" if seconds < 90 else f"{m:.0f} min" if m < 90 else f"{m / 60:.1f} h"
+
+
+class Wait:
+    """One wait for an unavailable agent. note(kind, reason) at each failed try: prints the first line, then a reminder
+    every REMIND seconds (and when the reason changes), and returns the seconds to wait before the next try (None once
+    `timeout` is over). end() prints the last line and returns the seconds waited. every: the caller checks again at
+    this fixed interval (a mailbox's poll) instead of after 60 s, then at most every 5 min."""
+
+    def __init__(self, say, what="Claude Code", training_waits=None, timeout=None, every=None):
+        self.say, self.what, self.timeout, self.every = say, what, timeout, every
+        self.training_waits = training_waits
+        self.since = self.reminded = time.time()
+        self.tries = 0
+        self._retry_wait = min(RECHECK, RECHECK_MAX)
+        self.kind = self.reason = None
+
+    @property
+    def seconds(self):
+        return time.time() - self.since
+
+    def over(self):
+        return self.timeout is not None and self.seconds >= self.timeout
+
+    def note(self, kind, reason):
+        """Seconds to wait before the next try, or None when unavailable_timeout is over (the decision then fails)."""
+        self.tries += 1
+        changed = (kind, reason) != (self.kind, self.reason)
+        self.kind, self.reason = kind, reason
+        if self.tries == 1:
+            run = (
+                "training waits for the agent"
+                if self.training_waits
+                else "training continues with the current programs"
+                if self.training_waits is False
+                else "the request is kept"
+            )
+            limit = "" if self.timeout is None else f", for at most {_minutes(self.timeout)} (unavailable_timeout)"
+            again = (
+                f"checks again every {self.every:g} s"
+                if self.every is not None
+                else f"tries again in {RECHECK:.0f} s, then at most every {RECHECK_MAX / 60:.0f} min"
+            )
+            self.say(
+                f"{self.what} is unavailable ({reason}); {run} and {again}{limit}. To go on: "
+                f"{HINTS.get(kind, HINTS['down'])}"
+            )
+        elif time.time() - self.reminded >= REMIND or (changed and kind != "down"):
+            self.reminded = time.time()
+            self.say(f"still waiting for {self.what} ({_minutes(self.seconds)} so far): {reason}")
+        if self.over():
+            self.say(f"{self.what} stayed unavailable for {_minutes(self.seconds)} (unavailable_timeout); giving up")
+            return None
+        wait = self._retry_wait
+        self._retry_wait = min(2 * wait, RECHECK_MAX)
+        reset = seconds_to_reset(reason)
+        if reset is not None:
+            wait = min(wait, reset)
+        if self.timeout is not None:
+            wait = min(wait, max(0.0, self.timeout - self.seconds))
+        return wait
+
+    def end(self):
+        waited = self.seconds
+        self.say(f"{self.what} is available again after {_minutes(waited)}; the decision goes on")
+        return waited

opal/cli.py

@@ -0,0 +1,396 @@
+"""opal broker | say | report | show | doctor | estimate"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+
+def _backend(args):
+    from .backends import make_backend
+
+    opts = {}
+    if args.model:
+        opts["model"] = args.model
+    if args.effort:
+        opts["effort"] = args.effort
+    if getattr(args, "timeout", None):
+        opts["timeout"] = args.timeout
+    if getattr(args, "retries", None) is not None:
+        opts["retries"] = args.retries
+    if getattr(args, "unavailable_timeout", None) is not None:
+        opts["unavailable_timeout"] = args.unavailable_timeout
+    if getattr(args, "max_tokens", None) is not None:
+        opts["max_tokens"] = args.max_tokens
+    return make_backend(args.backend, **opts)
+
+
+def _run_folder(path):
+    run = Path(path)
+    if (run / "events.jsonl").exists():
+        return run
+    inside = sorted(d.name for d in run.iterdir() if (d / "events.jsonl").exists()) if run.is_dir() else []
+    hint = f"; runs inside: {', '.join(inside[-5:])}" if inside else ""
+    sys.exit(f"opal: not a run folder: {run} (no events.jsonl){hint}")
+
+
+def _show(run):
+    from .report import _jsonl, _response
+
+    events = _jsonl(run / "events.jsonl")  # a line being written is skipped
+    start = next((e for e in events if e["kind"] == "start"), {})
+    print(f"run {run}")
+    if start.get("goal"):
+        print(f"goal: {start['goal']}")
+    from .report import decision_counts, describe_counts, rank_repair_summary
+
+    mode = start.get("blocking")
+    mode = "auto-blocking" if mode == "auto" else "blocking" if mode else "non-blocking"
+    print(f"{start.get('total_updates', '?')} updates; backend {start.get('backend')}; {mode}")
+    repaired = rank_repair_summary(events)
+    if repaired:
+        print(f"WARNING: {repaired}; the run's trajectory was changed.")
+    print(describe_counts(decision_counts(events)))
+    print(f"{'k':>4}  {'asked at':>9}  {'outcome':<11} {'at':>9}  {'secs':>6}  {'cost':>6}  rationale")
+    req = {e["k"]: e for e in events if e["kind"] == "requested"}
+    out = {
+        e["k"]: e
+        for e in events
+        if e["kind"] in ("applied", "kept", "rejected", "load_failed", "cancelled") and e.get("k")
+    }
+    for k in sorted(req):
+        o = out.get(k, {})
+        resp = _response(run, k)
+        said = " ".join((resp.get("text") or "").split())[:70]
+        kind = o.get("kind", "pending")
+        if kind == "pending" and any(
+            e["kind"] in ("resumed", "end") and e.get("utc", "") > req[k].get("utc", "") for e in events
+        ):
+            kind = "lost"
+        if kind == "applied":
+            kind = "changed"
+        elif o.get("failed") and not o.get("gave_up"):
+            kind = "failed"
+        cost = resp.get("cost_usd")
+        print(
+            f"{k:>4}  {req[k]['update']:>9,}  {kind:<11} {o.get('applied_at', ''):>9}  "
+            f"{resp.get('seconds', ''):>6}  {('$%.2f' % cost) if cost else '':>6}  {said}"
+        )
+    for e in events:
+        if e["kind"] in ("reverted", "rank_divergence", "ranks_resynced", "resumed", "interrupted", "crashed"):
+            print(f"{e['kind']} at update {e.get('update')}")
+    end = next((e for e in reversed(events) if e["kind"] == "end"), None)
+    if end:
+        tokens = "usage" in end or start.get("backend") == "codex"
+        cost = "USD unavailable" if tokens and end.get("cost_usd") is None else f"${end.get('cost_usd') or 0:.2f}"
+        print(
+            f"finished at update {end['update']:,}; {end.get('decisions')} decisions; {cost}"
+        )
+
+
+def doctor_mailbox(spec):
+    """A mailbox setup: no Claude Code here; the folder must be usable, and a broker should be serving it."""
+    import json
+    import shutil
+    import subprocess
+    import time
+    import uuid
+
+    from .backends import mailbox
+
+    box = Path(spec.split(":", 1)[1]).expanduser()
+    print(f"mailbox {box}: the agent runs on the broker's machine; no Claude Code is needed here")
+    python = shutil.which("python3")
+    if python is None:
+        print("the mailbox host needs python3 (Python 3.6+) on its non-interactive SSH PATH; install it or fix PATH")
+        return 1
+    try:
+        probe = subprocess.run([python, "-c", "import sys; sys.exit(sys.version_info < (3, 6))"],
+                               capture_output=True, timeout=10)
+    except (OSError, subprocess.TimeoutExpired) as exc:
+        print(f"the mailbox host's python3 could not run: {exc}")
+        return 1
+    if probe.returncode:
+        print("the mailbox host needs python3 version 3.6 or newer on its non-interactive SSH PATH")
+        return 1
+    print("the mailbox host's python3: Python 3.6+ available (also required on the broker's SSH PATH)")
+    try:
+        box.mkdir(mode=0o700, parents=True, exist_ok=True)
+        probe = box / f".doctor-{uuid.uuid4().hex[:8]}"
+        probe.write_text("ok")
+        probe.unlink()
+    except OSError as exc:
+        print(f"the mailbox folder cannot be written: {exc}")
+        return 1
+    print("the mailbox folder: writable")
+    beats = [box / "broker-alive"] + [d / "broker-alive" for d in box.iterdir() if d.is_dir()]
+    alive = []
+    for f in beats:
+        try:
+            if f.exists() and time.time() - f.stat().st_mtime < mailbox.HEARTBEAT_STALE:
+                alive.append(json.loads(f.read_text() or "{}").get("broker") or "a broker")
+        except (OSError, ValueError):
+            continue
+    command = (
+        f"opal broker {box} --ssh <you>@<this host> --backend claude --model claude-opus-5-5 --effort high "
+        "--keep-serving"
+    )
+    if alive:
+        print(f"a broker is serving this mailbox ({', '.join(sorted(set(alive)))})")
+        return 0
+    print(
+        "no broker is serving this mailbox yet. On a machine with Claude Code logged in, start one, and keep that "
+        f"machine awake for the whole run (on a Mac: caffeinate -i {command})"
+    )
+    return 1
+
+
+def doctor(a):
+    import platform
+    import shutil
+    import subprocess
+    import tempfile
+
+    import torch
+
+    ok = True
+    if a.backend.startswith("mailbox:"):
+        return doctor_mailbox(a.backend)
+    print(f"python {sys.version.split()[0]} ({sys.executable})")
+    print(
+        f"torch {torch.__version__}; CUDA available: {torch.cuda.is_available()}"
+        + (f" ({torch.cuda.device_count()} x {torch.cuda.get_device_name(0)})" if torch.cuda.is_available() else "")
+    )
+    for name in ("codex" if a.backend == "codex" else "claude",):
+        exe = shutil.which(name)
+        ver = ""
+        if exe:
+            try:
+                ver = subprocess.run([exe, "--version"], capture_output=True, text=True, timeout=30).stdout.strip()
+            except (OSError, subprocess.TimeoutExpired):
+                ver = "(version unknown)"
+        print(f"{name}: {exe + ' ' + ver if exe else 'not found'}")
+    if platform.system() == "Linux":
+        for tool, what in (("bwrap", "bubblewrap"), ("socat", "socat")):
+            print(f"{what} (for the sandbox of the agent's shell): {'found' if shutil.which(tool) else 'not found'}")
+    from .backends import make_backend
+
+    # one call: a failed one (not logged in, a limit) is reported at once, not waited for
+    be = make_backend(a.backend, retries=0, timeout=300, unavailable_timeout=0)
+    try:
+        be.preflight()
+        if a.backend == "codex":
+            print(f"isolated CODEX_HOME: {be.codex_home} (default home and linked auth rejected)")
+            logged_in = be.login_status()
+            print(f"isolated Codex login: {'ready' if logged_in else 'not logged in'}")
+            if not logged_in:
+                print("run CODEX_HOME=~/.codex-opal codex login --device-auth")
+                return 1
+        print(
+            f"backend {a.backend}: found"
+            + ("" if a.live or a.backend == "codex" else "; login not checked (run `opal doctor --live`, a few cents)")
+        )
+    except RuntimeError as exc:
+        print(str(exc))
+        return 1
+    problem = getattr(be, "sandbox_problem", None)
+    if problem:
+        print(f"the sandbox for the agent's shell cannot start ({problem}): the agent gets file tools only, no shell")
+    elif getattr(be, "sandbox", None) is not None:
+        print(f"the agent's shell: {'in the OS sandbox' if be.sandbox else 'none (file tools only)'}")
+    if a.live:
+        with tempfile.TemporaryDirectory() as tmp:
+            ws = Path(tmp) / "workspace"
+            ws.mkdir()
+            be.setup(Path(tmp), ws)
+            out = be.decide(ws, "Reply with the single word READY and do nothing else.", "You are a test.", Path(tmp))
+            reply = (out.get("text") or "").strip()
+            ok = out.get("returncode") == 0 and "READY" in reply
+            print(
+                f"live call: {'ok' if ok else 'failed'} ({out.get('seconds', 0):.0f} s; reply {reply[:80]!r}; error {out.get('error')})"
+            )
+            if a.backend == "codex":
+                print(f"Codex thread: {out.get('session')}; provider rollout: {out.get('rollout')}")
+                print(f"tokens (USD unavailable): {json.dumps(out.get('usage'))}")
+    return 0 if ok else 1
+
+
+def estimate_cmd(a):
+    from . import estimate
+    from .schedule import DEFAULT_DECISIONS, decision_updates
+
+    explicit = a.decisions is not None
+    n = len(decision_updates(a.total_steps, a.decisions if explicit else DEFAULT_DECISIONS, a.every))
+    capped_from = None
+    if not explicit and a.blocking:
+        k = estimate.capped_decisions(n, a.total_steps * a.sec_per_step, a.blocking, a.sec_per_step)
+        if k < n:
+            capped_from, n = n, len(decision_updates(a.total_steps, k, a.every))
+    print(f"OPAL estimate, from our runs of {estimate.MODEL} at effort {estimate.EFFORT} (yours may differ):")
+    for line in estimate.lines(
+        a.total_steps, a.sec_per_step, n, a.blocking, model=a.model, effort=a.effort, mailbox=a.mailbox,
+        capped_from=capped_from, explicit=explicit,
+    ):  # fmt: skip
+        print("  " + line)
+    print("  These are estimates from our runs (opal/estimate.py says which); the run prints its own at its first check.")
+    return 0
+
+
+def main(argv=None):
+    ap = argparse.ArgumentParser(prog="opal")
+    sub = ap.add_subparsers(dest="cmd", required=True)
+
+    b = sub.add_parser("broker", help="answer a training run's mailbox with a local agent (laptop or login node)")
+    b.add_argument("box", help="mailbox folder of the run (on the GPU machine when --ssh is given)")
+    b.add_argument("--ssh", help="user@host of the GPU machine; omit when the folder is on a shared filesystem")
+    b.add_argument(
+        "--backend",
+        default="claude",
+        choices=["claude", "codex", "keep"],
+        help="keep: answer every decision with 'keep the programs' (to test the mailbox path at no cost)",
+    )
+    b.add_argument("--model")
+    b.add_argument("--effort")
+    b.add_argument("--max-tokens", type=int, help="Codex only: soft run cap on input + output tokens; an in-flight turn may exceed it")
+    b.add_argument(
+        "--workdir", help="local folder for the agent's workspace and records (default ~/.opal/broker/<name>)"
+    )
+    b.add_argument("--poll", type=float, default=2.0, help="seconds between checks of the mailbox")
+    b.add_argument("--run", help="run id (subfolder of the box) to serve; default: the newest unfinished run")
+    b.add_argument(
+        "--timeout",
+        type=int,
+        help="seconds per decision; by default the limit the run's agent was told (the mailbox timeout of the run); "
+        "a shorter value here is kept, and the agent is told it in each message",
+    )
+    b.add_argument(
+        "--retries",
+        type=int,
+        help="how often a failed agent turn (for example a tool crash) is made again, after 15 s, 30 s, "
+        "... (at most 10 min apart) while the decision's time allows; default: as often as that time allows; 0: never",
+    )
+    b.add_argument(
+        "--unavailable-timeout",
+        type=float,
+        help="seconds an unavailable agent (not logged in, a usage or spend limit, the API overloaded, no network) is "
+        "waited for before the request is answered as failed; default: no deadline",
+    )
+    b.add_argument("--keep-serving", action="store_true", help="after a run ends, wait for the next run in the box")
+    b.add_argument(
+        "--give-up-after-min",
+        type=float,
+        help="exit with an error if the mailbox stays unreadable this long (default: retry forever)",
+    )
+
+    s = sub.add_parser("say", help="leave a note for the agent's next decision")
+    s.add_argument("run")
+    s.add_argument("text")
+
+    r = sub.add_parser("report", help="write <run>/report.html: loss, decisions, rationales and code diffs")
+    r.add_argument("run")
+    r.add_argument(
+        "--summarize",
+        action="store_true",
+        help="first describe and tag the code changes that have no summary yet (Claude Code CLI, one small call each)",
+    )
+    r.add_argument("--model", help="model for --summarize (default claude-opus-5-5 at low effort)")
+
+    w = sub.add_parser("show", help="print a run's decisions and where its final program is")
+    w.add_argument("run")
+
+    d = sub.add_parser("doctor", help="check that torch and an agent CLI are ready on this machine")
+    d.add_argument("--live", action="store_true", help="also make one tiny agent call (costs a few cents)")
+    d.add_argument(
+        "--backend",
+        default="claude",
+        help='claude (default), codex, or mailbox:<folder> on a GPU machine whose agent runs elsewhere (checks the folder and '
+        "that a broker serves it)",
+    )
+
+    e = sub.add_parser(
+        "estimate", help="before a run: how long training will wait for the agent, and what the agent will cost"
+    )
+    e.add_argument("--total-steps", type=int, required=True, help="updates in the run (wrap's total_steps)")
+    e.add_argument(
+        "--sec-per-step", type=float, required=True, help="seconds per update without the agent (time the smoke test)"
+    )
+    e.add_argument(
+        "--decisions",
+        type=int,
+        help="as in wrap(decisions=...); not given: OPAL's default, 50, lowered so that the expected waiting stays "
+        "within 4x the run's own time",
+    )
+    e.add_argument("--every", type=int, default=50, help="as in wrap(every=...): decisions fall on checks")
+    e.add_argument(
+        "--blocking",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="training waits for each decision (the default); --no-blocking: it goes on while the agent thinks",
+    )
+    e.add_argument("--model", default="claude-opus-5-5")
+    e.add_argument("--effort", default="high")
+    e.add_argument("--mailbox", action="store_true", help="the agent answers through a mailbox broker")
+
+    a = ap.parse_args(argv)
+    if a.cmd == "doctor" and a.backend not in ("claude", "codex") and not a.backend.startswith("mailbox:"):
+        ap.error("doctor --backend must be claude, codex or mailbox:<folder>")
+    if a.cmd == "estimate":
+        return estimate_cmd(a)
+    if a.cmd == "broker":
+        from .backends.mailbox import run_broker
+
+        name = (a.ssh or "local").replace("@", "_") + "_" + a.box.strip("/").replace("/", "_")
+        workdir = Path(a.workdir) if a.workdir else Path("~/.opal/broker").expanduser() / name
+        run_broker(
+            a.box,
+            _backend(a),
+            workdir,
+            ssh=a.ssh,
+            poll=a.poll,
+            run=a.run,
+            keep_serving=a.keep_serving,
+            give_up_after=60 * a.give_up_after_min if a.give_up_after_min else None,
+        )
+    elif a.cmd == "say":
+        run = _run_folder(a.run)
+        with open(run / "notes.txt", "a") as f:
+            f.write(a.text.strip() + "\n")
+        from .report import _jsonl
+
+        kinds = [e.get("kind") for e in _jsonl(run / "events.jsonl")]
+        last = lambda kind: max((i for i, k in enumerate(kinds) if k == kind), default=-1)  # noqa: E731
+        ended = last("end") > last("resumed")
+        print(
+            "noted, but this run has finished: no agent will read it"
+            if ended
+            else f"noted; the agent sees it at its next decision ({run / 'notes.txt'})"
+        )
+    elif a.cmd == "report":
+        from .report import summarize_missing, write_report
+
+        run = _run_folder(a.run)
+        if a.summarize:
+            from .summaries import Summarizer
+
+            failed = summarize_missing(run, Summarizer(model=a.model), log=print)
+            if failed:
+                print(
+                    f"{failed} summaries failed (see decisions/NNN/summary-error.json); their rows show the rationale"
+                )
+        print(write_report(run))
+    elif a.cmd == "doctor":
+        return doctor(a)
+    elif a.cmd == "show":
+        from .wrapper import final_program
+
+        run = _run_folder(a.run)
+        _show(run)
+        print(f"final program: {final_program(run)}")
+        print("rerun it with no agent: opal.replay(<that folder>, optimizer, model, total_steps=...)")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

opal/context.py

@@ -0,0 +1,111 @@
+"""What the agent is told about the run, built from the live objects so the user only writes a goal."""
+
+from __future__ import annotations
+
+import inspect
+import re
+import sys
+from pathlib import Path
+
+import torch
+
+_PREFIXES = ("module.", "_orig_mod.", "_fsdp_wrapped_module.")
+
+
+def clean_name(name: str) -> str:
+    """Strip wrapper prefixes (DDP's module., torch.compile's _orig_mod.) so names match the bare model."""
+    changed = True
+    while changed:
+        changed = False
+        for p in _PREFIXES:
+            if name.startswith(p):
+                name, changed = name[len(p) :], True
+    return name.replace("._orig_mod.", ".")
+
+
+def named_params(model, optimizers):
+    """(name, parameter) for every parameter the user's optimizers update, in the model's order."""
+    in_opt = {id(p) for o in optimizers for g in o.param_groups for p in g["params"]}
+    out, seen = [], set()
+    for name, p in model.named_parameters():
+        if id(p) in in_opt and id(p) not in seen:
+            out.append((clean_name(name), p))
+            seen.add(id(p))
+    extra = [p for o in optimizers for g in o.param_groups for p in g["params"] if id(p) not in seen]
+    for i, p in enumerate(extra):
+        if id(p) not in seen:
+            out.append((f"extra.{i}", p))
+            seen.add(id(p))
+    return out
+
+
+def param_table(named, names_by_group=None) -> str:
+    lines = []
+    for n, p in named:
+        g = f"  group {names_by_group[n]}" if names_by_group and n in names_by_group else ""
+        lines.append(f"  {n} {tuple(p.shape)} {str(p.dtype).replace('torch.', '')}{g}")
+    return "\n".join(lines)
+
+
+def _wrappers(model):
+    """The model and every wrapper inside it, outermost first: DDP and Accelerate (.module), torch.compile
+    (._orig_mod), in any number and order (Accelerate on top of a script's own DDP gives two)."""
+    chain, m = [model], model
+    for _ in range(16):
+        inner = getattr(m, "_orig_mod", None)
+        if not isinstance(inner, torch.nn.Module):
+            inner = getattr(m, "module", None)
+        if not isinstance(inner, torch.nn.Module) or inner is m:
+            break
+        chain.append(inner)
+        m = inner
+    return chain
+
+
+def bare_model(model):
+    return _wrappers(model)[-1]
+
+
+def model_summary(model, named, world_size, device) -> str:
+    cls = type(bare_model(model))
+    n = sum(p.numel() for _, p in named)
+    compiled = any(hasattr(m, "_orig_mod") for m in _wrappers(model))
+    gpu = torch.cuda.get_device_name(device) if device.type == "cuda" else str(device)
+    return (
+        f"Model class {cls.__module__}.{cls.__qualname__}; {n:,} parameters updated by the optimizer in {len(named)} "
+        f"tensors; data-parallel world size {world_size}; device {gpu}; torch.compile on the model: "
+        f"{'yes' if compiled else 'no'}; torch {torch.__version__}."
+    )
+
+
+def group_index(optimizers, named):
+    by_id = {id(p): n for n, p in named}
+    out = {}
+    for i, o in enumerate(optimizers):
+        for j, g in enumerate(o.param_groups):
+            for p in g["params"]:
+                if id(p) in by_id:
+                    out[by_id[id(p)]] = f"{i}.{j}" if len(optimizers) > 1 else str(j)
+    return out
+
+
+def reference_files(model, optimizers, schedulers, max_bytes=300_000) -> dict:
+    """Read-only copies for the agent: the training script and the source of user-defined classes."""
+    out = {}
+    main = sys.modules.get("__main__")
+    script = getattr(main, "__file__", None)
+    if script and Path(script).is_file() and Path(script).stat().st_size <= max_bytes:
+        out["training_script_" + Path(script).name] = Path(script).read_text(errors="replace")
+    objs = [getattr(getattr(model, "module", model), "_orig_mod", getattr(model, "module", model))]
+    objs += list(optimizers) + list(schedulers)
+    for obj in objs:
+        mod = type(obj).__module__ or ""
+        if mod == "torch" or mod.startswith("torch.") or mod == "__main__":
+            continue
+        try:
+            f = inspect.getsourcefile(type(obj))
+        except TypeError:
+            f = None
+        if f and Path(f).is_file() and Path(f).stat().st_size <= max_bytes:
+            out["source_" + re.sub(r"[^A-Za-z0-9_.]", "_", mod) + ".py"] = Path(f).read_text(errors="replace")
+    return out

opal/defaults.json

@@ -0,0 +1,11 @@
+{
+ "agent.py": "8851488afd1424a8d7810db9e3c03f8a44bb0cda6c385e1ef79877eab2cad33b",
+ "context.py": "8c90b3631f16edb066e7cb4cba67440011f9dd74f10d81ebf21ca19ab9932508",
+ "misc/aggressive.py": "f75269eb77ba534f3bb0d0d856c10ba95297ab3cf9c65a882f51d5aaaa3d6332",
+ "misc/neutral_analysis.py": "1caa3720dd424c34d6d3076551f6c6f24882e5e6724266cf4fac84b870db2cd2",
+ "misc/setup.py": "e52e4f391b40be19c56c76abe6db1e05dbbe494dbc0d3e249a47e1220ee570c6",
+ "misc/taskfacts.py": "9ee5148cc7e441bee6a08f01e49e07c0fb8dc67e42843e1813f846c50fcd15dc",
+ "programs.py": "9963a97acb49e48aea76e82a63530bea6a64bc2cab234475a4b970984677e130",
+ "prompts/example_updater.py": "e7548f935ce808c625bfdceee324173571c3ce2272eab66e490a87394da5911f",
+ "prompts/neutral.md": "87e7b83756acffc87c25e5b24fae32b203c3c1cdf9d2452b6716c1cf943ee57d"
+}