OpenMOSS-Team/SWE-bench-Science

1#!/usr/bin/env python32"""Pull prebuilt task images and run an explicit local selection with Pier."""3 4from __future__ import annotations5 6import argparse7import hashlib8import json9import os10import shlex11import subprocess12import sys13from pathlib import Path14 15try:16    from .provider_config import parse_dotenv, render_codex_config, resolve_codex_profile17except ImportError:  # Direct execution: python3 scripts/run_batch.py18    from provider_config import parse_dotenv, render_codex_config, resolve_codex_profile19 20try:21    from .summarize_results import write_summary22except ImportError:  # Direct execution: python3 scripts/run_batch.py23    from summarize_results import write_summary24 25try:26    import tomllib27except ModuleNotFoundError:  # Python 3.10 and earlier28    try:29        import tomli as tomllib30    except ModuleNotFoundError as exc:  # pragma: no cover - depends on host Python31        raise SystemExit(32            "run_batch.py requires Python 3.11+ or the backport: "33            "python3 -m pip install tomli"34        ) from exc35 36 37def task_dirs(root: Path) -> list[Path]:38    if root.name.startswith("task_") and (root / "task.toml").is_file():39        return [root]40    return sorted(41        (path for path in root.glob("task_*") if path.is_dir()),42        key=lambda path: path.name,43    )44 45 46def validate_artifact_hooks(task_dirs_: list[Path]) -> None:47    missing = [48        task_dir.name49        for task_dir in task_dirs_50        if not (task_dir / "pre_artifacts.sh").is_file()51    ]52    if missing:53        raise ValueError(54            "task bundles are missing pre_artifacts.sh; rematerialize them with "55            "the current release tools: " + ", ".join(missing)56        )57 58 59def load_image_refs(task_dir: Path) -> list[str]:60    config = tomllib.loads((task_dir / "task.toml").read_text(encoding="utf-8"))61    refs = [62        config.get("environment", {}).get("docker_image", ""),63        config.get("verifier", {}).get("environment", {}).get("docker_image", ""),64    ]65    missing = [ref for ref in refs if not ref or "pending" in ref]66    if missing:67        raise ValueError(f"{task_dir.name} has unpublished image references")68    return list(dict.fromkeys(refs))69 70 71def pull_images(task_dirs_: list[Path], *, platform: str) -> list[str]:72    refs: list[str] = []73    for task_dir in task_dirs_:74        refs.extend(load_image_refs(task_dir))75    refs = list(dict.fromkeys(refs))76    for ref in refs:77        command = ["docker", "pull", "--platform", platform, ref]78        print("+ " + shlex.join(command), flush=True)79        subprocess.run(command, check=True)80    return refs81 82 83def selection_payload(root: Path, dirs: list[Path]) -> dict[str, object]:84    selection_file = root / "selection.json"85    if selection_file.is_file():86        payload = json.loads(selection_file.read_text(encoding="utf-8"))87        task_ids = [str(value) for value in payload.get("task_ids", [])]88    else:89        task_ids = [path.name.removeprefix("task_") for path in dirs]90        payload = {91            "allow_restricted_licenses": None,92            "task_ids": task_ids,93        }94    canonical = json.dumps({"task_ids": task_ids}, sort_keys=True).encode("utf-8")95    payload["task_ids"] = task_ids96    payload["selection_sha256"] = hashlib.sha256(canonical).hexdigest()97    return payload98 99 100def redacted_command(command: list[str]) -> str:101    redacted: list[str] = []102    index = 0103    while index < len(command):104        value = command[index]105        if value == "--agent-env" and index + 1 < len(command):106            redacted.extend([value, "<redacted>"])107            index += 2108            continue109        if (110            value == "--agent-kwarg"111            and index + 1 < len(command)112            and command[index + 1].startswith("config_toml=")113        ):114            redacted.extend([value, "config_toml=<provider-config>"])115            index += 2116            continue117        if "=" in value:118            key = value.split("=", 1)[0].lower()119            if any(marker in key for marker in ("key", "token", "secret", "password", "authorization")):120                redacted.append(key + "=<redacted>")121                index += 1122                continue123        redacted.append(value)124        index += 1125    return shlex.join(redacted)126 127 128def pier_version(pier_bin: str) -> str | None:129    try:130        completed = subprocess.run(131            [pier_bin, "--version"], capture_output=True, text=True, check=False132        )133    except OSError:134        return None135    value = (completed.stdout or completed.stderr).strip()136    return value or None137 138 139def main() -> int:140    parser = argparse.ArgumentParser(141        description=__doc__,142        epilog="Use docs/run-batch.md for provider profiles, gateway routing, and result paths.",143    )144    parser.add_argument("--path", type=Path, required=True, help="Materialized task directory")145    parser.add_argument("--agent", default="nop", help="Pier harness, for example codex, claude-code, mini-swe-agent, or nop")146    parser.add_argument("--env", default="docker", help="Pier environment backend")147    parser.add_argument("--env-file", type=Path, help="Provider/harness dotenv file")148    parser.add_argument("--model", action="append", default=[], help="Model route; repeatable")149    parser.add_argument("--agent-env", action="append", default=[], help="Extra harness environment KEY=VALUE; repeatable")150    parser.add_argument("--agent-kwarg", action="append", default=[], help="Extra Pier agent keyword KEY=VALUE; repeatable")151    parser.add_argument("--n-concurrent", type=int, default=1, help="Simultaneous tasks (default: 1)")152    parser.add_argument("--n-attempts", type=int, default=1, help="Attempts per task (default: 1)")153    parser.add_argument("--max-retries", type=int, default=0, help="Retries after attempt-level failure (default: 0)")154    parser.add_argument("--agent-timeout-multiplier", type=float, help="Multiplier for the agent-stage timeout")155    parser.add_argument("--verifier-timeout-multiplier", type=float, help="Multiplier for verifier/build timeouts")156    parser.add_argument("--jobs-dir", type=Path, default=Path("jobs"), help="Pier jobs and summary directory")157    parser.add_argument("--job-name", help="Stable name used in result paths")158    parser.add_argument("--platform", default="linux/amd64", help="Docker platform (default: linux/amd64)")159    parser.add_argument("--pier-bin", default="pier", help="Pier executable or absolute path")160    parser.add_argument("--agent-import-path", help="Explicit Pier agent import path")161    parser.add_argument("--skip-pull", action="store_true", help="Skip Docker pulls for refs already present locally")162    parser.add_argument(163        "--no-auto-provider",164        action="store_true",165        help="Do not translate CODEX_* values from --env-file into native Pier kwargs",166    )167    parser.add_argument(168        "--no-auto-agent-adapter",169        action="store_true",170        help="Use Pier's built-in agent class instead of the runtime-only Codex adapter",171    )172    parser.add_argument("--dry-run", action="store_true", help="Validate, pull, and record metadata without invoking Pier")173    args = parser.parse_args()174 175    root = args.path.resolve()176    dirs = task_dirs(root)177    if not dirs:178        raise ValueError(f"no task_NNN directories found under {root}")179    validate_artifact_hooks(dirs)180    selection = selection_payload(root, dirs)181    image_refs: list[str] = []182    for task_dir in dirs:183        image_refs.extend(load_image_refs(task_dir))184    image_refs = list(dict.fromkeys(image_refs))185    if not args.skip_pull:186        pull_images(dirs, platform=args.platform)187    models = list(args.model)188    agent_kwargs = list(args.agent_kwarg)189    agent_import_path = args.agent_import_path190    if args.agent == "codex" and not args.no_auto_agent_adapter and not agent_import_path:191        package = Path(__file__).resolve().parent.name192        agent_import_path = f"{package}.pier_adapters:ScienceBenchCodex"193    provider_metadata: dict[str, object] | None = None194    if args.agent == "codex" and not args.no_auto_provider:195        profile_env = dict(os.environ)196        if args.env_file:197            profile_env.update(parse_dotenv(args.env_file))198        profile = resolve_codex_profile(profile_env)199        if not models:200            models.append(profile.model)201        if not any(value.startswith(("config_toml=", "config_toml_file=")) for value in agent_kwargs):202            agent_kwargs.append("config_toml=" + render_codex_config(profile))203        # Pier normally strips a provider prefix before invoking Codex. Gateways204        # may use that prefix for routing, so preserve the exact model identifier205        # unless the caller supplied an explicit command override.206        if not any(value.startswith("command_model_name=") for value in agent_kwargs):207            agent_kwargs.append("command_model_name=" + profile.model)208        if profile.version and not any(value.startswith("version=") for value in agent_kwargs):209            agent_kwargs.append("version=" + profile.version)210        if profile.reasoning_effort and not any(211            value.startswith("reasoning_effort=") for value in agent_kwargs212        ):213            agent_kwargs.append("reasoning_effort=" + profile.reasoning_effort)214        provider_metadata = {215            "protocol": profile.wire_api,216            "base_url": profile.safe_base_url,217            "credential_env": "OPENAI_API_KEY",218        }219 220    command = [221        args.pier_bin, "run", "--path", str(root), "--agent", args.agent, "--env", args.env,222        "--n-concurrent", str(args.n_concurrent), "--n-attempts", str(args.n_attempts),223        "--max-retries", str(args.max_retries), "--no-force-build", "--no-delete", "--yes",224    ]225    if args.agent_timeout_multiplier is not None:226        command.extend(["--agent-timeout-multiplier", str(args.agent_timeout_multiplier)])227    if args.verifier_timeout_multiplier is not None:228        command.extend(["--verifier-timeout-multiplier", str(args.verifier_timeout_multiplier)])229    if args.env_file:230        command.extend(["--env-file", str(args.env_file)])231    if agent_import_path:232        command.extend(["--agent-import-path", agent_import_path])233    for model in models:234        command.extend(["--model", model])235    for value in args.agent_env:236        command.extend(["--agent-env", value])237    for value in agent_kwargs:238        command.extend(["--agent-kwarg", value])239    if args.jobs_dir:240        command.extend(["--jobs-dir", str(args.jobs_dir)])241    if args.job_name:242        command.extend(["--job-name", args.job_name])243 244    metadata = {245        "task_ids": selection["task_ids"],246        "allow_restricted_licenses": selection.get("allow_restricted_licenses"),247        "selection_sha256": selection["selection_sha256"],248        "image_refs": image_refs,249        "platform": args.platform,250        "agent": args.agent,251        "models": models,252        "n_concurrent": args.n_concurrent,253        "n_attempts": args.n_attempts,254        "max_retries": args.max_retries,255        "pier_version": pier_version(args.pier_bin),256        "pier_command": redacted_command(command),257        "agent_import_path": agent_import_path,258        "provider": provider_metadata,259    }260    metadata_path = root / "batch-run.json"261    metadata_path.write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")262    print(json.dumps(metadata, indent=2, sort_keys=True))263    if args.dry_run:264        return 0265    pier_environment = os.environ.copy()266    # Pier may build an ephemeral environment+agent image. Keep that derived267    # build on the same architecture as the prebuilt task images.268    pier_environment["DOCKER_DEFAULT_PLATFORM"] = args.platform269    tool_root = str(Path(__file__).resolve().parent.parent)270    existing_pythonpath = pier_environment.get("PYTHONPATH", "")271    pier_environment["PYTHONPATH"] = os.pathsep.join(272        value for value in (tool_root, existing_pythonpath) if value273    )274    returncode = subprocess.run(command, check=False, env=pier_environment).returncode275    try:276        summary_json, summary_csv = write_summary(args.jobs_dir)277        print(json.dumps({"summary_json": str(summary_json), "summary_csv": str(summary_csv)}, indent=2))278    except (OSError, ValueError) as exc:279        print(f"warning: unable to write result summary: {exc}", file=sys.stderr)280    return returncode281 282 283if __name__ == "__main__":284    try:285        raise SystemExit(main())286    except (FileNotFoundError, ValueError, tomllib.TOMLDecodeError) as exc:287        print(f"error: {exc}", file=sys.stderr)288        raise SystemExit(2)289