tenant.openagents/omega

1from __future__ import annotations
2
3import argparse
4import json
5import sys
6import uuid
7from pathlib import Path
8from typing import Any
9
10from . import benchmarks, config, harness_command, run_index, source
11from .builds import prepare_build_request, resolve_source, validate_build_id
12from .common import (
13    dedupe_preserving_order,
14    default_namespace,
15    deployed_function,
16    modal_call_id,
17    parse_parts,
18    print_json,
19    utc_now,
20    utc_timestamp,
21)
22from .volume import build_ready_on_volume
23
24# Local scratch dir used only for dry-run/plan previews of harness commands.
25PREVIEW_JOBS_DIR = "/tmp/agent-evals/harbor-jobs"
26
27
28def read_task_file(path: str | None) -> list[str]:
29    if not path:
30        return []
31    return [
32        task_name
33        for line in Path(path).read_text().splitlines()
34        if (task_name := line.strip()) and not task_name.startswith("#")
35    ]
36
37
38def baseten_provider_json(
39    *,
40    model_id: str,
41    api_url: str,
42    display_name: str | None,
43    max_tokens: int,
44    max_output_tokens: int,
45) -> str:
46    provider = {
47        config.BASETEN_PROVIDER_ID: {
48            "api_url": api_url,
49            "available_models": [
50                {
51                    "name": model_id,
52                    "display_name": display_name or model_id,
53                    "max_tokens": max_tokens,
54                    "max_output_tokens": max_output_tokens,
55                    "capabilities": {
56                        "tools": True,
57                        "images": False,
58                        "parallel_tool_calls": False,
59                        "prompt_cache_key": False,
60                    },
61                }
62            ],
63        }
64    }
65    return json.dumps(provider, separators=(",", ":"))
66
67
68def resolve_model_preset(model: str) -> str:
69    resolved = config.resolve_model_preset(model)
70    if resolved != model:
71        return resolved
72    if model.startswith("baseten:"):
73        return f"{config.BASETEN_PROVIDER_ID}/{model.split(':', 1)[1]}"
74    return model
75
76
77def resolve_model_options(
78    args: argparse.Namespace,
79) -> tuple[str, str | None, list[str]]:
80    raw_model = getattr(args, "model", None) or config.DEFAULT_MODEL
81    model = resolve_model_preset(raw_model)
82    openai_compatible_provider_json = getattr(
83        args, "openai_compatible_provider_json", None
84    )
85    extra_api_secrets = list(getattr(args, "extra_api_secret", None) or [])
86
87    if getattr(args, "model_provider", None) == "baseten":
88        baseten_model = getattr(args, "baseten_model", None) or model
89        if baseten_model.startswith(f"{config.BASETEN_PROVIDER_ID}/"):
90            baseten_model = baseten_model.split("/", 1)[1]
91        if baseten_model.startswith("baseten:"):
92            baseten_model = baseten_model.split(":", 1)[1]
93        model = f"{config.BASETEN_PROVIDER_ID}/{baseten_model}"
94
95    if model.startswith(f"{config.BASETEN_PROVIDER_ID}/"):
96        baseten_model = model.split("/", 1)[1]
97        if not openai_compatible_provider_json:
98            openai_compatible_provider_json = baseten_provider_json(
99                model_id=baseten_model,
100                api_url=getattr(args, "baseten_api_url", config.BASETEN_API_URL),
101                display_name=getattr(args, "baseten_model_display_name", None),
102                max_tokens=getattr(
103                    args, "baseten_model_max_tokens", config.BASETEN_DEFAULT_MAX_TOKENS
104                ),
105                max_output_tokens=getattr(
106                    args,
107                    "baseten_model_max_output_tokens",
108                    config.BASETEN_DEFAULT_MAX_OUTPUT_TOKENS,
109                ),
110            )
111
112    return model, openai_compatible_provider_json, extra_api_secrets
113
114
115def derive_run_id(
116    base_run_id: str | None, suite_id: str | None, suffix: str, index: int
117) -> str:
118    """Pick a run id for one leg of a (possibly multi-target) launch.
119
120    An explicit --run-id is used verbatim for the first leg and suffixed for the
121    rest; otherwise legs hang off the suite id, falling back to a timestamped id
122    for a lone run.
123    """
124    if base_run_id:
125        return base_run_id if index == 0 else f"{base_run_id}-{suffix}"
126    if suite_id:
127        return f"{suite_id}-{suffix}-{uuid.uuid4().hex[:6]}"
128    return f"{utc_timestamp()}-{uuid.uuid4().hex[:6]}"
129
130
131def mint_suite_id(args: argparse.Namespace, run_count: int) -> str | None:
132    """A suite id groups multiple runs from one invocation; None for a lone run."""
133    if run_count <= 1:
134        return None
135    explicit = getattr(args, "suite_id", None)
136    if explicit:
137        return explicit
138    prefix = getattr(args, "run_id", None) or "run"
139    return f"{source.sanitize_namespace(prefix)}-{utc_timestamp()}"
140
141
142def build_rejudge_request(args: argparse.Namespace) -> dict[str, Any]:
143    """Build the request for re-grading an existing run with a different judge.
144
145    The positional `run_id` is the *parent* run; the derived run gets a new id
146    under the same experiment so `report`/`list` group them together. Only the
147    judge differs — no build, no agent model, no task selection.
148    """
149    judge_preset = args.judge
150    config.get_judge(judge_preset)  # validate before spawning anything remote
151    if getattr(args, "experiment_name", None):
152        experiment_name = source.sanitize_namespace(args.experiment_name)
153        namespace = default_namespace(args)
154    else:
155        entry = run_index.lookup(args.run_id)
156        if not entry:
157            raise ValueError(
158                f"could not locate run '{args.run_id}' in the local run index "
159                f"({run_index.index_path()}). Pass --experiment-name (and "
160                "--namespace if it isn't yours)."
161            )
162        experiment_name = source.sanitize_namespace(entry["experiment_name"])
163        namespace = source.sanitize_namespace(
164            getattr(args, "namespace", None) or entry["namespace"]
165        )
166    parent_namespace = (
167        source.sanitize_namespace(args.parent_namespace)
168        if getattr(args, "parent_namespace", None)
169        else namespace
170    )
171    parent_run_id = args.run_id
172    judge_slug = source.sanitize_namespace(judge_preset)
173    new_run_id = (
174        args.new_run_id
175        or f"{parent_run_id}-rejudge-{judge_slug}-{uuid.uuid4().hex[:6]}"
176    )
177    return {
178        "namespace": namespace,
179        "experiment_name": experiment_name,
180        "run_id": new_run_id,
181        "parent": {
182            "namespace": parent_namespace,
183            "experiment_name": experiment_name,
184            "run_id": parent_run_id,
185        },
186        "judge_preset": judge_preset,
187        "judge_model": getattr(args, "judge_model", None),
188        "volume_name": args.volume,
189        "api_secret_name": args.api_secret,
190        "created_at": utc_now(),
191    }
192
193
194def command_rejudge(args: argparse.Namespace) -> int:
195    rejudge_request = build_rejudge_request(args)
196    if getattr(args, "dry_run", False) or getattr(args, "plan", False):
197        print_json(rejudge_request)
198        return 0
199
200    controller = deployed_function(args, "rejudge_controller")
201    call = controller.spawn(rejudge_request)
202    run_index.record_run(
203        {**rejudge_request, "volume_name": args.volume, "kind": "rejudge"}
204    )
205    parent = rejudge_request["parent"]
206    print(f"Namespace:  {rejudge_request['namespace']}")
207    print(f"Experiment: {rejudge_request['experiment_name']}")
208    print(
209        f"Source run: {parent['namespace']}/{parent['experiment_name']}/"
210        f"{parent['run_id']}"
211    )
212    print(f"New run id: {rejudge_request['run_id']}")
213    print(f"Judge:      {rejudge_request['judge_preset']}")
214    print(f"Spawned rejudge controller: {modal_call_id(call)}")
215    print("\nNext steps (run id alone is enough):")
216    print(f"  zed-eval status {rejudge_request['run_id']}")
217    print(f"  zed-eval report {rejudge_request['run_id']} --fetch")
218    return 0
219
220
221def common_run_request_fields(
222    args: argparse.Namespace,
223    *,
224    namespace: str,
225    run_id: str,
226    experiment_name: str,
227    judge_preset: str | None,
228    build_id: str | None,
229    suite_id: str | None,
230) -> dict[str, Any]:
231    """Fields shared by every benchmark run request."""
232    agent_model, openai_compatible_provider_json, extra_api_secrets = (
233        resolve_model_options(args)
234    )
235    task_names = dedupe_preserving_order(
236        read_task_file(getattr(args, "tasks", None))
237        + (getattr(args, "include_task_name", None) or [])
238    )
239    return {
240        "created_at": utc_now(),
241        "namespace": namespace,
242        "run_id": run_id,
243        "experiment_name": experiment_name,
244        "volume_name": args.volume,
245        "api_secret_name": args.api_secret,
246        "modal_token_secret_name": args.modal_token_secret,
247        "orchestration": config.orchestration_info(),
248        "agent_model": agent_model,
249        "judge_preset": judge_preset,
250        "judge_model": getattr(args, "judge_model", None),
251        "build_id": build_id,
252        "task_names": task_names,
253        "n_tasks": getattr(args, "n_tasks", None),
254        "n_concurrent": getattr(args, "n_concurrent", config.DEFAULT_N_CONCURRENT),
255        "override_cpus": getattr(args, "override_cpus", config.DEFAULT_OVERRIDE_CPUS),
256        "override_memory_mb": getattr(
257            args, "override_memory_mb", config.DEFAULT_OVERRIDE_MEMORY_MB
258        ),
259        "sandbox_timeout_secs": getattr(
260            args, "sandbox_timeout_secs", config.DEFAULT_SANDBOX_TIMEOUT_SECS
261        ),
262        "sandbox_idle_timeout_secs": getattr(
263            args, "sandbox_idle_timeout_secs", config.DEFAULT_SANDBOX_IDLE_TIMEOUT_SECS
264        ),
265        "build_wait_timeout_secs": getattr(args, "build_wait_timeout_secs", 7200),
266        "extra_harbor_args": list(getattr(args, "extra_harbor_arg", None) or []),
267        "openai_compatible_provider_json": openai_compatible_provider_json,
268        "anthropic_available_models_json": getattr(
269            args, "anthropic_available_models_json", None
270        ),
271        "extra_api_secrets": dedupe_preserving_order(extra_api_secrets),
272        "suite_id": suite_id,
273    }
274
275
276def benchmark_plan_entry(
277    benchmark_id: str, run_request: dict[str, Any], build_request: dict[str, Any] | None
278) -> dict[str, Any]:
279    return {
280        "benchmark": benchmark_id,
281        "run_id": run_request["run_id"],
282        "harness": run_request["benchmark"]["harness"],
283        "model": run_request["agent_model"],
284        "judge": run_request.get("judge_preset"),
285        "build_id": run_request.get("build_id"),
286        "will_build": build_request is not None,
287        "n_tasks": run_request.get("n_tasks"),
288    }
289
290
291def print_plan(
292    summaries: list[dict[str, Any]],
293    details: list[tuple[str, dict[str, Any], dict[str, Any] | None]],
294    *,
295    verbose: bool,
296) -> None:
297    print("Plan:")
298    print_json(summaries)
299    if verbose:
300        for header, run_request, build_request in details:
301            print(f"\n=== {header} ===")
302            print_dry_run(run_request, build_request)
303
304
305def print_dry_run(
306    run_request: dict[str, Any], build_request: dict[str, Any] | None
307) -> None:
308    print("Run request:")
309    print_json(run_request)
310    if build_request:
311        print("\nBuild request:")
312        build_preview = dict(build_request)
313        patch = build_preview.pop("patch", "")
314        build_preview["patch_line_count"] = len(patch.splitlines())
315        build_preview["source"] = source.public_source_info(
316            build_request.get("source") or {}
317        )
318        print_json(build_preview)
319    command = config.redacted_command(
320        harness_command.build_harness_command(run_request, PREVIEW_JOBS_DIR)
321    )
322    print("\nHarness command:")
323    print(command)
324
325
326def execute_prepared_runs(
327    args: argparse.Namespace,
328    prepared: list[tuple[str, dict[str, Any], dict[str, Any] | None]],
329    plan_entry,
330) -> int:
331    if args.plan:
332        print_plan(
333            [
334                plan_entry(label, run_request, build_request)
335                for label, run_request, build_request in prepared
336            ],
337            prepared,
338            verbose=args.verbose,
339        )
340        return 0
341    if args.dry_run:
342        for label, run_request, build_request in prepared:
343            print(f"\n=== {label} ===")
344            print_dry_run(run_request, build_request)
345        return 0
346
347    spawned_builds: set[str] = set()
348    for label, run_request, build_request in prepared:
349        print(f"\n=== Launching {label} ===")
350        launch_prepared_run(args, run_request, build_request, spawned_builds)
351    return 0
352
353
354def launch_prepared_run(
355    args: argparse.Namespace,
356    run_request: dict[str, Any],
357    build_request: dict[str, Any] | None,
358    spawned_builds: set[str] | None = None,
359) -> None:
360    build_function = None
361    if build_request:
362        run_request["source"] = source.public_source_info(build_request["source"])
363        print_untracked_warning(build_request)
364        build_function = deployed_function(args, "build_eval_cli")
365    record_function = deployed_function(args, "create_run_record")
366    controller_function = deployed_function(args, "run_controller")
367
368    record_state = record_function.remote(run_request)
369    run_index.record_run(run_request)
370
371    print(f"Namespace:  {run_request['namespace']}")
372    print(f"Experiment: {run_request['experiment_name']}")
373    print(f"Run id:     {run_request['run_id']}")
374    print(f"Volume:     {run_request['volume_name']}")
375    print(f"Run state:  {record_state['status']}")
376    print(f"Model:      {run_request['agent_model']}")
377    if run_request.get("judge_preset"):
378        print(f"Judge:      {run_request['judge_preset']}")
379    if run_request.get("suite_id"):
380        print(f"Suite:      {run_request['suite_id']}")
381    if run_request.get("build_id"):
382        print(f"Build id:  {run_request['build_id']}")
383    if run_request.get("task_names"):
384        print(f"Tasks:     {len(run_request['task_names'])} explicit task(s)")
385    elif run_request.get("n_tasks"):
386        print(f"Tasks:     Harbor --n-tasks {run_request['n_tasks']}")
387    else:
388        print("Tasks:     full dataset selection")
389
390    build_id = run_request.get("build_id")
391    if build_function is not None and build_id not in (spawned_builds or set()):
392        build_call = build_function.spawn(build_request)
393        print(f"Spawned build:      {modal_call_id(build_call)}")
394        if spawned_builds is not None and build_id:
395            spawned_builds.add(build_id)
396    controller_call = controller_function.spawn(run_request)
397    print(f"Spawned controller: {modal_call_id(controller_call)}")
398
399    run_id = run_request["run_id"]
400    print(
401        "\nNext steps (run id alone is enough; namespace/experiment are resolved "
402        "from this machine's local run index):"
403    )
404    print(f"  zed-eval status {run_id}")
405    print(f"  zed-eval logs {run_id}")
406    print(f"  zed-eval report {run_id} --fetch")
407
408
409def resolve_benchmark_judge(
410    args: argparse.Namespace, benchmark: benchmarks.Benchmark
411) -> str | None:
412    if not benchmark.needs_judge:
413        return None
414    judge = getattr(args, "judge", None) or config.DEFAULT_JUDGE_PRESET
415    if judge == "auto":
416        return benchmark.default_judge or "leaderboard"
417    config.get_judge(judge)
418    return judge
419
420
421def benchmark_metadata_for_run(
422    args: argparse.Namespace, benchmark: benchmarks.Benchmark
423) -> dict[str, object]:
424    metadata = benchmarks.benchmark_metadata(benchmark)
425    dataset = metadata.get("dataset")
426    if (
427        isinstance(dataset, dict)
428        and dataset.get("repo_url") == benchmarks.SWE_ATLAS_REPO_URL
429    ):
430        dataset["repo_url"] = (
431            getattr(args, "swe_atlas_repo_url", None) or benchmarks.SWE_ATLAS_REPO_URL
432        )
433        dataset["repo_ref"] = (
434            getattr(args, "swe_atlas_repo_ref", None) or benchmarks.SWE_ATLAS_REPO_REF
435        )
436    return metadata
437
438
439def print_untracked_warning(build_request: dict[str, Any]) -> None:
440    build_source = build_request.get("source") or {}
441    untracked_files = build_source.get("untracked_files") or []
442    if untracked_files:
443        print(
444            f"Warning: proceeding with {len(untracked_files)} untracked file(s) "
445            "not included in the build patch.",
446            file=sys.stderr,
447        )
448
449
450def prepare_shared_build(
451    args: argparse.Namespace,
452) -> tuple[str | None, dict[str, Any] | None]:
453    """Resolve the build once for a whole (possibly multi-benchmark) run.
454
455    Returns `(build_id, build_request)`. `build_request` is None when the target
456    build already exists; otherwise the caller should spawn the build.
457    """
458    explicit_build_id = getattr(args, "build", None)
459    validate_build_id(explicit_build_id)
460    if (
461        explicit_build_id
462        and not getattr(args, "plan", False)
463        and not getattr(args, "dry_run", False)
464        and build_ready_on_volume(args, explicit_build_id)
465    ):
466        print(f"Reusing existing build {explicit_build_id} (already on volume)")
467        return explicit_build_id, None
468
469    base_sha, clean_source, source_label, pre_resolved = resolve_source(args)
470
471    build_request = prepare_build_request(
472        base_sha=base_sha,
473        patch_path=getattr(args, "patch_path", None),
474        build_id=explicit_build_id,
475        allow_untracked=getattr(args, "allow_untracked", False),
476        require_clean=getattr(args, "require_clean", False),
477        repo_url=getattr(args, "repo_url", None),
478        clean_source=clean_source,
479        source_label=source_label,
480        pre_resolved_base_sha=pre_resolved,
481    )
482    build_id = build_request["build_id"]
483
484    if not getattr(args, "plan", False) and not getattr(args, "dry_run", False):
485        if build_ready_on_volume(args, build_id):
486            print(f"Reusing existing build {build_id} (already on volume)")
487            return build_id, None
488
489    return build_id, build_request
490
491
492def command_build(args: argparse.Namespace) -> int:
493    build_id, build_request = prepare_shared_build(args)
494    if build_request is None:
495        return 0
496
497    print(f"Build id: {build_id}")
498    print(f"Base sha: {build_request['base_sha']}")
499    print(f"Patch sha256: {build_request['patch_sha256'] or '(none)'}")
500    print_untracked_warning(build_request)
501
502    build_function = deployed_function(args, "build_eval_cli")
503    if args.detach:
504        call = build_function.spawn(build_request)
505        print(f"Spawned build: {modal_call_id(call)}")
506    else:
507        result = build_function.remote(build_request)
508        print_json(result)
509    return 0
510
511
512def build_benchmark_run_request(
513    args: argparse.Namespace,
514    *,
515    benchmark_id: str,
516    build_id: str | None,
517    suite_id: str | None,
518    index: int,
519    run_id_suffix: str | None = None,
520) -> dict[str, Any]:
521    benchmark = benchmarks.get_benchmark(benchmark_id)
522    run_id = derive_run_id(
523        getattr(args, "run_id", None), suite_id, run_id_suffix or benchmark_id, index
524    )
525    # Staff mode is off by default for remote runs: it enables the sandboxed
526    # terminal, which hangs inside Modal sandboxes.
527    extra_env = {"EVAL_CLI_STAFF": "true" if getattr(args, "staff", False) else "false"}
528
529    return {
530        **common_run_request_fields(
531            args,
532            namespace=default_namespace(args),
533            run_id=run_id,
534            # The benchmark id doubles as the experiment name for run storage
535            # paths (runs/<namespace>/<experiment_name>/<run_id>), keeping
536            # monitoring and fetch uniform across benchmarks.
537            experiment_name=source.sanitize_namespace(benchmark_id),
538            judge_preset=resolve_benchmark_judge(args, benchmark),
539            build_id=build_id,
540            suite_id=suite_id,
541        ),
542        "benchmark": benchmark_metadata_for_run(args, benchmark),
543        "eval_cli_timeout": getattr(args, "eval_cli_timeout", None),
544        "extra_env": extra_env,
545    }
546
547
548def prepare_runs_for_benchmarks(
549    args: argparse.Namespace,
550    benchmark_ids: list[str],
551    *,
552    suite_id: str | None,
553    label_for_benchmark,
554    mark_swe_atlas_parts: bool = False,
555) -> list[tuple[str, dict[str, Any], dict[str, Any] | None]]:
556    if not benchmark_ids:
557        raise ValueError("choose at least one benchmark to run")
558
559    build_id, build_request = prepare_shared_build(args)
560    prepared: list[tuple[str, dict[str, Any], dict[str, Any] | None]] = []
561    for index, benchmark_id in enumerate(benchmark_ids):
562        label = label_for_benchmark(benchmark_id)
563        run_request = build_benchmark_run_request(
564            args,
565            benchmark_id=benchmark_id,
566            build_id=build_id,
567            suite_id=suite_id,
568            index=index,
569            run_id_suffix=label,
570        )
571        if (
572            mark_swe_atlas_parts
573            and benchmark_id in benchmarks.SWE_ATLAS_PART_BENCHMARKS.values()
574        ):
575            run_request["suite_part"] = label
576        # Only the first run carries the build_request; launch_prepared_run
577        # dedups the actual spawn via spawned_builds anyway.
578        prepared.append((label, run_request, build_request if index == 0 else None))
579    return prepared
580
581
582def prepare_benchmark_runs(
583    args: argparse.Namespace,
584) -> list[tuple[str, dict[str, Any], dict[str, Any] | None]]:
585    benchmark_ids = benchmarks.resolve_benchmarks(
586        list(getattr(args, "benchmark", None) or [])
587    )
588    return prepare_runs_for_benchmarks(
589        args,
590        benchmark_ids,
591        suite_id=mint_suite_id(args, len(benchmark_ids)),
592        label_for_benchmark=lambda benchmark_id: benchmark_id,
593    )
594
595
596def command_run(args: argparse.Namespace) -> int:
597    return execute_prepared_runs(
598        args,
599        prepare_benchmark_runs(args),
600        benchmark_plan_entry,
601    )
602
603
604def resolve_suite_parts(args: argparse.Namespace) -> list[str]:
605    parts = parse_parts([args.parts] if getattr(args, "parts", None) else [])
606    if parts:
607        return parts
608    raise ValueError(
609        "choose at least one SWE-Atlas part with --parts (e.g. --parts rf,qna or --parts all)"
610    )
611
612
613def suite_plan_entry(
614    part: str, run_request: dict[str, Any], build_request: dict[str, Any] | None
615) -> dict[str, Any]:
616    return {
617        "part": part,
618        **benchmark_plan_entry(
619            run_request["benchmark"]["id"], run_request, build_request
620        ),
621    }
622
623
624def suite_entry_label(benchmark_id: str) -> str:
625    for part, part_benchmark_id in benchmarks.SWE_ATLAS_PART_BENCHMARKS.items():
626        if part_benchmark_id == benchmark_id:
627            return part
628    return benchmark_id
629
630
631def prepare_benchmark_suite(
632    args: argparse.Namespace, selectors: list[str]
633) -> list[tuple[str, dict[str, Any], dict[str, Any] | None]]:
634    benchmark_ids = benchmarks.resolve_benchmarks(selectors)
635    if not benchmark_ids:
636        raise ValueError("choose at least one benchmark to run")
637    if getattr(args, "zed_version", None):
638        args.clean_source = True
639        args.require_clean = True
640    timestamp = utc_timestamp()
641    prefix_seed = args.run_id_prefix or args.experiment_prefix or "swe-atlas"
642    if getattr(args, "zed_version", None):
643        prefix_seed = f"{prefix_seed}-{source.sanitize_namespace(args.zed_version)}"
644    suite_id = args.suite_id or f"{source.sanitize_namespace(prefix_seed)}-{timestamp}"
645
646    return prepare_runs_for_benchmarks(
647        args,
648        benchmark_ids,
649        suite_id=suite_id,
650        label_for_benchmark=suite_entry_label,
651        mark_swe_atlas_parts=True,
652    )
653
654
655def prepare_suite(
656    args: argparse.Namespace,
657) -> list[tuple[str, dict[str, Any], dict[str, Any] | None]]:
658    return prepare_benchmark_suite(args, resolve_suite_parts(args))
659
660
661def command_swe_atlas(args: argparse.Namespace) -> int:
662    from .interactive import configure_interactive_suite
663
664    configure_interactive_suite(args)
665    prepared = (
666        prepare_benchmark_suite(args, args.benchmark)
667        if getattr(args, "benchmark", None)
668        else prepare_suite(args)
669    )
670    return execute_prepared_runs(args, prepared, suite_plan_entry)
671