670 lines · 23.4 KB · Python
RawDownload
from __future__ import annotations import argparseimport jsonimport sysimport uuidfrom pathlib import Pathfrom typing import Any from . import benchmarks, config, harness_command, run_index, sourcefrom .builds import prepare_build_request, resolve_source, validate_build_idfrom .common import (    dedupe_preserving_order,    default_namespace,    deployed_function,    modal_call_id,    parse_parts,    print_json,    utc_now,    utc_timestamp,)from .volume import build_ready_on_volume # Local scratch dir used only for dry-run/plan previews of harness commands.PREVIEW_JOBS_DIR = "/tmp/agent-evals/harbor-jobs"  def read_task_file(path: str | None) -> list[str]:    if not path:        return []    return [        task_name        for line in Path(path).read_text().splitlines()        if (task_name := line.strip()) and not task_name.startswith("#")    ]  def baseten_provider_json(    *,    model_id: str,    api_url: str,    display_name: str | None,    max_tokens: int,    max_output_tokens: int,) -> str:    provider = {        config.BASETEN_PROVIDER_ID: {            "api_url": api_url,            "available_models": [                {                    "name": model_id,                    "display_name": display_name or model_id,                    "max_tokens": max_tokens,                    "max_output_tokens": max_output_tokens,                    "capabilities": {                        "tools": True,                        "images": False,                        "parallel_tool_calls": False,                        "prompt_cache_key": False,                    },                }            ],        }    }    return json.dumps(provider, separators=(",", ":"))  def resolve_model_preset(model: str) -> str:    resolved = config.resolve_model_preset(model)    if resolved != model:        return resolved    if model.startswith("baseten:"):        return f"{config.BASETEN_PROVIDER_ID}/{model.split(':', 1)[1]}"    return model  def resolve_model_options(    args: argparse.Namespace,) -> tuple[str, str | None, list[str]]:    raw_model = getattr(args, "model", None) or config.DEFAULT_MODEL    model = resolve_model_preset(raw_model)    openai_compatible_provider_json = getattr(        args, "openai_compatible_provider_json", None    )    extra_api_secrets = list(getattr(args, "extra_api_secret", None) or [])     if getattr(args, "model_provider", None) == "baseten":        baseten_model = getattr(args, "baseten_model", None) or model        if baseten_model.startswith(f"{config.BASETEN_PROVIDER_ID}/"):            baseten_model = baseten_model.split("/", 1)[1]        if baseten_model.startswith("baseten:"):            baseten_model = baseten_model.split(":", 1)[1]        model = f"{config.BASETEN_PROVIDER_ID}/{baseten_model}"     if model.startswith(f"{config.BASETEN_PROVIDER_ID}/"):        baseten_model = model.split("/", 1)[1]        if not openai_compatible_provider_json:            openai_compatible_provider_json = baseten_provider_json(                model_id=baseten_model,                api_url=getattr(args, "baseten_api_url", config.BASETEN_API_URL),                display_name=getattr(args, "baseten_model_display_name", None),                max_tokens=getattr(                    args, "baseten_model_max_tokens", config.BASETEN_DEFAULT_MAX_TOKENS                ),                max_output_tokens=getattr(                    args,                    "baseten_model_max_output_tokens",                    config.BASETEN_DEFAULT_MAX_OUTPUT_TOKENS,                ),            )     return model, openai_compatible_provider_json, extra_api_secrets  def derive_run_id(    base_run_id: str | None, suite_id: str | None, suffix: str, index: int) -> str:    """Pick a run id for one leg of a (possibly multi-target) launch.     An explicit --run-id is used verbatim for the first leg and suffixed for the    rest; otherwise legs hang off the suite id, falling back to a timestamped id    for a lone run.    """    if base_run_id:        return base_run_id if index == 0 else f"{base_run_id}-{suffix}"    if suite_id:        return f"{suite_id}-{suffix}-{uuid.uuid4().hex[:6]}"    return f"{utc_timestamp()}-{uuid.uuid4().hex[:6]}"  def mint_suite_id(args: argparse.Namespace, run_count: int) -> str | None:    """A suite id groups multiple runs from one invocation; None for a lone run."""    if run_count <= 1:        return None    explicit = getattr(args, "suite_id", None)    if explicit:        return explicit    prefix = getattr(args, "run_id", None) or "run"    return f"{source.sanitize_namespace(prefix)}-{utc_timestamp()}"  def build_rejudge_request(args: argparse.Namespace) -> dict[str, Any]:    """Build the request for re-grading an existing run with a different judge.     The positional `run_id` is the *parent* run; the derived run gets a new id    under the same experiment so `report`/`list` group them together. Only the    judge differs — no build, no agent model, no task selection.    """    judge_preset = args.judge    config.get_judge(judge_preset)  # validate before spawning anything remote    if getattr(args, "experiment_name", None):        experiment_name = source.sanitize_namespace(args.experiment_name)        namespace = default_namespace(args)    else:        entry = run_index.lookup(args.run_id)        if not entry:            raise ValueError(                f"could not locate run '{args.run_id}' in the local run index "                f"({run_index.index_path()}). Pass --experiment-name (and "                "--namespace if it isn't yours)."            )        experiment_name = source.sanitize_namespace(entry["experiment_name"])        namespace = source.sanitize_namespace(            getattr(args, "namespace", None) or entry["namespace"]        )    parent_namespace = (        source.sanitize_namespace(args.parent_namespace)        if getattr(args, "parent_namespace", None)        else namespace    )    parent_run_id = args.run_id    judge_slug = source.sanitize_namespace(judge_preset)    new_run_id = (        args.new_run_id        or f"{parent_run_id}-rejudge-{judge_slug}-{uuid.uuid4().hex[:6]}"    )    return {        "namespace": namespace,        "experiment_name": experiment_name,        "run_id": new_run_id,        "parent": {            "namespace": parent_namespace,            "experiment_name": experiment_name,            "run_id": parent_run_id,        },        "judge_preset": judge_preset,        "judge_model": getattr(args, "judge_model", None),        "volume_name": args.volume,        "api_secret_name": args.api_secret,        "created_at": utc_now(),    }  def command_rejudge(args: argparse.Namespace) -> int:    rejudge_request = build_rejudge_request(args)    if getattr(args, "dry_run", False) or getattr(args, "plan", False):        print_json(rejudge_request)        return 0     controller = deployed_function(args, "rejudge_controller")    call = controller.spawn(rejudge_request)    run_index.record_run(        {**rejudge_request, "volume_name": args.volume, "kind": "rejudge"}    )    parent = rejudge_request["parent"]    print(f"Namespace:  {rejudge_request['namespace']}")    print(f"Experiment: {rejudge_request['experiment_name']}")    print(        f"Source run: {parent['namespace']}/{parent['experiment_name']}/"        f"{parent['run_id']}"    )    print(f"New run id: {rejudge_request['run_id']}")    print(f"Judge:      {rejudge_request['judge_preset']}")    print(f"Spawned rejudge controller: {modal_call_id(call)}")    print("\nNext steps (run id alone is enough):")    print(f"  zed-eval status {rejudge_request['run_id']}")    print(f"  zed-eval report {rejudge_request['run_id']} --fetch")    return 0  def common_run_request_fields(    args: argparse.Namespace,    *,    namespace: str,    run_id: str,    experiment_name: str,    judge_preset: str | None,    build_id: str | None,    suite_id: str | None,) -> dict[str, Any]:    """Fields shared by every benchmark run request."""    agent_model, openai_compatible_provider_json, extra_api_secrets = (        resolve_model_options(args)    )    task_names = dedupe_preserving_order(        read_task_file(getattr(args, "tasks", None))        + (getattr(args, "include_task_name", None) or [])    )    return {        "created_at": utc_now(),        "namespace": namespace,        "run_id": run_id,        "experiment_name": experiment_name,        "volume_name": args.volume,        "api_secret_name": args.api_secret,        "modal_token_secret_name": args.modal_token_secret,        "orchestration": config.orchestration_info(),        "agent_model": agent_model,        "judge_preset": judge_preset,        "judge_model": getattr(args, "judge_model", None),        "build_id": build_id,        "task_names": task_names,        "n_tasks": getattr(args, "n_tasks", None),        "n_concurrent": getattr(args, "n_concurrent", config.DEFAULT_N_CONCURRENT),        "override_cpus": getattr(args, "override_cpus", config.DEFAULT_OVERRIDE_CPUS),        "override_memory_mb": getattr(            args, "override_memory_mb", config.DEFAULT_OVERRIDE_MEMORY_MB        ),        "sandbox_timeout_secs": getattr(            args, "sandbox_timeout_secs", config.DEFAULT_SANDBOX_TIMEOUT_SECS        ),        "sandbox_idle_timeout_secs": getattr(            args, "sandbox_idle_timeout_secs", config.DEFAULT_SANDBOX_IDLE_TIMEOUT_SECS        ),        "build_wait_timeout_secs": getattr(args, "build_wait_timeout_secs", 7200),        "extra_harbor_args": list(getattr(args, "extra_harbor_arg", None) or []),        "openai_compatible_provider_json": openai_compatible_provider_json,        "anthropic_available_models_json": getattr(            args, "anthropic_available_models_json", None        ),        "extra_api_secrets": dedupe_preserving_order(extra_api_secrets),        "suite_id": suite_id,    }  def benchmark_plan_entry(    benchmark_id: str, run_request: dict[str, Any], build_request: dict[str, Any] | None) -> dict[str, Any]:    return {        "benchmark": benchmark_id,        "run_id": run_request["run_id"],        "harness": run_request["benchmark"]["harness"],        "model": run_request["agent_model"],        "judge": run_request.get("judge_preset"),        "build_id": run_request.get("build_id"),        "will_build": build_request is not None,        "n_tasks": run_request.get("n_tasks"),    }  def print_plan(    summaries: list[dict[str, Any]],    details: list[tuple[str, dict[str, Any], dict[str, Any] | None]],    *,    verbose: bool,) -> None:    print("Plan:")    print_json(summaries)    if verbose:        for header, run_request, build_request in details:            print(f"\n=== {header} ===")            print_dry_run(run_request, build_request)  def print_dry_run(    run_request: dict[str, Any], build_request: dict[str, Any] | None) -> None:    print("Run request:")    print_json(run_request)    if build_request:        print("\nBuild request:")        build_preview = dict(build_request)        patch = build_preview.pop("patch", "")        build_preview["patch_line_count"] = len(patch.splitlines())        build_preview["source"] = source.public_source_info(            build_request.get("source") or {}        )        print_json(build_preview)    command = config.redacted_command(        harness_command.build_harness_command(run_request, PREVIEW_JOBS_DIR)    )    print("\nHarness command:")    print(command)  def execute_prepared_runs(    args: argparse.Namespace,    prepared: list[tuple[str, dict[str, Any], dict[str, Any] | None]],    plan_entry,) -> int:    if args.plan:        print_plan(            [                plan_entry(label, run_request, build_request)                for label, run_request, build_request in prepared            ],            prepared,            verbose=args.verbose,        )        return 0    if args.dry_run:        for label, run_request, build_request in prepared:            print(f"\n=== {label} ===")            print_dry_run(run_request, build_request)        return 0     spawned_builds: set[str] = set()    for label, run_request, build_request in prepared:        print(f"\n=== Launching {label} ===")        launch_prepared_run(args, run_request, build_request, spawned_builds)    return 0  def launch_prepared_run(    args: argparse.Namespace,    run_request: dict[str, Any],    build_request: dict[str, Any] | None,    spawned_builds: set[str] | None = None,) -> None:    build_function = None    if build_request:        run_request["source"] = source.public_source_info(build_request["source"])        print_untracked_warning(build_request)        build_function = deployed_function(args, "build_eval_cli")    record_function = deployed_function(args, "create_run_record")    controller_function = deployed_function(args, "run_controller")     record_state = record_function.remote(run_request)    run_index.record_run(run_request)     print(f"Namespace:  {run_request['namespace']}")    print(f"Experiment: {run_request['experiment_name']}")    print(f"Run id:     {run_request['run_id']}")    print(f"Volume:     {run_request['volume_name']}")    print(f"Run state:  {record_state['status']}")    print(f"Model:      {run_request['agent_model']}")    if run_request.get("judge_preset"):        print(f"Judge:      {run_request['judge_preset']}")    if run_request.get("suite_id"):        print(f"Suite:      {run_request['suite_id']}")    if run_request.get("build_id"):        print(f"Build id:  {run_request['build_id']}")    if run_request.get("task_names"):        print(f"Tasks:     {len(run_request['task_names'])} explicit task(s)")    elif run_request.get("n_tasks"):        print(f"Tasks:     Harbor --n-tasks {run_request['n_tasks']}")    else:        print("Tasks:     full dataset selection")     build_id = run_request.get("build_id")    if build_function is not None and build_id not in (spawned_builds or set()):        build_call = build_function.spawn(build_request)        print(f"Spawned build:      {modal_call_id(build_call)}")        if spawned_builds is not None and build_id:            spawned_builds.add(build_id)    controller_call = controller_function.spawn(run_request)    print(f"Spawned controller: {modal_call_id(controller_call)}")     run_id = run_request["run_id"]    print(        "\nNext steps (run id alone is enough; namespace/experiment are resolved "        "from this machine's local run index):"    )    print(f"  zed-eval status {run_id}")    print(f"  zed-eval logs {run_id}")    print(f"  zed-eval report {run_id} --fetch")  def resolve_benchmark_judge(    args: argparse.Namespace, benchmark: benchmarks.Benchmark) -> str | None:    if not benchmark.needs_judge:        return None    judge = getattr(args, "judge", None) or config.DEFAULT_JUDGE_PRESET    if judge == "auto":        return benchmark.default_judge or "leaderboard"    config.get_judge(judge)    return judge  def benchmark_metadata_for_run(    args: argparse.Namespace, benchmark: benchmarks.Benchmark) -> dict[str, object]:    metadata = benchmarks.benchmark_metadata(benchmark)    dataset = metadata.get("dataset")    if (        isinstance(dataset, dict)        and dataset.get("repo_url") == benchmarks.SWE_ATLAS_REPO_URL    ):        dataset["repo_url"] = (            getattr(args, "swe_atlas_repo_url", None) or benchmarks.SWE_ATLAS_REPO_URL        )        dataset["repo_ref"] = (            getattr(args, "swe_atlas_repo_ref", None) or benchmarks.SWE_ATLAS_REPO_REF        )    return metadata  def print_untracked_warning(build_request: dict[str, Any]) -> None:    build_source = build_request.get("source") or {}    untracked_files = build_source.get("untracked_files") or []    if untracked_files:        print(            f"Warning: proceeding with {len(untracked_files)} untracked file(s) "            "not included in the build patch.",            file=sys.stderr,        )  def prepare_shared_build(    args: argparse.Namespace,) -> tuple[str | None, dict[str, Any] | None]:    """Resolve the build once for a whole (possibly multi-benchmark) run.     Returns `(build_id, build_request)`. `build_request` is None when the target    build already exists; otherwise the caller should spawn the build.    """    explicit_build_id = getattr(args, "build", None)    validate_build_id(explicit_build_id)    if (        explicit_build_id        and not getattr(args, "plan", False)        and not getattr(args, "dry_run", False)        and build_ready_on_volume(args, explicit_build_id)    ):        print(f"Reusing existing build {explicit_build_id} (already on volume)")        return explicit_build_id, None     base_sha, clean_source, source_label, pre_resolved = resolve_source(args)     build_request = prepare_build_request(        base_sha=base_sha,        patch_path=getattr(args, "patch_path", None),        build_id=explicit_build_id,        allow_untracked=getattr(args, "allow_untracked", False),        require_clean=getattr(args, "require_clean", False),        repo_url=getattr(args, "repo_url", None),        clean_source=clean_source,        source_label=source_label,        pre_resolved_base_sha=pre_resolved,    )    build_id = build_request["build_id"]     if not getattr(args, "plan", False) and not getattr(args, "dry_run", False):        if build_ready_on_volume(args, build_id):            print(f"Reusing existing build {build_id} (already on volume)")            return build_id, None     return build_id, build_request  def command_build(args: argparse.Namespace) -> int:    build_id, build_request = prepare_shared_build(args)    if build_request is None:        return 0     print(f"Build id: {build_id}")    print(f"Base sha: {build_request['base_sha']}")    print(f"Patch sha256: {build_request['patch_sha256'] or '(none)'}")    print_untracked_warning(build_request)

Lines 1–500 of 670.