Files
2026-07-13 12:49:17 +08:00

277 lines
10 KiB
Python

# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import json
import os
import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable
from datasets import load_dataset
from docker.models.containers import Container, ExecResult
def docker_exec(container: Container, command: str):
"""
Execute a shell command inside a Docker container.
Args:
container: Docker container object.
command: Shell command to execute.
Returns:
Tuple (return_code, output_str).
"""
exec_result: ExecResult = container.exec_run(cmd=command)
return_code = exec_result[0]
output = exec_result[1].decode("utf-8")
return return_code, output
def swebench_evaluate_harness_after(benchmark_harness_path, task_id):
src_base = f"{benchmark_harness_path}/logs/run_evaluation/{task_id}/trae-agent"
dst_base = f"results/{task_id}"
json_src = f"{benchmark_harness_path}/trae-agent.{task_id}.json"
json_dst = os.path.join(dst_base, "results.json")
if not os.path.exists(src_base):
print(f"Source directory does not exist: {src_base}")
return
for folder_name in os.listdir(src_base):
src_folder = os.path.join(src_base, folder_name)
dst_folder = os.path.join(dst_base, folder_name)
if os.path.isdir(src_folder):
os.makedirs(dst_folder, exist_ok=True)
for file_name in os.listdir(src_folder):
src_file = os.path.join(src_folder, file_name)
dst_file = os.path.join(dst_folder, file_name)
if not os.path.exists(dst_file):
shutil.copy2(src_file, dst_file)
os.makedirs(dst_base, exist_ok=True)
if not os.path.exists(json_dst):
shutil.copy2(json_src, json_dst)
def multi_swebench_evaluate_harness_after(benchmark_harness_path, task_id):
task_results_dir = Path("results") / task_id
output_dir = (task_results_dir / "dataset").resolve()
src_file = output_dir / "final_report.json"
dst_file = task_results_dir / "results.json"
if not src_file.exists():
raise FileNotFoundError(f"{src_file} not found")
shutil.copyfile(src_file, dst_file)
def _write_problem_statement(instance_dir: Path, content: str) -> int:
"""Helper function to write problem statement using context manager."""
with open(instance_dir / "problem_statement.txt", "w", encoding="utf-8") as f:
return f.write(content)
def _load_jsonl_dataset(dataset_name: str) -> list[dict]:
"""Helper function to load JSONL dataset using context manager."""
result = []
with open(f"{dataset_name.lower().replace('-', '_')}.jsonl", "r", encoding="utf-8") as f:
for line in f:
if line.strip():
result.append(json.loads(line))
return result
def _write_multi_problem_statement(instance_dir: Path, resolved_issues: list[dict]) -> int:
"""Helper function to write multi-issue problem statement using context manager."""
content = "\n".join(
issue.get("title", "") + "\n" + issue.get("body", "") for issue in resolved_issues
)
with open(instance_dir / "problem_statement.txt", "w", encoding="utf-8") as f:
return f.write(content)
def multi_swebench_evaluate_harness_before(task_results_dir, dataset_name, max_workers):
task_results_dir = Path(task_results_dir)
pred_json_path = task_results_dir / "predictions.json"
pred_jsonl_path = task_results_dir / "predictions.jsonl"
dataset_file_path = f"{dataset_name.lower().replace('-', '_')}.jsonl"
instance_map = {}
with open(dataset_file_path, "r", encoding="utf-8") as f:
for line in f:
if not line.strip():
continue
item = json.loads(line)
instance_id = item.get("instance_id")
org = item.get("org")
repo = item.get("repo")
number = item.get("number")
instance_map[instance_id] = {"org": org, "repo": repo, "number": number}
with open(pred_json_path, "r", encoding="utf-8") as f:
preds = json.load(f)
with open(pred_jsonl_path, "w", encoding="utf-8") as f:
for item in preds:
instance_id = item["instance_id"]
patch = item["model_patch"]
info = instance_map.get(instance_id, {})
new_item = {
"org": info.get("org"),
"repo": info.get("repo"),
"number": info.get("number"),
"fix_patch": patch,
}
f.write(json.dumps(new_item, ensure_ascii=False) + "\n")
base_dir = Path(__file__).resolve().parent
task_results_dir = base_dir / task_results_dir
patch_file_path = str((base_dir / pred_jsonl_path).resolve())
dataset_file_path = str((base_dir / dataset_file_path).resolve())
output_dir = (task_results_dir / "dataset").resolve()
repo_dir = (task_results_dir / "repos").resolve()
log_dir = (task_results_dir / "logs").resolve()
workdir = (task_results_dir / "workdir").resolve()
output_dir.mkdir(parents=True, exist_ok=True)
repo_dir.mkdir(parents=True, exist_ok=True)
log_dir.mkdir(parents=True, exist_ok=True)
workdir.mkdir(parents=True, exist_ok=True)
output_dir = str(output_dir)
repo_dir = str(repo_dir)
log_dir = str(log_dir)
workdir = str(workdir)
config = {
"mode": "evaluation",
"workdir": workdir,
"patch_files": [patch_file_path],
"dataset_files": [dataset_file_path],
"force_build": False,
"output_dir": output_dir,
"specifics": [],
"skips": [],
"repo_dir": repo_dir,
"need_clone": False,
"global_env": [],
"clear_env": True,
"stop_on_error": True,
"max_workers": max_workers,
"max_workers_build_image": max_workers,
"max_workers_run_instance": max_workers,
"log_dir": log_dir,
"log_level": "DEBUG",
}
config_path = task_results_dir / "evaluate_config.json"
with open(config_path, "w", encoding="utf-8") as f:
json.dump(config, f, indent=2)
@dataclass
class BenchmarkConfig:
valid_datasets: list[str]
load_dataset: Callable[[str], Any]
image_name: Callable[[str], str]
problem_statement: Callable[[dict, Path], Any]
working_dir: Callable[[str], str]
evaluate_harness: Callable[..., list[str]]
evaluate_harness_before: Callable[..., Any]
evaluate_harness_after: Callable[..., Any]
BENCHMARK_CONFIG: dict[str, BenchmarkConfig] = {
# SWE-bench
"SWE-bench": BenchmarkConfig(
valid_datasets=["SWE-bench", "SWE-bench_Lite", "SWE-bench_Verified"],
load_dataset=lambda dataset_name: load_dataset(
f"princeton-nlp/{dataset_name}", split="test"
),
image_name=lambda instance_id: (
f"swebench/sweb.eval.x86_64.{instance_id.lower()}:latest".replace("__", "_1776_")
),
problem_statement=lambda instance, instance_dir: (
_write_problem_statement(instance_dir, instance.get("problem_statement", ""))
),
working_dir=lambda instance_id: "/testbed/",
evaluate_harness=lambda dataset_name, task_results_dir, task_id, max_workers: [
"swebench_venv/bin/python",
"-m",
"swebench.harness.run_evaluation",
"--dataset_name",
f"princeton-nlp/{dataset_name}",
"--predictions_path",
(task_results_dir / "predictions.json").absolute().as_posix(),
"--max_workers",
str(max_workers),
"--run_id",
task_id,
"--cache_level",
"instance",
"--instance_image_tag",
"latest",
],
evaluate_harness_before=lambda *args, **kwargs: None,
evaluate_harness_after=swebench_evaluate_harness_after,
),
# SWE-bench-Live
"SWE-bench-Live": BenchmarkConfig(
valid_datasets=["SWE-bench-Live/lite", "SWE-bench-Live/verified", "SWE-bench-Live/full"],
load_dataset=lambda dataset_name: load_dataset(
"SWE-bench-Live/SWE-bench-Live", split=dataset_name.split("/")[-1]
),
image_name=lambda instance_id: (
f"starryzhang/sweb.eval.x86_64.{instance_id.lower()}:latest".replace("__", "_1776_")
),
problem_statement=lambda instance, instance_dir: (
_write_problem_statement(instance_dir, instance.get("problem_statement", ""))
),
working_dir=lambda instance_id: "/testbed/",
evaluate_harness=lambda dataset_name, task_results_dir, task_id, max_workers: [
"swebench_live_venv/bin/python",
"-m",
"swebench.harness.run_evaluation",
"--dataset_name",
"SWE-bench-Live/SWE-bench-Live",
"--namespace",
"starryzhang",
"--split",
dataset_name.split("/")[-1],
"--predictions_path",
(task_results_dir / "predictions.json").absolute().as_posix(),
"--run_id",
task_id,
"--max_workers",
str(max_workers),
],
evaluate_harness_before=lambda *args, **kwargs: None,
evaluate_harness_after=swebench_evaluate_harness_after,
),
# Multi-SWE-bench
"Multi-SWE-bench": BenchmarkConfig(
valid_datasets=["Multi-SWE-bench-flash", "Multi-SWE-bench_mini"],
load_dataset=lambda dataset_name: _load_jsonl_dataset(dataset_name),
image_name=lambda instance_id: (
(lambda key: key.rpartition("-")[0] + ":pr-" + key.rpartition("-")[2])(
f"mswebench/{instance_id.lower()}".replace("__", "_m_")
)
),
problem_statement=lambda instance, instance_dir: (
_write_multi_problem_statement(instance_dir, instance.get("resolved_issues", []))
),
working_dir=lambda instance_id: (
f"/home/{'-'.join(instance_id.split('__')[-1].split('-')[:-1])}/"
),
evaluate_harness=lambda dataset_name, task_results_dir, task_id, max_workers: [
"multi_swebench_venv/bin/python",
"-m",
"multi_swe_bench.harness.run_evaluation",
"--config",
os.path.join(
os.path.dirname(os.path.abspath(__file__)),
task_results_dir / "evaluate_config.json",
),
],
evaluate_harness_before=multi_swebench_evaluate_harness_before,
evaluate_harness_after=multi_swebench_evaluate_harness_after,
),
}