277 lines
10 KiB
Python
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,
|
|
),
|
|
}
|