tenant.openagents/omega

1"""Build the harness command (Harbor or Pier) for a benchmark run request.
2
3The command is driven by the self-describing `benchmark` block embedded in a run
4request (see `benchmarks.benchmark_metadata`), so it works for any registered
5benchmark without a separate experiment registry lookup.
6"""
7
8from __future__ import annotations
9
10import json
11from typing import Any
12
13from . import benchmarks, config
14
15
16def _benchmark_block(run_request: dict[str, Any]) -> dict[str, Any]:
17    block = run_request.get("benchmark")
18    if not isinstance(block, dict):
19        raise ValueError("run request is missing a 'benchmark' block")
20    return block
21
22
23def dataset_args(benchmark: dict[str, Any]) -> list[str]:
24    dataset = benchmark.get("dataset") or {}
25    kind = dataset.get("kind")
26    if kind == benchmarks.DATASET_REGISTRY:
27        name = dataset.get("name")
28        if not name:
29            raise ValueError("registry dataset requires a name")
30        return ["-d", name]
31    if kind in (benchmarks.DATASET_PATH, benchmarks.DATASET_PIER_PATH):
32        data_dir = dataset.get("data_dir")
33        if not data_dir:
34            raise ValueError("path dataset requires data_dir")
35        return ["-p", dataset_path(benchmark)]
36    raise ValueError(f"unsupported dataset kind: {kind}")
37
38
39def dataset_clone_dir(benchmark: dict[str, Any]) -> str:
40    """Where the controller clones the dataset repo for path datasets."""
41    return f"/tmp/datasets/{benchmark['id']}"
42
43
44def dataset_path(benchmark: dict[str, Any]) -> str:
45    """Repo-relative task directory inside the cloned dataset repo."""
46    dataset = benchmark.get("dataset") or {}
47    return f"{dataset_clone_dir(benchmark)}/{dataset.get('data_dir')}"
48
49
50def harness_binary(benchmark: dict[str, Any]) -> str:
51    harness = benchmark.get("harness")
52    if harness not in (benchmarks.HARNESS_HARBOR, benchmarks.HARNESS_PIER):
53        raise ValueError(f"unsupported harness: {harness}")
54    return harness
55
56
57def eval_cli_timeout(run_request: dict[str, Any], benchmark: dict[str, Any]) -> int:
58    return int(
59        run_request.get("eval_cli_timeout")
60        or benchmark.get("default_timeout_secs")
61        or config.DEFAULT_SANDBOX_TIMEOUT_SECS
62    )
63
64
65def build_harness_command(run_request: dict[str, Any], jobs_dir: str) -> list[str]:
66    benchmark = _benchmark_block(run_request)
67    build_id = run_request.get("build_id")
68    if not build_id:
69        raise ValueError("zed benchmarks require build_id")
70    volume_name = run_request["volume_name"]
71    api_secret_name = run_request["api_secret_name"]
72    run_id = run_request["run_id"]
73    agent_model = run_request.get("agent_model") or config.DEFAULT_MODEL
74    n_concurrent = int(run_request.get("n_concurrent") or config.DEFAULT_N_CONCURRENT)
75    sandbox_timeout_secs = int(
76        run_request.get("sandbox_timeout_secs") or config.DEFAULT_SANDBOX_TIMEOUT_SECS
77    )
78    sandbox_idle_timeout_secs = int(
79        run_request.get("sandbox_idle_timeout_secs")
80        or config.DEFAULT_SANDBOX_IDLE_TIMEOUT_SECS
81    )
82    override_cpus = run_request.get("override_cpus") or config.DEFAULT_OVERRIDE_CPUS
83    override_memory_mb = (
84        run_request.get("override_memory_mb") or config.DEFAULT_OVERRIDE_MEMORY_MB
85    )
86
87    secret_names = [api_secret_name]
88    for extra_secret_name in run_request.get("extra_api_secrets") or []:
89        if extra_secret_name not in secret_names:
90            secret_names.append(extra_secret_name)
91
92    command = [
93        harness_binary(benchmark),
94        "run",
95        *dataset_args(benchmark),
96        "-m",
97        agent_model,
98        "--env",
99        "modal",
100        "--n-concurrent",
101        str(n_concurrent),
102        "--job-name",
103        run_id,
104        "--jobs-dir",
105        jobs_dir,
106        "--ek",
107        f"volumes={json.dumps({'/data': volume_name})}",
108        "--ek",
109        f"secrets={json.dumps(secret_names)}",
110        "--ek",
111        f"sandbox_timeout_secs={sandbox_timeout_secs}",
112        "--ek",
113        f"sandbox_idle_timeout_secs={sandbox_idle_timeout_secs}",
114        # Harbor and Pier have different installed-agent interfaces, so each gets
115        # its own thin agent shell (the sanctioned cross-framework exception).
116        "--agent-import-path",
117        (
118            "zed_eval.pier_agent:ZedPierAgent"
119            if harness_binary(benchmark) == benchmarks.HARNESS_PIER
120            else "zed_eval.agent:ZedAgent"
121        ),
122        "--ae",
123        f"EVAL_CLI_CONTAINER_PATH=/data/builds/{build_id}/eval-cli",
124        "--ae",
125        f"EVAL_CLI_TIMEOUT={eval_cli_timeout(run_request, benchmark)}",
126    ]
127    # Omit --override-cpus/--override-memory-mb unless explicitly set, so Harbor
128    # applies each task's declared cpus/memory. Overriding below the declared
129    # values OOM-kills memory-heavy tasks (SIGKILL / exit 137).
130    if override_cpus is not None:
131        command += ["--override-cpus", str(int(override_cpus))]
132    if override_memory_mb is not None:
133        command += ["--override-memory-mb", str(int(override_memory_mb))]
134
135    openai_compatible_provider_json = run_request.get("openai_compatible_provider_json")
136    if openai_compatible_provider_json:
137        command.extend(
138            [
139                "--ae",
140                f"ZED_OPENAI_COMPATIBLE_PROVIDERS={openai_compatible_provider_json}",
141            ]
142        )
143    anthropic_available_models_json = run_request.get("anthropic_available_models_json")
144    if anthropic_available_models_json:
145        command.extend(
146            [
147                "--ae",
148                f"ZED_ANTHROPIC_AVAILABLE_MODELS={anthropic_available_models_json}",
149            ]
150        )
151
152    # NOTE: Pier's network allowlist for air-gapped DeepSWE tasks is declared by
153    # the agent (`network_allowlist()` on a Pier-native agent class), not via a
154    # CLI flag — see benchmarks.AGENT_API_HOSTS and the DeepSWE agent. There is
155    # no `--agent-allow-host` option on `pier run`.
156    if benchmark.get("needs_judge"):
157        judge = config.get_judge(run_request["judge_preset"])
158        judge_model = run_request.get("judge_model") or judge.model
159        command.extend(config.judge_verifier_args(judge, judge_model))
160
161    for key, value in (benchmark.get("env") or {}).items():
162        command.extend(["--ae", f"{key}={value}"])
163    for key, value in (run_request.get("extra_env") or {}).items():
164        command.extend(["--ae", f"{key}={value}"])
165
166    for task_name in run_request.get("task_names") or []:
167        command.extend(["--include-task-name", task_name])
168
169    n_tasks = run_request.get("n_tasks")
170    if n_tasks is not None:
171        command.extend(["--n-tasks", str(n_tasks)])
172
173    command.extend(run_request.get("extra_harbor_args") or [])
174    return command
175
176
177def run_metadata(run_request: dict[str, Any]) -> dict[str, Any]:
178    benchmark = _benchmark_block(run_request)
179    metadata: dict[str, Any] = {
180        "run_id": run_request["run_id"],
181        "namespace": run_request["namespace"],
182        "benchmark": benchmark,
183        "suite_id": run_request.get("suite_id"),
184        "harness": benchmark.get("harness"),
185        "scoring": benchmark.get("scoring"),
186        "build_id": run_request.get("build_id"),
187        "agent_model": run_request.get("agent_model") or config.DEFAULT_MODEL,
188        "orchestration": run_request.get("orchestration")
189        or config.orchestration_info(),
190        "build_toolchain": run_request.get("build_toolchain"),
191        "volume_name": run_request["volume_name"],
192        "api_secret_name": run_request["api_secret_name"],
193        "extra_api_secrets": run_request.get("extra_api_secrets") or [],
194        "openai_compatible_provider_json": run_request.get(
195            "openai_compatible_provider_json"
196        ),
197        "anthropic_available_models_json": run_request.get(
198            "anthropic_available_models_json"
199        ),
200        "n_concurrent": run_request.get("n_concurrent") or config.DEFAULT_N_CONCURRENT,
201        "override_cpus": run_request.get("override_cpus")
202        or config.DEFAULT_OVERRIDE_CPUS,
203        "override_memory_mb": run_request.get("override_memory_mb")
204        or config.DEFAULT_OVERRIDE_MEMORY_MB,
205        "eval_cli_timeout": eval_cli_timeout(run_request, benchmark),
206        "sandbox_timeout_secs": run_request.get("sandbox_timeout_secs")
207        or config.DEFAULT_SANDBOX_TIMEOUT_SECS,
208        "sandbox_idle_timeout_secs": run_request.get("sandbox_idle_timeout_secs")
209        or config.DEFAULT_SANDBOX_IDLE_TIMEOUT_SECS,
210        "task_count": len(run_request.get("task_names") or []),
211        "n_tasks": run_request.get("n_tasks"),
212    }
213    if benchmark.get("needs_judge"):
214        judge = config.get_judge(run_request["judge_preset"])
215        metadata["judge_preset"] = run_request["judge_preset"]
216        metadata["judge_model"] = run_request.get("judge_model") or judge.model
217        metadata["judge_upstream"] = judge.upstream
218        metadata["judge_auth_env"] = judge.auth_env
219        metadata["judge_eval_base_url"] = config.EVAL_BASE_URL_IN_SANDBOX
220    return metadata
221