tenant.openagents/omega

1from __future__ import annotations
2
3import hashlib
4import json
5import os
6import pathlib
7import shlex
8import shutil
9import subprocess
10import tarfile
11import time
12import traceback
13import uuid
14from typing import Any
15
16import modal
17
18from . import cleanup, config, harness_command, rejudge, source
19from .common import load_json, safe_extract_archive, utc_now, write_json
20from .source import (
21    BUILD_TARGET,
22    CARGO_ZIGBUILD_VERSION,
23    RUST_IMAGE,
24    RUST_VERSION,
25    ZIG_VERSION,
26    build_toolchain_info,
27)
28
29APP_NAME = os.environ.get("AGENT_EVALS_APP_NAME", "agent-evals")
30VOLUME_NAME = os.environ.get("AGENT_EVALS_VOLUME", "agent-evals")
31MODAL_TOKEN_SECRET_NAME = os.environ.get(
32    "AGENT_EVALS_MODAL_TOKEN_SECRET", "agent-evals-modal-token"
33)
34# LLM-providers secret mounted into the rejudge controller so the in-controller
35# judge proxy can read the judge's API key (e.g. BASETEN_API_KEY) the same way a
36# trial sandbox does. Override at deploy time if the secret is named differently.
37LLM_PROVIDERS_SECRET_NAME = os.environ.get(
38    "AGENT_EVALS_LLM_PROVIDERS_SECRET", config.DEFAULT_LLM_PROVIDERS_SECRET_NAME
39)
40REPO_URL = os.environ.get(
41    "AGENT_EVALS_REPO_URL", "https://github.com/zed-industries/zed.git"
42)
43ZIG_URL = (
44    f"https://ziglang.org/download/{ZIG_VERSION}/zig-x86_64-linux-{ZIG_VERSION}.tar.xz"
45)
46CLEANUP_BUILD_RETENTION_DAYS = float(
47    os.environ.get("AGENT_EVALS_BUILD_RETENTION_DAYS", "14")
48)
49
50app = modal.App(APP_NAME)
51volume = modal.Volume.from_name(VOLUME_NAME, create_if_missing=True)
52
53
54def reload_volume() -> None:
55    try:
56        volume.reload()
57    except Exception:
58        pass
59
60
61build_image = (
62    modal.Image.from_registry(RUST_IMAGE, add_python="3.13")
63    .apt_install("cmake", "build-essential", "curl", "xz-utils", "git")
64    .run_commands(
65        f"rustup toolchain install {RUST_VERSION} --profile minimal"
66        " --component rustfmt --component clippy"
67        " --component rust-analyzer --component rust-src"
68        " --target wasm32-wasip2 --target wasm32-unknown-unknown"
69        " --target x86_64-unknown-linux-musl --target x86_64-unknown-linux-gnu",
70        f"mkdir -p /opt/zig && curl -fsSL {ZIG_URL}"
71        " | tar -xJ -C /opt/zig --strip-components=1"
72        " && ln -s /opt/zig/zig /usr/local/bin/zig",
73        f"cargo install --locked cargo-zigbuild --version {CARGO_ZIGBUILD_VERSION}",
74    )
75)
76
77controller_image = (
78    modal.Image.debian_slim(python_version="3.13")
79    .apt_install("bash", "ca-certificates", "curl", "git", "tar")
80    .pip_install(f"modal=={config.MODAL_VERSION}", "uv")
81    .run_commands(
82        f"uv tool install harbor=={config.HARBOR_VERSION}"
83        f" --with 'modal=={config.MODAL_VERSION}' && "
84        "ln -sf /root/.local/bin/harbor /usr/local/bin/harbor",
85        # Pier is a Harbor fork that adds per-agent network allowlists, required
86        # to run DeepSWE's air-gapped (`allow_internet = false`) tasks while
87        # still letting eval-cli reach the model API. Kept non-fatal so a Pier
88        # install problem can't break the Harbor benchmarks that share this image;
89        # DeepSWE runs surface a clear "pier not found" error at runtime instead.
90        "(uv tool install git+https://github.com/datacurve-ai/pier"
91        f" --with 'modal=={config.MODAL_VERSION}' && "
92        "ln -sf /root/.local/bin/pier /usr/local/bin/pier) || "
93        "echo 'WARNING: pier install failed; DeepSWE runs will not work'",
94    )
95)
96
97
98def provision_benchmark_dataset(run_request: dict[str, Any], log: Any) -> None:
99    """Clone a path-backed benchmark dataset (SWE-Atlas tw, DeepSWE) into the
100    location the harness command expects. Registry datasets need no provisioning;
101    the harness pulls them from its hub."""
102    benchmark = run_request["benchmark"]
103    dataset = benchmark.get("dataset") or {}
104    kind = dataset.get("kind")
105    if kind not in ("path", "pier_path"):
106        return
107
108    repo_url = dataset.get("repo_url")
109    repo_ref = dataset.get("repo_ref") or "main"
110    if not repo_url:
111        raise ValueError(f"benchmark {benchmark['id']} path dataset requires repo_url")
112
113    clone_dir = pathlib.Path(harness_command.dataset_clone_dir(benchmark))
114    if clone_dir.exists():
115        shutil.rmtree(clone_dir)
116    clone_dir.mkdir(parents=True, exist_ok=True)
117    log(f"Fetching {benchmark['id']} dataset {repo_url}@{repo_ref}")
118    subprocess.run(["git", "init", "-q", str(clone_dir)], check=True)
119    subprocess.run(
120        ["git", "fetch", "--depth", "1", repo_url, repo_ref],
121        cwd=clone_dir,
122        check=True,
123    )
124    subprocess.run(["git", "checkout", "-q", "FETCH_HEAD"], cwd=clone_dir, check=True)
125    data_dir = pathlib.Path(harness_command.dataset_path(benchmark))
126    if not data_dir.exists():
127        raise FileNotFoundError(f"benchmark dataset directory not found: {data_dir}")
128
129
130def output_of(command: list[str]) -> str:
131    try:
132        return subprocess.run(
133            command,
134            check=True,
135            capture_output=True,
136            text=True,
137        ).stdout.strip()
138    except subprocess.CalledProcessError as error:
139        return f"unknown ({error})"
140
141
142@app.function(
143    image=build_image,
144    # Right-sized from the original 16 cpu / 32 GB after observing peak usage of
145    # ~10 cores and ~10 GB. ephemeral_disk stays at Modal's 512 GiB floor (the
146    # minimum allowed for a function), so it isn't a tunable here.
147    cpu=12,
148    memory=24576,
149    ephemeral_disk=524288,
150    timeout=7200,
151    volumes={"/data": volume},
152)
153def build_eval_cli(build_request: dict[str, Any]) -> dict[str, Any]:
154    def run(command: str, **kwargs: Any) -> None:
155        print(f"+ {command}", flush=True)
156        subprocess.run(command, shell=True, check=True, **kwargs)
157
158    reload_volume()
159
160    build_id = build_request["build_id"]
161    source_info = build_request.get("source") or {"type": "git_patch"}
162    base_sha = source_info.get("base_sha") or build_request["base_sha"]
163    patch = build_request.get("patch") or ""
164    build_dir = pathlib.Path("/data/builds") / build_id
165    ready_path = build_dir / "READY"
166    binary_path = build_dir / "eval-cli"
167    build_info_path = build_dir / "build-info.json"
168    # The lease lives outside build_dir so it never interferes with the atomic
169    # move of the finished build directory below.
170    building_path = pathlib.Path("/data/build-locks") / f"{build_id}.json"
171    lease_ttl = int(build_request.get("build_wait_timeout_secs") or 7200)
172    owner = uuid.uuid4().hex
173
174    def reuse_existing() -> dict[str, Any]:
175        build_info = load_json(build_info_path) or {}
176        build_info.setdefault("build_id", build_id)
177        build_info["reused"] = True
178        volume.commit()
179        print(f"Reusing existing build {build_id}", flush=True)
180        return build_info
181
182    if ready_path.exists() and binary_path.exists():
183        return reuse_existing()
184    if build_dir.exists():
185        raise RuntimeError(
186            f"build directory already exists but is not ready; refusing to overwrite {build_id}"
187        )
188
189    # Single-flight lease: if another invocation is already compiling this exact
190    # build, wait for it to finish rather than running a second multi-minute
191    # compile. The lease is best-effort (the volume has no atomic compare-and-swap);
192    # the atomic move below still guarantees correctness if two builds slip through.
193    deadline = time.time() + lease_ttl
194    while True:
195        reload_volume()
196        if ready_path.exists() and binary_path.exists():
197            return reuse_existing()
198        if build_dir.exists():
199            raise RuntimeError(
200                f"build directory already exists but is not ready; refusing to overwrite {build_id}"
201            )
202        lease = load_json(building_path)
203        now = time.time()
204        held_by_other = (
205            isinstance(lease, dict)
206            and lease.get("owner") != owner
207            and (now - float(lease.get("epoch") or 0)) < lease_ttl
208        )
209        if held_by_other and time.time() < deadline:
210            print(
211                f"Build {build_id} is being compiled elsewhere; waiting",
212                flush=True,
213            )
214            time.sleep(15)
215            continue
216        break
217    # Claim the lease (write_json creates /data/build-locks). build_dir is
218    # intentionally NOT created here so the atomic move below still works.
219    write_json(building_path, {"owner": owner, "epoch": time.time(), "at": utc_now()})
220    volume.commit()
221
222    rustc_version = output_of(["rustc", "--version"])
223    zig_version = output_of(["zig", "version"])
224    cargo_zigbuild_version = output_of(["cargo-zigbuild", "--version"])
225
226    workdir = pathlib.Path("/build/zed")
227    if workdir.exists():
228        shutil.rmtree(workdir)
229    workdir.mkdir(parents=True, exist_ok=True)
230    os.chdir(workdir)
231
232    repo_url = source_info.get("repo_url") or REPO_URL
233    run("git init -q .")
234    run(f"git fetch --depth 1 {shlex.quote(repo_url)} {shlex.quote(base_sha)}")
235    run("git checkout -q FETCH_HEAD")
236
237    if patch.strip():
238        patch_file = pathlib.Path("/build/source.patch")
239        patch_file.write_text(patch)
240        run(f"git apply --stat {shlex.quote(str(patch_file))}")
241        run(f"git apply {shlex.quote(str(patch_file))}")
242
243    patch_sha256 = source_info.get("patch_sha256")
244
245    run(
246        "cargo zigbuild --release --package eval_cli --target x86_64-unknown-linux-musl"
247    )
248
249    built = workdir / "target/x86_64-unknown-linux-musl/release/eval-cli"
250    run(f"strip {built}")
251    binary_bytes = built.read_bytes()
252    binary_sha256 = hashlib.sha256(binary_bytes).hexdigest()
253
254    temporary_dir = pathlib.Path("/data/tmp/builds") / f"{build_id}-{uuid.uuid4().hex}"
255    temporary_dir.mkdir(parents=True, exist_ok=True)
256    shutil.copy(built, temporary_dir / "eval-cli")
257    (temporary_dir / "eval-cli").chmod(0o755)
258    if patch.strip():
259        (temporary_dir / "source.patch").write_text(patch)
260    write_json(temporary_dir / "source-info.json", source_info)
261
262    build_info = {
263        "build_id": build_id,
264        "base_sha": base_sha,
265        "patch_sha256": patch_sha256,
266        "rustc_version": rustc_version,
267        "zig_version": zig_version,
268        "cargo_zigbuild_version": cargo_zigbuild_version,
269        "built_at_utc": utc_now(),
270        "binary_size_bytes": len(binary_bytes),
271        "binary_sha256": binary_sha256,
272        "target": BUILD_TARGET,
273        "rust_image": RUST_IMAGE,
274        "toolchain": build_toolchain_info(),
275        "source": source_info,
276    }
277    write_json(temporary_dir / "build-info.json", build_info)
278    (temporary_dir / "READY").write_text(utc_now() + "\n")
279
280    if not build_dir.exists():
281        build_dir.parent.mkdir(parents=True, exist_ok=True)
282        shutil.move(str(temporary_dir), str(build_dir))
283    elif ready_path.exists() and binary_path.exists():
284        shutil.rmtree(temporary_dir, ignore_errors=True)
285        build_info = load_json(build_info_path) or build_info
286        build_info["reused_after_race"] = True
287    else:
288        shutil.rmtree(temporary_dir, ignore_errors=True)
289        raise RuntimeError(
290            f"build directory already exists but is not ready; refusing to overwrite {build_id}"
291        )
292
293    # Release the single-flight lease now that READY exists.
294    try:
295        building_path.unlink()
296    except OSError:
297        pass
298    volume.commit()
299    print(f"Committed build {build_id} to volume '{VOLUME_NAME}'", flush=True)
300    return build_info
301
302
303@app.function(
304    image=controller_image,
305    cpu=1,
306    memory=512,
307    timeout=300,
308    volumes={"/data": volume},
309)
310def list_builds(limit: int = 50) -> list[dict[str, Any]]:
311    reload_volume()
312    builds_dir = pathlib.Path("/data/builds")
313    rows = []
314    for build_dir in builds_dir.iterdir() if builds_dir.exists() else []:
315        if not build_dir.is_dir():
316            continue
317        build_info = load_json(build_dir / "build-info.json") or {}
318        ready = (build_dir / "READY").exists() and (build_dir / "eval-cli").exists()
319        rows.append(
320            {
321                "build_id": build_dir.name,
322                "ready": ready,
323                "base_sha": build_info.get("base_sha"),
324                "patch_sha256": build_info.get("patch_sha256"),
325                "built_at_utc": build_info.get("built_at_utc"),
326                "binary_sha256": build_info.get("binary_sha256"),
327                "source": build_info.get("source"),
328            }
329        )
330    rows.sort(key=lambda row: row.get("built_at_utc") or "", reverse=True)
331    return rows[:limit]
332
333
334def run_row(
335    namespace_dir: pathlib.Path, experiment_dir: pathlib.Path, run_dir: pathlib.Path
336) -> dict[str, Any]:
337    state = load_json(run_dir / "state.json") or {}
338    request = load_json(run_dir / "request.json") or {}
339    metadata = load_json(run_dir / "run-metadata.json") or {}
340    summary = load_json(run_dir / "summary.json") or {}
341    benchmark = request.get("benchmark") or metadata.get("benchmark") or {}
342    benchmark_id = benchmark.get("id") if isinstance(benchmark, dict) else None
343    suite_id = request.get("suite_id") or metadata.get("suite_id")
344    return {
345        "namespace": namespace_dir.name,
346        "experiment_name": experiment_dir.name,
347        "run_id": run_dir.name,
348        "status": state.get("status"),
349        "updated_at": state.get("updated_at"),
350        "created_at": metadata.get("created_at")
351        or state.get("created_at")
352        or request.get("created_at"),
353        "agent_model": metadata.get("agent_model"),
354        "judge_preset": metadata.get("judge_preset"),
355        "judge_model": metadata.get("judge_model"),
356        "build_id": metadata.get("build_id") or state.get("build_id"),
357        "trial_count": summary.get("trial_count"),
358        "has_archive": summary.get("has_archive"),
359        "suite_id": suite_id,
360        "part": request.get("suite_part"),
361        "benchmark": benchmark_id,
362    }
363
364
365def scan_runs(
366    *,
367    namespace: str | None = None,
368    experiment_name: str | None = None,
369    suite_id: str | None = None,
370) -> list[dict[str, Any]]:
371    runs_root = pathlib.Path("/data/runs")
372    rows = []
373    if not runs_root.exists():
374        return []
375    namespaces = (
376        [runs_root / namespace] if namespace else [path for path in runs_root.iterdir()]
377    )
378    for namespace_dir in namespaces:
379        if not namespace_dir.is_dir():
380            continue
381        experiments = (
382            [namespace_dir / experiment_name]
383            if experiment_name
384            else [path for path in namespace_dir.iterdir()]
385        )
386        for experiment_dir in experiments:
387            if not experiment_dir.is_dir():
388                continue
389            for run_dir in experiment_dir.iterdir():
390                if not run_dir.is_dir():
391                    continue
392                row = run_row(namespace_dir, experiment_dir, run_dir)
393                if suite_id and row.get("suite_id") != suite_id:
394                    continue
395                rows.append(row)
396    rows.sort(
397        key=lambda row: row.get("updated_at") or row.get("created_at") or "",
398        reverse=True,
399    )
400    return rows
401
402
403@app.function(
404    image=controller_image,
405    cpu=1,
406    memory=512,
407    timeout=300,
408    volumes={"/data": volume},
409)
410def list_runs(
411    namespace: str | None = None,
412    experiment_name: str | None = None,
413    limit: int = 50,
414) -> list[dict[str, Any]]:
415    reload_volume()
416    rows = scan_runs(namespace=namespace, experiment_name=experiment_name)
417    return rows[:limit]
418
419
420def _run_cleanup(
421    *,
422    dry_run: bool,
423    build_retention_days: float | None = None,
424) -> dict[str, Any]:
425    reload_volume()
426    result = cleanup.prune_artifacts(
427        pathlib.Path("/data"),
428        dry_run=dry_run,
429        build_retention_days=(
430            build_retention_days
431            if build_retention_days is not None
432            else CLEANUP_BUILD_RETENTION_DAYS
433        ),
434    )
435    if not dry_run:
436        volume.commit()
437    print(json.dumps(result, indent=2), flush=True)
438    return result
439
440
441@app.function(
442    image=controller_image,
443    cpu=1,
444    memory=2048,
445    timeout=1800,
446    volumes={"/data": volume},
447    schedule=modal.Period(days=1),
448)
449def cleanup_scheduled() -> dict[str, Any]:
450    """Daily prune of stale build artifacts. Never touches eval results."""
451    return _run_cleanup(dry_run=False)
452
453
454@app.function(
455    image=controller_image,
456    cpu=1,
457    memory=2048,
458    timeout=1800,
459    volumes={"/data": volume},
460)
461def cleanup_artifacts(request: dict[str, Any]) -> dict[str, Any]:
462    """On-demand prune (the `zed-eval cleanup` command), supporting --dry-run
463    and retention overrides."""
464    return _run_cleanup(
465        dry_run=bool(request.get("dry_run")),
466        build_retention_days=request.get("build_retention_days"),
467    )
468
469
470@app.function(
471    image=controller_image,
472    cpu=1,
473    memory=512,
474    timeout=300,
475    volumes={"/data": volume},
476)
477def read_run_provenance(
478    namespace: str, experiment_name: str, run_id: str
479) -> dict[str, Any]:
480    """Provenance a baseline record needs for one run: the launch request
481    (model/judge/resources/build_id), the run summary (status), and the build's
482    build-info (base_sha/patch_sha256/source label)."""
483    reload_volume()
484    run_dir = pathlib.Path("/data/runs") / namespace / experiment_name / run_id
485    request = load_json(run_dir / "request.json") or {}
486    summary = load_json(run_dir / "summary.json") or {}
487    metadata = load_json(run_dir / "run-metadata.json") or {}
488    build_id = request.get("build_id") or metadata.get("build_id")
489    build_info: dict[str, Any] = {}
490    if build_id:
491        build_info = (
492            load_json(pathlib.Path("/data/builds") / build_id / "build-info.json") or {}
493        )
494    return {
495        "request": request,
496        "summary": summary,
497        "metadata": metadata,
498        "build_info": build_info,
499    }
500
501
502@app.function(
503    image=controller_image,
504    cpu=1,
505    memory=512,
506    timeout=300,
507    volumes={"/data": volume},
508)
509def record_baseline(record: dict[str, Any]) -> dict[str, Any]:
510    """Write a baseline-of-record keyed by (experiment, model). Supersedes any
511    existing current baseline (archived under history/) and refreshes the
512    discoverable index."""
513    reload_volume()
514    experiment_slug = source.sanitize_namespace(record["experiment"])
515    model_slug = source.sanitize_namespace(record.get("model_slug") or record["model"])
516    base_dir = pathlib.Path("/data/baselines") / experiment_slug / model_slug
517    current_path = base_dir / "current.json"
518
519    existing = load_json(current_path)
520    if existing:
521        previous_sha = source.sanitize_namespace(
522            str(existing.get("base_sha") or "unknown")
523        )
524        write_json(base_dir / "history" / f"{previous_sha}.json", existing)
525
526    payload = {**record, "updated_at": utc_now()}
527    write_json(current_path, payload)
528
529    index_path = pathlib.Path("/data/baselines/index.json")
530    index = load_json(index_path) or {}
531    entries = index.get("baselines")
532    if not isinstance(entries, list):
533        entries = []
534    entries = [
535        entry
536        for entry in entries
537        if not (
538            entry.get("experiment") == record["experiment"]
539            and entry.get("model") == record["model"]
540        )
541    ]
542    entries.append(
543        {
544            "experiment": record["experiment"],
545            "model": record["model"],
546            "base_sha": record.get("base_sha"),
547            "base_ref": record.get("base_ref"),
548            "on_main": record.get("on_main"),
549            "clean": record.get("clean"),
550            "judge": record.get("judge"),
551            "run_id": (record.get("run") or {}).get("run_id"),
552            "recorded_at": record.get("recorded_at"),
553            "path": f"baselines/{experiment_slug}/{model_slug}/current.json",
554        }
555    )
556    entries.sort(
557        key=lambda entry: (entry.get("experiment") or "", entry.get("model") or "")
558    )
559    write_json(index_path, {"baselines": entries})
560    volume.commit()
561    return payload
562
563
564@app.function(
565    image=controller_image,
566    cpu=1,
567    memory=512,
568    timeout=300,
569    volumes={"/data": volume},
570)
571def read_baselines(
572    experiment: str | None = None,
573    model_slug: str | None = None,
574    include_history: bool = False,
575) -> dict[str, Any]:
576    """List all current baselines (no args) or show one (experiment+model_slug),
577    optionally with superseded history."""
578    reload_volume()
579    root = pathlib.Path("/data/baselines")
580    if experiment and model_slug:
581        base_dir = (
582            root
583            / source.sanitize_namespace(experiment)
584            / source.sanitize_namespace(model_slug)
585        )
586        result: dict[str, Any] = {"current": load_json(base_dir / "current.json")}
587        if include_history:
588            history = []
589            history_dir = base_dir / "history"
590            if history_dir.is_dir():
591                for path in sorted(history_dir.glob("*.json")):
592                    entry = load_json(path)
593                    if entry:
594                        history.append(entry)
595            result["history"] = history
596        return result
597    return load_json(root / "index.json") or {"baselines": []}
598
599
600@app.function(
601    image=controller_image,
602    cpu=1,
603    memory=512,
604    timeout=300,
605    volumes={"/data": volume},
606)
607def suite_status(namespace: str, suite_id: str) -> list[dict[str, Any]]:
608    reload_volume()
609    rows = scan_runs(namespace=namespace, suite_id=suite_id)
610    if not rows:
611        raise FileNotFoundError(f"suite not found: {namespace}/{suite_id}")
612    rows.sort(key=lambda row: row.get("created_at") or row.get("run_id") or "")
613    return rows
614
615
616def write_run_inputs(run_dir: pathlib.Path, run_request: dict[str, Any]) -> None:
617    write_json(run_dir / "request.json", run_request)
618    write_json(run_dir / "run-metadata.json", harness_command.run_metadata(run_request))
619    task_names = run_request.get("task_names") or []
620    (run_dir / "selected-tasks.txt").write_text(
621        "\n".join(task_names) + ("\n" if task_names else "")
622    )
623
624
625@app.function(
626    image=controller_image,
627    cpu=1,
628    memory=512,
629    timeout=300,
630    volumes={"/data": volume},
631)
632def create_run_record(run_request: dict[str, Any]) -> dict[str, Any]:
633    namespace = run_request["namespace"]
634    experiment_name = run_request["experiment_name"]
635    run_id = run_request["run_id"]
636    run_dir = pathlib.Path("/data/runs") / namespace / experiment_name / run_id
637    state_path = run_dir / "state.json"
638
639    reload_volume()
640
641    if state_path.exists():
642        raise FileExistsError(
643            f"run record already exists: {namespace}/{experiment_name}/{run_id}"
644        )
645
646    run_dir.mkdir(parents=True, exist_ok=True)
647    write_run_inputs(run_dir, run_request)
648    state = {
649        "run_id": run_id,
650        "namespace": namespace,
651        "experiment_name": experiment_name,
652        "status": "pending",
653        "created_at": run_request.get("created_at"),
654        "updated_at": utc_now(),
655        "build_id": run_request.get("build_id"),
656    }
657    write_json(state_path, state)
658    volume.commit()
659    return state
660
661
662@app.function(
663    image=controller_image,
664    cpu=2,
665    memory=4096,
666    timeout=86_400,
667    volumes={"/data": volume},
668    secrets=[modal.Secret.from_name(MODAL_TOKEN_SECRET_NAME)],
669)
670def run_controller(run_request: dict[str, Any]) -> dict[str, Any]:
671    namespace = run_request["namespace"]
672    experiment_name = run_request["experiment_name"]
673    run_id = run_request["run_id"]
674    run_dir = pathlib.Path("/data/runs") / namespace / experiment_name / run_id
675    run_dir.mkdir(parents=True, exist_ok=True)
676    log_path = run_dir / "controller.log"
677    state_path = run_dir / "state.json"
678    started_at = utc_now()
679
680    def commit() -> None:
681        volume.commit()
682
683    def log(message: str) -> None:
684        line = f"[{utc_now()}] {message}"
685        print(line, flush=True)
686        with log_path.open("a") as log_file:
687            log_file.write(line + "\n")
688
689    def state(status: str, **extra: Any) -> None:
690        payload = {
691            "run_id": run_id,
692            "namespace": namespace,
693            "experiment_name": experiment_name,
694            "status": status,
695            "updated_at": utc_now(),
696            "started_at": started_at,
697            **extra,
698        }
699        write_json(state_path, payload)
700        commit()
701
702    try:
703        write_run_inputs(run_dir, run_request)
704        state("starting", build_id=run_request.get("build_id"))
705        log(f"Starting run {namespace}/{experiment_name}/{run_id}")
706
707        build_id = run_request.get("build_id")
708        if build_id:
709            build_dir = pathlib.Path("/data/builds") / build_id
710            ready_path = build_dir / "READY"
711            build_info_path = build_dir / "build-info.json"
712            source_patch_path = build_dir / "source.patch"
713            source_info_path = build_dir / "source-info.json"
714            state("waiting_for_build", build_id=build_id)
715            deadline = time.time() + int(
716                run_request.get("build_wait_timeout_secs") or 7200
717            )
718            while not ready_path.exists():
719                if time.time() >= deadline:
720                    raise TimeoutError(
721                        f"build {build_id} was not ready before the wait timeout"
722                    )
723                log(f"Waiting for build {build_id} to become ready")
724                time.sleep(30)
725                reload_volume()
726            shutil.copy(build_info_path, run_dir / "build-info.json")
727            if source_patch_path.exists():
728                shutil.copy(source_patch_path, run_dir / "source.patch")
729            if source_info_path.exists():
730                shutil.copy(source_info_path, run_dir / "source-info.json")
731            log(f"Using build {build_id}")
732
733        provision_benchmark_dataset(run_request, log)
734
735        jobs_parent = pathlib.Path("/tmp/agent-evals/harbor-jobs")
736        jobs_parent.mkdir(parents=True, exist_ok=True)
737        command = harness_command.build_harness_command(run_request, str(jobs_parent))
738        redacted = config.redacted_command(command)
739        (run_dir / "harbor-command.txt").write_text(redacted + "\n")
740        state(
741            "running",
742            build_id=build_id,
743            harness_command=redacted,
744        )
745        commit()
746        log("Launching Harbor")
747        log(redacted)
748
749        import zed_eval
750
751        package_file = zed_eval.__file__
752        if package_file is None:
753            raise RuntimeError("could not resolve zed_eval package location")
754        package_parent = str(pathlib.Path(package_file).resolve().parent.parent)
755        environment = os.environ.copy()
756        existing_pythonpath = environment.get("PYTHONPATH")
757        environment["PYTHONPATH"] = (
758            package_parent
759            if not existing_pythonpath
760            else f"{package_parent}:{existing_pythonpath}"
761        )
762
763        process = subprocess.Popen(
764            command,
765            stdout=subprocess.PIPE,
766            stderr=subprocess.STDOUT,
767            text=True,
768            env=environment,
769        )
770        last_commit = time.time()
771        assert process.stdout is not None
772        with log_path.open("a") as log_file:
773            for line in process.stdout:
774                print(line, end="", flush=True)
775                log_file.write(line)
776                if time.time() - last_commit >= 10:
777                    log_file.flush()
778                    commit()
779                    last_commit = time.time()
780        return_code = process.wait()
781        commit()
782        log(f"Harbor exited with status {return_code}")
783
784        job_dir = jobs_parent / run_id
785        archive_path = run_dir / "harbor-job.tar.gz"
786        if job_dir.exists():
787            with tarfile.open(archive_path, "w:gz") as archive:
788                archive.add(job_dir, arcname=job_dir.name)
789            # Count trial dirs by the presence of a trial result.json rather than
790            # a name prefix: Harbor names them `task-<id>__<suffix>` while Pier
791            # uses `<task-id>__<suffix>`.
792            trial_count = sum(
793                1
794                for path in job_dir.iterdir()
795                if path.is_dir() and (path / "result.json").exists()
796            )
797            job_result = load_json(job_dir / "result.json")
798        else:
799            trial_count = 0
800            job_result = None
801            log(f"Harbor job dir was not found at {job_dir}")
802
803        summary = {
804            "run_id": run_id,
805            "namespace": namespace,
806            "experiment_name": experiment_name,
807            "status": "completed" if return_code == 0 else "failed",
808            "harbor_return_code": return_code,
809            "trial_count": trial_count,
810            "has_archive": archive_path.exists(),
811            "job_result": job_result,
812            "started_at": started_at,
813            "completed_at": utc_now(),
814        }
815        write_json(run_dir / "summary.json", summary)
816
817        if return_code == 0:
818            (run_dir / "READY").write_text(utc_now() + "\n")
819            state("completed", summary=summary)
820        else:
821            (run_dir / "FAILED").write_text(utc_now() + "\n")
822            state("failed", summary=summary)
823        commit()
824        return summary
825    except Exception as error:
826        formatted = traceback.format_exc()
827        log(f"FAILED: {error}\n{formatted}")
828        (run_dir / "FAILED").write_text(utc_now() + "\n")
829        state("failed", error=str(error), traceback=formatted)
830        commit()
831        raise
832
833
834def rejudge_run_metadata(rejudge_request: dict[str, Any]) -> dict[str, Any]:
835    judge = config.get_judge(rejudge_request["judge_preset"])
836    parent = rejudge_request["parent"]
837    return {
838        "kind": "rejudge",
839        "run_id": rejudge_request["run_id"],
840        "namespace": rejudge_request["namespace"],
841        "experiment_name": rejudge_request["experiment_name"],
842        "source_run": parent,
843        "judge_preset": rejudge_request["judge_preset"],
844        "judge_model": rejudge_request.get("judge_model") or judge.model,
845        "judge_upstream": judge.upstream,
846        "judge_auth_env": judge.auth_env,
847        "orchestration": config.orchestration_info(),
848        "volume_name": rejudge_request["volume_name"],
849        "api_secret_name": rejudge_request.get("api_secret_name"),
850        "created_at": rejudge_request.get("created_at"),
851    }
852
853
854@app.function(
855    image=controller_image,
856    cpu=2,
857    memory=4096,
858    timeout=86_400,
859    volumes={"/data": volume},
860    secrets=[
861        modal.Secret.from_name(MODAL_TOKEN_SECRET_NAME),
862        modal.Secret.from_name(LLM_PROVIDERS_SECRET_NAME),
863    ],
864)
865def rejudge_controller(rejudge_request: dict[str, Any]) -> dict[str, Any]:
866    """Re-grade a finished parent run with a different judge, producing a new
867    derived run. The parent is read only: its stored agent outputs are re-scored
868    by the real cached verifier through the judge proxy, and only the verdicts
869    change. See `rejudge.py` for the per-trial grading."""
870    namespace = rejudge_request["namespace"]
871    experiment_name = rejudge_request["experiment_name"]
872    run_id = rejudge_request["run_id"]
873    parent = rejudge_request["parent"]
874    judge_preset = rejudge_request["judge_preset"]
875    judge = config.get_judge(judge_preset)
876    judge_model = rejudge_request.get("judge_model") or judge.model
877
878    run_dir = pathlib.Path("/data/runs") / namespace / experiment_name / run_id
879    run_dir.mkdir(parents=True, exist_ok=True)
880    log_path = run_dir / "controller.log"
881    state_path = run_dir / "state.json"
882    started_at = utc_now()
883
884    def log(message: str) -> None:
885        line = f"[{utc_now()}] {message}"
886        print(line, flush=True)
887        with log_path.open("a") as log_file:
888            log_file.write(line + "\n")
889
890    def state(status: str, **extra: Any) -> None:
891        write_json(
892            state_path,
893            {
894                "run_id": run_id,
895                "namespace": namespace,
896                "experiment_name": experiment_name,
897                "kind": "rejudge",
898                "status": status,
899                "updated_at": utc_now(),
900                "started_at": started_at,
901                **extra,
902            },
903        )
904        volume.commit()
905
906    try:
907        write_json(run_dir / "request.json", rejudge_request)
908        write_json(run_dir / "run-metadata.json", rejudge_run_metadata(rejudge_request))
909        state("starting", source_run=parent)
910        log(
911            f"Rejudging {parent['namespace']}/{parent['experiment_name']}/"
912            f"{parent['run_id']} as {namespace}/{experiment_name}/{run_id} "
913            f"with judge {judge_preset} ({judge_model})"
914        )
915
916        reload_volume()
917        parent_dir = (
918            pathlib.Path("/data/runs")
919            / parent["namespace"]
920            / parent["experiment_name"]
921            / parent["run_id"]
922        )
923        parent_archive = parent_dir / "harbor-job.tar.gz"
924        if not parent_archive.exists():
925            raise FileNotFoundError(
926                f"parent run has no harbor-job.tar.gz: {parent_archive}. "
927                "Rejudge needs a completed parent run with a stored job archive."
928            )
929
930        metadata = load_json(parent_dir / "run-metadata.json") or {}
931        benchmark = metadata.get("benchmark") or {}
932        dataset = benchmark.get("dataset") if isinstance(benchmark, dict) else {}
933        dataset_kind = dataset.get("kind") if isinstance(dataset, dict) else None
934        dataset_name = dataset.get("name") if isinstance(dataset, dict) else None
935        if dataset_kind != "registry" or not isinstance(dataset_name, str):
936            raise NotImplementedError(
937                "rejudge currently supports registry datasets (SWE-Atlas rf/qna). "
938                f"Parent run dataset_kind={dataset_kind!r} is not yet supported."
939            )
940
941        work = pathlib.Path("/tmp/agent-evals/rejudge") / run_id
942        if work.exists():
943            shutil.rmtree(work)
944        extract_root = work / "parent"
945        extract_root.mkdir(parents=True, exist_ok=True)
946        log(f"Extracting parent archive {parent_archive}")
947        with tarfile.open(parent_archive, "r:gz") as archive:
948            safe_extract_archive(archive, extract_root)
949        parent_job_dir = extract_root / parent["run_id"]
950        if not parent_job_dir.is_dir():
951            # Fall back to the single top-level dir if the arcname differs.
952            candidates = [p for p in extract_root.iterdir() if p.is_dir()]
953            if len(candidates) != 1:
954                raise FileNotFoundError(
955                    f"could not locate job dir inside {parent_archive}"
956                )
957            parent_job_dir = candidates[0]
958
959        log(f"Downloading task packages for dataset {dataset_name}")
960        subprocess.run(["harbor", "datasets", "download", dataset_name], check=True)
961        tasks_root = pathlib.Path.home() / ".cache" / "harbor" / "tasks" / "packages"
962
963        state("running", source_run=parent, judge_model=judge_model)
964        out_job_dir = work / "job" / run_id
965        summary = rejudge.rejudge_job(
966            parent_job_dir=parent_job_dir,
967            out_job_dir=out_job_dir,
968            tasks_root=tasks_root,
969            judge=judge,
970            judge_model=judge_model,
971            log=log,
972        )
973        volume.commit()
974
975        archive_path = run_dir / "harbor-job.tar.gz"
976        with tarfile.open(archive_path, "w:gz") as archive:
977            archive.add(out_job_dir, arcname=out_job_dir.name)
978
979        summary = {
980            **summary,
981            "run_id": run_id,
982            "namespace": namespace,
983            "experiment_name": experiment_name,
984            "kind": "rejudge",
985            "source_run": parent,
986            "judge_preset": judge_preset,
987            "status": "completed",
988            "has_archive": archive_path.exists(),
989            "started_at": started_at,
990            "completed_at": utc_now(),
991        }
992        write_json(run_dir / "summary.json", summary)
993        (run_dir / "READY").write_text(utc_now() + "\n")
994        log(
995            f"Rejudge complete: {summary['passed_count']}/"
996            f"{summary['rejudged_count']} passed "
997            f"({summary['failed_count']} trials could not be rejudged)"
998        )
999        state("completed", summary=summary)
1000        volume.commit()
1001        return summary
1002    except Exception as error:
1003        formatted = traceback.format_exc()
1004        log(f"FAILED: {error}\n{formatted}")
1005        (run_dir / "FAILED").write_text(utc_now() + "\n")
1006        state("failed", error=str(error), traceback=formatted)
1007        volume.commit()
1008        raise
1009