modal_run.py

#!/usr/bin/env python3
"""Run swe_bench_run.py on Modal — beefy remote Linux, your API keys, your harness.

Setup (one-time):
  pip install modal swebench
  modal setup
  modal secret create openai-keys OPENAI_API_KEY=sk-...
  modal secret create anthropic-keys ANTHROPIC_API_KEY=sk-ant-...
    (required even if you only ever pass --provider openai -- both secrets are
    attached to run_instance/run_queue at decoration time, so both must exist)
  modal secret create modal-token MODAL_TOKEN_ID=<id> MODAL_TOKEN_SECRET=<secret>
    (values from ~/.modal.toml under your active profile)

Run one instance:
  modal run modal_run.py::run_one --instance-id django__django-12345

Switch the agent's LLM provider (default openai; IMPLEMENT/REPRODUCE use
claude-sonnet-4-6, everything else uses claude-haiku-4-5 -- see cloud_agent/config.py):
  modal run modal_run.py::run_one --instance-id django__django-12345 --provider anthropic

Run many in parallel (reads instance IDs from a file, one per line):
  modal run modal_run.py::run_batch --ids-file ids.txt

Results are printed to stdout; set RESULTS_DIR below to also write JSON files locally.

Both the agent loop AND the official swebench eval run inside `run_instance` --
one Modal function call per instance, nothing decided back on your machine.
`--modal true` (in swe_bench_run.evaluate_patch) makes the eval itself spin up
a *second*, nested Modal client session from inside this container to build
and run the per-instance Docker image. That needs `~/.modal.toml` on disk
(swebench's own credential check doesn't accept env vars), which
`_write_modal_token_file()` materializes at the top of `run_instance` from the
`modal-token` secret above -- and a writable cwd, which is why eval output
goes to /tmp same as DB_PATH/TRACES_DIR below.
"""

from __future__ import annotations

import json
import queue
import sys
import tempfile
import threading
from pathlib import Path

import modal
from modal_proto import api_pb2

# ── Image: Ubuntu + git + your harness deps ───────────────────────────────────

HARNESS_DIR = Path(__file__).parent  # local repo root

image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install(["git", "gcc", "python3-dev"])
    .pip_install([
        "openai>=1.0",
        "anthropic>=0.40",
        "pydantic>=2.0",
        "python-dotenv>=1.0",
        "datasets",
        "pytest",
        # Pinned to match the local client (`modal --version`) so the in-image
        # SDK doesn't silently drift to whatever's newest on PyPI at build time.
        # This alone did NOT fix the nested-eval crash below -- see the
        # patch_swebench_modal_sandbox_python.py run_commands call, the real fix.
        "modal==1.5.1",
        "swebench",
    ])
    .run_commands(
        "git config --global user.email 'agent@swe-bench'",
        "git config --global user.name 'SWE-bench Agent'",
    )
    .add_local_dir(
        str(HARNESS_DIR),
        remote_path="/app",
        ignore=[".git", "__pycache__", ".pytest_cache", "traces", "build", "*.db"],
        copy=True,  # required to run a build step (below) after add_local_dir
    )
    # swebench's own REPO_BASE_COMMIT_BRANCH table pins a git branch per (repo,
    # base_commit) that no longer exists for every instance (upstream repos delete/
    # rename release branches over time -- confirmed for sympy/sympy's 1.7 branch,
    # see setup_scripts/patch_swebench_stale_branches.py) -- without this, the
    # official eval's own Docker build fails on `git clone --branch <stale>` before
    # our agent's patch is ever considered, regardless of whether the fix is right.
    .run_commands("python3 /app/setup_scripts/patch_swebench_stale_branches.py")
    # swebench's own nested-eval Sandbox (swebench.harness.modal_eval.run_evaluation_modal)
    # provisions its orchestration Python via `modal.Image.from_registry("ubuntu:22.04",
    # add_python="3.11")` -- that specific add_python build's `typing` module is broken
    # (missing `Iterable`), crashing the Sandbox entrypoint's `import asyncio` before the
    # real per-instance eval (which runs under its own separately-installed miniconda env,
    # untouched by this) ever gets to run -- confirmed on django__django-11099, see
    # report-7-8.txt. Nothing on our side (agent code, our own outer image's Python 3.11)
    # is involved -- see setup_scripts/patch_swebench_modal_sandbox_python.py.
    .run_commands("python3 /app/setup_scripts/patch_swebench_modal_sandbox_python.py")
)

app = modal.App("swe-bench-harness")

# Per-(repo,version) micromamba env cache built by setup_scripts/build_env_cache.py
# (see that script and agent/env_cache.py) -- gives repos like scikit-learn 0.20
# a real Python 3.6 with real cp36m wheels for numpy==1.19.2, which have zero
# wheels for this image's own Python 3.11 (confirmed: hard cythonize failure).
# Mount path must match agent.env_cache.CACHE_VOLUME_PATH exactly. Must be a v2
# Volume (see build_env_cache.py's comment) -- v1 caps out at 500,000 inodes
# total and a handful of conda envs alone can exceed that.
ENV_CACHE_VOLUME_NAME = "swebench-envs-v2"
ENV_CACHE_MOUNT_PATH = "/vol/envs"
env_cache_volume = modal.Volume.from_name(
    ENV_CACHE_VOLUME_NAME, create_if_missing=True, version=api_pb2.VolumeFsVersion.VOLUME_FS_VERSION_V2,
)

# ── Core function ─────────────────────────────────────────────────────────────

def _write_modal_token_file() -> None:
    """Materialize ~/.modal.toml in this container from the modal-token secret.

    The nested `swebench.harness.run_evaluation --modal true` call below is
    itself a Modal client session -- its own credential check
    (validate_modal_credentials) only looks for this file on disk, not
    MODAL_TOKEN_ID/MODAL_TOKEN_SECRET env vars, even though the underlying
    client accepts both.
    """
    import os

    token_id = os.environ["MODAL_TOKEN_ID"]
    token_secret = os.environ["MODAL_TOKEN_SECRET"]
    toml_path = Path.home() / ".modal.toml"
    toml_path.write_text(
        f'[nested-eval]\ntoken_id = "{token_id}"\ntoken_secret = "{token_secret}"\nactive = true\n'
    )


@app.function(
    image=image,
    secrets=[
        modal.Secret.from_name("openai-keys"),
        modal.Secret.from_name("anthropic-keys"),
        modal.Secret.from_name("modal-token"),
    ],
    volumes={ENV_CACHE_MOUNT_PATH: env_cache_volume},
    cpu=4,
    memory=8192,  # 8 GB — plenty for git clone + pytest; bump to 16384 if needed
    timeout=60 * 60 * 2,  # 1 hour for the agent + up to 1 hour for the nested eval build/run
)
def run_instance(
    instance_id: str | None = None,
    index: int | None = None,
    instance_json: str | None = None,  # pre-serialized instance dict
    max_turns: int = 130,
    max_cost: float = 50.0,
    dataset: str = "princeton-nlp/SWE-bench_Lite",
    split: str = "test",
    provider: str = "openai",  # "openai" | "anthropic" -- see cloud_agent.config.Settings
) -> dict:
    """Runs one SWE-bench instance end-to-end and returns a result dict."""
    import os
    import sys
    from pathlib import Path

    sys.path.insert(0, "/app")
    os.chdir("/app")

    # Modal mounts are read-only; redirect writable state to /tmp
    os.environ.setdefault("DB_PATH", "/tmp/cloud_agent.db")
    os.environ.setdefault("TRACES_DIR", "/tmp/traces")
    Path("/tmp/traces").mkdir(exist_ok=True)
    # Must be set before importing swe_bench_run below -- cloud_agent.config's
    # module-level `settings` singleton reads LLM_PROVIDER (and picks its model
    # defaults) at import time, the first time anything imports that module.
    os.environ["LLM_PROVIDER"] = provider
    _write_modal_token_file()

    from swe_bench_run import (
        build_model_patch,
        evaluate_patch,
        load_instance,
        run_agent,
        setup_workspace,
        summarize_test_results,
    )

    # ── Load instance ──────────────────────────────────────────────────────
    if instance_json:
        instance = json.loads(instance_json)
    else:
        instance = load_instance(
            instance_id=instance_id,
            index=index,
            dataset=dataset,
            split=split,
        )

    iid = instance["instance_id"]
    print(f"\n=== {iid} ===")
    print(f"repo: {instance['repo']}  @{instance['base_commit'][:8]}")

    with tempfile.TemporaryDirectory() as work_dir:
        # ── Setup ──────────────────────────────────────────────────────────
        try:
            setup_workspace(instance, work_dir)
            # Capture HEAD now, before run_agent() -- CHECKPOINT commits the
            # agent's changes at the end of the run, so this must happen
            # before that (see build_model_patch's docstring).
            import subprocess
            base_commit = subprocess.run(
                ["git", "rev-parse", "HEAD"], cwd=work_dir, capture_output=True, text=True,
            ).stdout.strip()
        except Exception as exc:
            return {"instance_id": iid, "outcome": "setup_error", "error": str(exc)}

        # ── Agent ──────────────────────────────────────────────────────────
        try:
            final = run_agent(instance, work_dir, max_turns, max_cost)
            agent_meta = {
                "status": final.task_status,
                "turns": final.budgets.used_llm_turns,
                "cost_usd": round(final.budgets.used_cost_usd, 4),
            }
            trace_path = Path("/tmp/traces") / f"{final.session_id}.jsonl"
            trace_jsonl = trace_path.read_text() if trace_path.exists() else ""
        except Exception as exc:
            return {"instance_id": iid, "outcome": "agent_error", "error": str(exc)}

        # ── Diff ───────────────────────────────────────────────────────────
        try:
            model_patch = build_model_patch(work_dir, base_commit)
        except Exception as exc:
            return {
                "instance_id": iid,
                "outcome": "diff_error",
                "error": str(exc),
                "trace_jsonl": trace_jsonl,
                **agent_meta,
            }

        # ── Eval ───────────────────────────────────────────────────────────
        # Nested Modal client session (--modal true): builds/runs the real
        # per-instance Docker image via the official swebench harness, right
        # here in this same container. Output goes to /tmp (this container's
        # own cwd, /app, is a read-only mount).
        try:
            passed, report = evaluate_patch(
                iid, model_patch, dataset=dataset, split=split, cwd="/tmp",
            )
        except Exception as exc:
            return {
                "instance_id": iid,
                "outcome": "eval_error",
                "error": str(exc),
                "trace_jsonl": trace_jsonl,
                **agent_meta,
            }

    outcome = "pass" if passed else "fail"
    print(f"{outcome.upper()}  {iid}  cost=${agent_meta['cost_usd']}  turns={agent_meta['turns']}")

    return {
        "instance_id": iid,
        "outcome": outcome,
        "test_results": summarize_test_results(iid, report),
        "trace_jsonl": trace_jsonl,
        "report": report,  # full swebench receipt -- would otherwise die with /tmp
        **agent_meta,
    }


@app.function(
    image=image,
    secrets=[
        modal.Secret.from_name("openai-keys"),
        modal.Secret.from_name("anthropic-keys"),
        modal.Secret.from_name("modal-token"),
    ],
    volumes={ENV_CACHE_MOUNT_PATH: env_cache_volume},
    cpu=4,
    memory=8192,
    timeout=60 * 60 * 6,  # a queue processes several instances sequentially
)
def run_queue(ids: list[str], **kwargs):
    """Run a list of instance IDs one at a time in a single container.

    Used to cap OpenAI request/token concurrency: N queues running in parallel via
    .map() means at most N instances are in flight at once, regardless of --limit.
    Its own volumes={} mount above is what actually matters for the env cache --
    run_instance.local() below runs in *this* container, so it inherits this
    decorator's mounts, not run_instance's own (those only apply when
    run_instance is invoked directly via .remote(), e.g. from run_one).

    Yields each instance's result as it finishes (rather than returning a list at
    the end) so .map() streams completed instances -- trace and all -- back to the
    caller immediately. If this container is later killed or times out mid-queue,
    everything already yielded has already reached durable disk; only the
    in-flight instance is lost, not the whole queue.
    """
    for iid in ids:
        yield run_instance.local(instance_id=iid, **kwargs)


def _chunk(ids: list[str], concurrency: int) -> list[list[str]]:
    """Split ids into up to `concurrency` round-robin groups, each run sequentially."""
    groups = [ids[i::concurrency] for i in range(concurrency)]
    return [g for g in groups if g]


def _save_trace(result: dict, traces_dir: Path) -> None:
    """Pop the per-session trace_jsonl/report blobs off a result dict and write them locally."""
    trace_jsonl = result.pop("trace_jsonl", None)
    report = result.pop("report", None)
    if not trace_jsonl and not report:
        return
    traces_dir.mkdir(parents=True, exist_ok=True)
    if trace_jsonl:
        (traces_dir / f"{result['instance_id']}_trace.txt").write_text(trace_jsonl)
    if report:
        report_path = traces_dir / f"{result['instance_id']}_report.json"
        report_path.write_text(json.dumps(report, indent=2))


def _run_queues_to_file(
    ids: list[str], concurrency: int, static_kwargs: dict, results_file: str,
) -> None:
    chunks = _chunk(ids, concurrency)
    print(f"Submitting {len(ids)} instances across {len(chunks)} queue(s) "
          f"(max {concurrency} concurrent)...")

    out = Path(results_file)
    out.parent.mkdir(parents=True, exist_ok=True)
    passed = failed = errors = 0

    # run_queue is a generator, and Modal's .map() refuses generator functions
    # outright ("cannot be called with .map()") -- so each chunk's remote_gen()
    # is drained on its own thread, funneled through this queue. Every instance
    # reaches `stream` (and disk) the moment it finishes, regardless of which
    # chunk it came from or whether that chunk's container later dies.
    stream: queue.Queue = queue.Queue()

    def _drain(chunk: list[str]) -> None:
        try:
            for result in run_queue.remote_gen(chunk, **static_kwargs):
                stream.put(result)
        except Exception as exc:
            stream.put(exc)
        finally:
            stream.put(None)  # sentinel: this chunk is done (ok or not)

    threads = [threading.Thread(target=_drain, args=(c,), daemon=True) for c in chunks]
    for t in threads:
        t.start()

    with out.open("w") as f:
        pending = len(threads)
        while pending:
            result = stream.get()
            if result is None:
                pending -= 1
                continue
            if isinstance(result, BaseException):
                errors += 1
                record = {"outcome": "map_error", "error": str(result)}
                f.write(json.dumps(record) + "\n")
                f.flush()
                print(f"  [{passed+failed+errors}/{len(ids)}] QUEUE ERROR: {result}")
                continue
            _save_trace(result, out.parent)
            f.write(json.dumps(result) + "\n")
            f.flush()
            o = result.get("outcome", "?")
            if o == "pass":
                passed += 1
            elif o == "fail":
                failed += 1
            else:
                errors += 1
            print(f"  [{passed+failed+errors}/{len(ids)}] {result['instance_id']} → {o}")

    print(f"\nDone. pass={passed} fail={failed} errors={errors}")
    print(f"Results written to {out}")


# ── Local entrypoints ─────────────────────────────────────────────────────────

@app.local_entrypoint()
def run_one(
    instance_id: str = "",
    index: int = -1,
    max_turns: int = 130,
    max_cost: float = 50.0,
    dataset: str = "princeton-nlp/SWE-bench_Lite",
    split: str = "test",
    provider: str = "openai",  # "openai" | "anthropic"
) -> None:
    """modal run modal_run.py::run_one --instance-id django__django-12345 [--provider anthropic]"""
    kwargs: dict = {
        "max_turns": max_turns, "max_cost": max_cost, "dataset": dataset, "split": split,
        "provider": provider,
    }
    if instance_id:
        kwargs["instance_id"] = instance_id
    elif index >= 0:
        kwargs["index"] = index
    else:
        print("ERROR: pass --instance-id or --index", file=sys.stderr)
        sys.exit(1)

    result = run_instance.remote(**kwargs)
    _save_trace(result, Path("test_traces"))
    print(json.dumps(result))


@app.local_entrypoint()
def run_repo(
    repo: str = "sympy/sympy",
    limit: int = 10,
    concurrency: int = 10,
    max_turns: int = 130,
    max_cost: float = 50.0,
    dataset: str = "princeton-nlp/SWE-bench_Lite",
    split: str = "test",
    results_file: str = "test_traces/results.jsonl",
    provider: str = "openai",  # "openai" | "anthropic"
) -> None:
    """Run the first --limit instances from a specific repo on Modal.

    --concurrency caps how many instances run at once (default 10, matched to
    OpenAI Tier 2 headroom) by splitting --limit instances into that many queues,
    each processed sequentially in its own container.

    modal run modal_run.py::run_repo --repo sympy/sympy --limit 50 --concurrency 10
    """
    from datasets import load_dataset

    ds = load_dataset(dataset, split=split)
    ids = [row["instance_id"] for row in ds if row["repo"] == repo][:limit]
    if not ids:
        print(f"ERROR: no instances found for repo {repo!r}", file=sys.stderr)
        sys.exit(1)

    print(f"Found {len(ids)} instances for {repo}: {ids}")

    static_kwargs = {
        "max_turns": max_turns, "max_cost": max_cost, "dataset": dataset, "split": split,
        "provider": provider,
    }
    _run_queues_to_file(ids, concurrency, static_kwargs, results_file)


@app.local_entrypoint()
def run_batch(
    ids_file: str = "ids.txt",
    concurrency: int = 10,
    max_turns: int = 130,
    max_cost: float = 50.0,
    dataset: str = "princeton-nlp/SWE-bench_Lite",
    split: str = "test",
    results_file: str = "test_traces/results.jsonl",
    provider: str = "openai",  # "openai" | "anthropic"
) -> None:
    """Run all instance IDs in ids_file on Modal, --concurrency at a time.

    ids.txt format: one instance_id per line, blank lines / # comments ignored.

    modal run modal_run.py::run_batch --ids-file ids.txt --results-file results.jsonl
    """
    ids = [
        line.strip()
        for line in Path(ids_file).read_text().splitlines()
        if line.strip() and not line.startswith("#")
    ]
    if not ids:
        print("ERROR: ids_file is empty", file=sys.stderr)
        sys.exit(1)

    static_kwargs = {
        "max_turns": max_turns, "max_cost": max_cost, "dataset": dataset, "split": split,
        "provider": provider,
    }
    _run_queues_to_file(ids, concurrency, static_kwargs, results_file)