1008 lines · 35.8 KB · Python
RawDownload
from __future__ import annotations import hashlibimport jsonimport osimport pathlibimport shleximport shutilimport subprocessimport tarfileimport timeimport tracebackimport uuidfrom typing import Any import modal from . import cleanup, config, harness_command, rejudge, sourcefrom .common import load_json, safe_extract_archive, utc_now, write_jsonfrom .source import (    BUILD_TARGET,    CARGO_ZIGBUILD_VERSION,    RUST_IMAGE,    RUST_VERSION,    ZIG_VERSION,    build_toolchain_info,) APP_NAME = os.environ.get("AGENT_EVALS_APP_NAME", "agent-evals")VOLUME_NAME = os.environ.get("AGENT_EVALS_VOLUME", "agent-evals")MODAL_TOKEN_SECRET_NAME = os.environ.get(    "AGENT_EVALS_MODAL_TOKEN_SECRET", "agent-evals-modal-token")# LLM-providers secret mounted into the rejudge controller so the in-controller# judge proxy can read the judge's API key (e.g. BASETEN_API_KEY) the same way a# trial sandbox does. Override at deploy time if the secret is named differently.LLM_PROVIDERS_SECRET_NAME = os.environ.get(    "AGENT_EVALS_LLM_PROVIDERS_SECRET", config.DEFAULT_LLM_PROVIDERS_SECRET_NAME)REPO_URL = os.environ.get(    "AGENT_EVALS_REPO_URL", "https://github.com/zed-industries/zed.git")ZIG_URL = (    f"https://ziglang.org/download/{ZIG_VERSION}/zig-x86_64-linux-{ZIG_VERSION}.tar.xz")CLEANUP_BUILD_RETENTION_DAYS = float(    os.environ.get("AGENT_EVALS_BUILD_RETENTION_DAYS", "14")) app = modal.App(APP_NAME)volume = modal.Volume.from_name(VOLUME_NAME, create_if_missing=True)  def reload_volume() -> None:    try:        volume.reload()    except Exception:        pass  build_image = (    modal.Image.from_registry(RUST_IMAGE, add_python="3.13")    .apt_install("cmake", "build-essential", "curl", "xz-utils", "git")    .run_commands(        f"rustup toolchain install {RUST_VERSION} --profile minimal"        " --component rustfmt --component clippy"        " --component rust-analyzer --component rust-src"        " --target wasm32-wasip2 --target wasm32-unknown-unknown"        " --target x86_64-unknown-linux-musl --target x86_64-unknown-linux-gnu",        f"mkdir -p /opt/zig && curl -fsSL {ZIG_URL}"        " | tar -xJ -C /opt/zig --strip-components=1"        " && ln -s /opt/zig/zig /usr/local/bin/zig",        f"cargo install --locked cargo-zigbuild --version {CARGO_ZIGBUILD_VERSION}",    )) controller_image = (    modal.Image.debian_slim(python_version="3.13")    .apt_install("bash", "ca-certificates", "curl", "git", "tar")    .pip_install(f"modal=={config.MODAL_VERSION}", "uv")    .run_commands(        f"uv tool install harbor=={config.HARBOR_VERSION}"        f" --with 'modal=={config.MODAL_VERSION}' && "        "ln -sf /root/.local/bin/harbor /usr/local/bin/harbor",        # Pier is a Harbor fork that adds per-agent network allowlists, required        # to run DeepSWE's air-gapped (`allow_internet = false`) tasks while        # still letting eval-cli reach the model API. Kept non-fatal so a Pier        # install problem can't break the Harbor benchmarks that share this image;        # DeepSWE runs surface a clear "pier not found" error at runtime instead.        "(uv tool install git+https://github.com/datacurve-ai/pier"        f" --with 'modal=={config.MODAL_VERSION}' && "        "ln -sf /root/.local/bin/pier /usr/local/bin/pier) || "        "echo 'WARNING: pier install failed; DeepSWE runs will not work'",    ))  def provision_benchmark_dataset(run_request: dict[str, Any], log: Any) -> None:    """Clone a path-backed benchmark dataset (SWE-Atlas tw, DeepSWE) into the    location the harness command expects. Registry datasets need no provisioning;    the harness pulls them from its hub."""    benchmark = run_request["benchmark"]    dataset = benchmark.get("dataset") or {}    kind = dataset.get("kind")    if kind not in ("path", "pier_path"):        return     repo_url = dataset.get("repo_url")    repo_ref = dataset.get("repo_ref") or "main"    if not repo_url:        raise ValueError(f"benchmark {benchmark['id']} path dataset requires repo_url")     clone_dir = pathlib.Path(harness_command.dataset_clone_dir(benchmark))    if clone_dir.exists():        shutil.rmtree(clone_dir)    clone_dir.mkdir(parents=True, exist_ok=True)    log(f"Fetching {benchmark['id']} dataset {repo_url}@{repo_ref}")    subprocess.run(["git", "init", "-q", str(clone_dir)], check=True)    subprocess.run(        ["git", "fetch", "--depth", "1", repo_url, repo_ref],        cwd=clone_dir,        check=True,    )    subprocess.run(["git", "checkout", "-q", "FETCH_HEAD"], cwd=clone_dir, check=True)    data_dir = pathlib.Path(harness_command.dataset_path(benchmark))    if not data_dir.exists():        raise FileNotFoundError(f"benchmark dataset directory not found: {data_dir}")  def output_of(command: list[str]) -> str:    try:        return subprocess.run(            command,            check=True,            capture_output=True,            text=True,        ).stdout.strip()    except subprocess.CalledProcessError as error:        return f"unknown ({error})"  @app.function(    image=build_image,    # Right-sized from the original 16 cpu / 32 GB after observing peak usage of    # ~10 cores and ~10 GB. ephemeral_disk stays at Modal's 512 GiB floor (the    # minimum allowed for a function), so it isn't a tunable here.    cpu=12,    memory=24576,    ephemeral_disk=524288,    timeout=7200,    volumes={"/data": volume},)def build_eval_cli(build_request: dict[str, Any]) -> dict[str, Any]:    def run(command: str, **kwargs: Any) -> None:        print(f"+ {command}", flush=True)        subprocess.run(command, shell=True, check=True, **kwargs)     reload_volume()     build_id = build_request["build_id"]    source_info = build_request.get("source") or {"type": "git_patch"}    base_sha = source_info.get("base_sha") or build_request["base_sha"]    patch = build_request.get("patch") or ""    build_dir = pathlib.Path("/data/builds") / build_id    ready_path = build_dir / "READY"    binary_path = build_dir / "eval-cli"    build_info_path = build_dir / "build-info.json"    # The lease lives outside build_dir so it never interferes with the atomic    # move of the finished build directory below.    building_path = pathlib.Path("/data/build-locks") / f"{build_id}.json"    lease_ttl = int(build_request.get("build_wait_timeout_secs") or 7200)    owner = uuid.uuid4().hex     def reuse_existing() -> dict[str, Any]:        build_info = load_json(build_info_path) or {}        build_info.setdefault("build_id", build_id)        build_info["reused"] = True        volume.commit()        print(f"Reusing existing build {build_id}", flush=True)        return build_info     if ready_path.exists() and binary_path.exists():        return reuse_existing()    if build_dir.exists():        raise RuntimeError(            f"build directory already exists but is not ready; refusing to overwrite {build_id}"        )     # Single-flight lease: if another invocation is already compiling this exact    # build, wait for it to finish rather than running a second multi-minute    # compile. The lease is best-effort (the volume has no atomic compare-and-swap);    # the atomic move below still guarantees correctness if two builds slip through.    deadline = time.time() + lease_ttl    while True:        reload_volume()        if ready_path.exists() and binary_path.exists():            return reuse_existing()        if build_dir.exists():            raise RuntimeError(                f"build directory already exists but is not ready; refusing to overwrite {build_id}"            )        lease = load_json(building_path)        now = time.time()        held_by_other = (            isinstance(lease, dict)            and lease.get("owner") != owner            and (now - float(lease.get("epoch") or 0)) < lease_ttl        )        if held_by_other and time.time() < deadline:            print(                f"Build {build_id} is being compiled elsewhere; waiting",                flush=True,            )            time.sleep(15)            continue        break    # Claim the lease (write_json creates /data/build-locks). build_dir is    # intentionally NOT created here so the atomic move below still works.    write_json(building_path, {"owner": owner, "epoch": time.time(), "at": utc_now()})    volume.commit()     rustc_version = output_of(["rustc", "--version"])    zig_version = output_of(["zig", "version"])    cargo_zigbuild_version = output_of(["cargo-zigbuild", "--version"])     workdir = pathlib.Path("/build/zed")    if workdir.exists():        shutil.rmtree(workdir)    workdir.mkdir(parents=True, exist_ok=True)    os.chdir(workdir)     repo_url = source_info.get("repo_url") or REPO_URL    run("git init -q .")    run(f"git fetch --depth 1 {shlex.quote(repo_url)} {shlex.quote(base_sha)}")    run("git checkout -q FETCH_HEAD")     if patch.strip():        patch_file = pathlib.Path("/build/source.patch")        patch_file.write_text(patch)        run(f"git apply --stat {shlex.quote(str(patch_file))}")        run(f"git apply {shlex.quote(str(patch_file))}")     patch_sha256 = source_info.get("patch_sha256")     run(        "cargo zigbuild --release --package eval_cli --target x86_64-unknown-linux-musl"    )     built = workdir / "target/x86_64-unknown-linux-musl/release/eval-cli"    run(f"strip {built}")    binary_bytes = built.read_bytes()    binary_sha256 = hashlib.sha256(binary_bytes).hexdigest()     temporary_dir = pathlib.Path("/data/tmp/builds") / f"{build_id}-{uuid.uuid4().hex}"    temporary_dir.mkdir(parents=True, exist_ok=True)    shutil.copy(built, temporary_dir / "eval-cli")    (temporary_dir / "eval-cli").chmod(0o755)    if patch.strip():        (temporary_dir / "source.patch").write_text(patch)    write_json(temporary_dir / "source-info.json", source_info)     build_info = {        "build_id": build_id,        "base_sha": base_sha,        "patch_sha256": patch_sha256,        "rustc_version": rustc_version,        "zig_version": zig_version,        "cargo_zigbuild_version": cargo_zigbuild_version,        "built_at_utc": utc_now(),        "binary_size_bytes": len(binary_bytes),        "binary_sha256": binary_sha256,        "target": BUILD_TARGET,        "rust_image": RUST_IMAGE,        "toolchain": build_toolchain_info(),        "source": source_info,    }    write_json(temporary_dir / "build-info.json", build_info)    (temporary_dir / "READY").write_text(utc_now() + "\n")     if not build_dir.exists():        build_dir.parent.mkdir(parents=True, exist_ok=True)        shutil.move(str(temporary_dir), str(build_dir))    elif ready_path.exists() and binary_path.exists():        shutil.rmtree(temporary_dir, ignore_errors=True)        build_info = load_json(build_info_path) or build_info        build_info["reused_after_race"] = True    else:        shutil.rmtree(temporary_dir, ignore_errors=True)        raise RuntimeError(            f"build directory already exists but is not ready; refusing to overwrite {build_id}"        )     # Release the single-flight lease now that READY exists.    try:        building_path.unlink()    except OSError:        pass    volume.commit()    print(f"Committed build {build_id} to volume '{VOLUME_NAME}'", flush=True)    return build_info  @app.function(    image=controller_image,    cpu=1,    memory=512,    timeout=300,    volumes={"/data": volume},)def list_builds(limit: int = 50) -> list[dict[str, Any]]:    reload_volume()    builds_dir = pathlib.Path("/data/builds")    rows = []    for build_dir in builds_dir.iterdir() if builds_dir.exists() else []:        if not build_dir.is_dir():            continue        build_info = load_json(build_dir / "build-info.json") or {}        ready = (build_dir / "READY").exists() and (build_dir / "eval-cli").exists()        rows.append(            {                "build_id": build_dir.name,                "ready": ready,                "base_sha": build_info.get("base_sha"),                "patch_sha256": build_info.get("patch_sha256"),                "built_at_utc": build_info.get("built_at_utc"),                "binary_sha256": build_info.get("binary_sha256"),                "source": build_info.get("source"),            }        )    rows.sort(key=lambda row: row.get("built_at_utc") or "", reverse=True)    return rows[:limit]  def run_row(    namespace_dir: pathlib.Path, experiment_dir: pathlib.Path, run_dir: pathlib.Path) -> dict[str, Any]:    state = load_json(run_dir / "state.json") or {}    request = load_json(run_dir / "request.json") or {}    metadata = load_json(run_dir / "run-metadata.json") or {}    summary = load_json(run_dir / "summary.json") or {}    benchmark = request.get("benchmark") or metadata.get("benchmark") or {}    benchmark_id = benchmark.get("id") if isinstance(benchmark, dict) else None    suite_id = request.get("suite_id") or metadata.get("suite_id")    return {        "namespace": namespace_dir.name,        "experiment_name": experiment_dir.name,        "run_id": run_dir.name,        "status": state.get("status"),        "updated_at": state.get("updated_at"),        "created_at": metadata.get("created_at")        or state.get("created_at")        or request.get("created_at"),        "agent_model": metadata.get("agent_model"),        "judge_preset": metadata.get("judge_preset"),        "judge_model": metadata.get("judge_model"),        "build_id": metadata.get("build_id") or state.get("build_id"),        "trial_count": summary.get("trial_count"),        "has_archive": summary.get("has_archive"),        "suite_id": suite_id,        "part": request.get("suite_part"),        "benchmark": benchmark_id,    }  def scan_runs(    *,    namespace: str | None = None,    experiment_name: str | None = None,    suite_id: str | None = None,) -> list[dict[str, Any]]:    runs_root = pathlib.Path("/data/runs")    rows = []    if not runs_root.exists():        return []    namespaces = (        [runs_root / namespace] if namespace else [path for path in runs_root.iterdir()]    )    for namespace_dir in namespaces:        if not namespace_dir.is_dir():            continue        experiments = (            [namespace_dir / experiment_name]            if experiment_name            else [path for path in namespace_dir.iterdir()]        )        for experiment_dir in experiments:            if not experiment_dir.is_dir():                continue            for run_dir in experiment_dir.iterdir():                if not run_dir.is_dir():                    continue                row = run_row(namespace_dir, experiment_dir, run_dir)                if suite_id and row.get("suite_id") != suite_id:                    continue                rows.append(row)    rows.sort(        key=lambda row: row.get("updated_at") or row.get("created_at") or "",        reverse=True,    )    return rows  @app.function(    image=controller_image,    cpu=1,    memory=512,    timeout=300,    volumes={"/data": volume},)def list_runs(    namespace: str | None = None,    experiment_name: str | None = None,    limit: int = 50,) -> list[dict[str, Any]]:    reload_volume()    rows = scan_runs(namespace=namespace, experiment_name=experiment_name)    return rows[:limit]  def _run_cleanup(    *,    dry_run: bool,    build_retention_days: float | None = None,) -> dict[str, Any]:    reload_volume()    result = cleanup.prune_artifacts(        pathlib.Path("/data"),        dry_run=dry_run,        build_retention_days=(            build_retention_days            if build_retention_days is not None            else CLEANUP_BUILD_RETENTION_DAYS        ),    )    if not dry_run:        volume.commit()    print(json.dumps(result, indent=2), flush=True)    return result  @app.function(    image=controller_image,    cpu=1,    memory=2048,    timeout=1800,    volumes={"/data": volume},    schedule=modal.Period(days=1),)def cleanup_scheduled() -> dict[str, Any]:    """Daily prune of stale build artifacts. Never touches eval results."""    return _run_cleanup(dry_run=False)  @app.function(    image=controller_image,    cpu=1,    memory=2048,    timeout=1800,    volumes={"/data": volume},)def cleanup_artifacts(request: dict[str, Any]) -> dict[str, Any]:    """On-demand prune (the `zed-eval cleanup` command), supporting --dry-run    and retention overrides."""    return _run_cleanup(        dry_run=bool(request.get("dry_run")),        build_retention_days=request.get("build_retention_days"),    )  @app.function(    image=controller_image,    cpu=1,    memory=512,    timeout=300,    volumes={"/data": volume},)def read_run_provenance(    namespace: str, experiment_name: str, run_id: str) -> dict[str, Any]:    """Provenance a baseline record needs for one run: the launch request    (model/judge/resources/build_id), the run summary (status), and the build's    build-info (base_sha/patch_sha256/source label)."""    reload_volume()    run_dir = pathlib.Path("/data/runs") / namespace / experiment_name / run_id    request = load_json(run_dir / "request.json") or {}    summary = load_json(run_dir / "summary.json") or {}    metadata = load_json(run_dir / "run-metadata.json") or {}    build_id = request.get("build_id") or metadata.get("build_id")    build_info: dict[str, Any] = {}    if build_id:        build_info = (            load_json(pathlib.Path("/data/builds") / build_id / "build-info.json") or {}        )    return {        "request": request,        "summary": summary,        "metadata": metadata,        "build_info": build_info,    } 

Lines 1–500 of 1008.