"""Orchestrates benchmark runs for SWE-bench instances."""
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
import structlog
from benchmark.server import ServerProcess
from benchmark.swebench.evaluator import (
EvaluationError,
EvaluationResult,
Evaluator,
)
from benchmark.swebench.generator import BenchmarkRunner, PatchGenerator
from benchmark.swebench.instance import SWEBenchInstance
from benchmark.swebench.prediction import Prediction, PredictionWriter
from benchmark.swebench.repository import RepositoryPreparer
from benchmark.swebench.retriever import KoditRetriever
@dataclass(frozen=True)
class ComparisonResult:
"""Result comparing baseline vs kodit-augmented predictions."""
instance_id: str
baseline_prediction: Prediction
kodit_prediction: Prediction
baseline_evaluation: EvaluationResult | None
kodit_evaluation: EvaluationResult | None
@property
def baseline_resolved(self) -> bool:
"""Whether baseline prediction resolved the issue."""
if self.baseline_evaluation is None:
return False
return self.baseline_evaluation.resolved > 0
@property
def kodit_resolved(self) -> bool:
"""Whether kodit prediction resolved the issue."""
if self.kodit_evaluation is None:
return False
return self.kodit_evaluation.resolved > 0
def as_dict(self) -> dict:
"""Convert to dictionary for JSON serialization."""
return {
"instance_id": self.instance_id,
"baseline_resolved": self.baseline_resolved,
"kodit_resolved": self.kodit_resolved,
"baseline_prediction": {
"instance_id": self.baseline_prediction.instance_id,
"model_name_or_path": self.baseline_prediction.model_name_or_path,
"patch_length": len(self.baseline_prediction.model_patch),
},
"kodit_prediction": {
"instance_id": self.kodit_prediction.instance_id,
"model_name_or_path": self.kodit_prediction.model_name_or_path,
"patch_length": len(self.kodit_prediction.model_patch),
},
"baseline_evaluation": (
self.baseline_evaluation.as_dict()
if self.baseline_evaluation
else None
),
"kodit_evaluation": (
self.kodit_evaluation.as_dict() if self.kodit_evaluation else None
),
}
class BenchmarkOperations:
"""Reusable benchmark operations for an already-running Kodit server."""
def __init__( # noqa: PLR0913
self,
kodit_base_url: str,
repos_dir: Path,
results_dir: Path,
model: str,
api_key: str | None = None,
top_k: int = 10,
llm_timeout: int = 60,
) -> None:
"""Initialize operations with server URL and paths."""
self._base_url = kodit_base_url
self._repos_dir = repos_dir
self._results_dir = results_dir
self._model = model
self._api_key = api_key
self._top_k = top_k
self._llm_timeout = llm_timeout
self._log = structlog.get_logger(__name__)
def prepare(self, instance: SWEBenchInstance) -> int:
"""Prepare an instance by cloning and indexing."""
self._log.info("Preparing instance", instance_id=instance.instance_id)
preparer = RepositoryPreparer(
kodit_base_url=self._base_url,
repos_dir=self._repos_dir,
)
return preparer.prepare(instance)
def generate(
self,
instance: SWEBenchInstance,
condition: str,
) -> Prediction:
"""Generate a prediction for baseline or kodit condition."""
generator = PatchGenerator(
model=self._model,
api_key=self._api_key,
timeout=self._llm_timeout,
)
if condition == "baseline":
self._log.info("Generating baseline prediction")
return generator.generate_baseline(instance)
retriever = KoditRetriever(kodit_base_url=self._base_url)
runner = BenchmarkRunner(generator=generator, retriever=retriever)
self._log.info("Generating kodit prediction", top_k=self._top_k)
return runner.run_kodit(instance, top_k=self._top_k)
def generate_and_write(
self,
instance: SWEBenchInstance,
condition: str,
output_path: Path | None = None,
) -> tuple[Prediction, Path]:
"""Generate a prediction and write to file."""
prediction = self.generate(instance, condition)
if output_path is None:
self._results_dir.mkdir(parents=True, exist_ok=True)
output_path = self._results_dir / f"{condition}.jsonl"
PredictionWriter(output_path).write(prediction)
self._log.info("Prediction written", output=str(output_path))
return prediction, output_path
def evaluate(
self,
predictions_path: Path,
run_id: str,
max_workers: int = 4,
) -> EvaluationResult:
"""Evaluate predictions using SWE-bench harness."""
self._log.info("Evaluating predictions", path=str(predictions_path))
evaluator = Evaluator()
return evaluator.evaluate_full(
predictions_path=predictions_path,
run_id=run_id,
max_workers=max_workers,
)
def run_comparison(
self,
instance: SWEBenchInstance,
*,
skip_evaluation: bool = False,
) -> ComparisonResult:
"""Run both conditions and compare results."""
# Generate predictions
baseline_prediction = self.generate(instance, "baseline")
kodit_prediction = self.generate(instance, "kodit")
# Write predictions
self._results_dir.mkdir(parents=True, exist_ok=True)
baseline_path = self._results_dir / f"{instance.instance_id}.baseline.jsonl"
kodit_path = self._results_dir / f"{instance.instance_id}.kodit.jsonl"
PredictionWriter(baseline_path).write(baseline_prediction)
PredictionWriter(kodit_path).write(kodit_prediction)
self._log.info(
"Predictions written",
baseline=str(baseline_path),
kodit=str(kodit_path),
)
# Evaluate if requested
baseline_evaluation = None
kodit_evaluation = None
if not skip_evaluation:
try:
baseline_evaluation = self.evaluate(
baseline_path,
run_id=f"{instance.instance_id}_baseline",
)
except EvaluationError as e:
self._log.warning("Baseline evaluation failed", error=str(e))
try:
kodit_evaluation = self.evaluate(
kodit_path,
run_id=f"{instance.instance_id}_kodit",
)
except EvaluationError as e:
self._log.warning("Kodit evaluation failed", error=str(e))
return ComparisonResult(
instance_id=instance.instance_id,
baseline_prediction=baseline_prediction,
kodit_prediction=kodit_prediction,
baseline_evaluation=baseline_evaluation,
kodit_evaluation=kodit_evaluation,
)
class InstanceRunner:
"""Runs a complete benchmark with server lifecycle management."""
def __init__(
self,
server: ServerProcess,
operations: BenchmarkOperations,
) -> None:
"""Initialize runner with server and operations."""
self._server = server
self._operations = operations
self._log = structlog.get_logger(__name__)
def run(
self,
instance: SWEBenchInstance,
*,
skip_evaluation: bool = False,
) -> ComparisonResult:
"""Run the full benchmark workflow for an instance."""
self._log.info(
"Starting benchmark run",
instance_id=instance.instance_id,
repo=instance.repo,
)
# Start server
self._log.info("Starting kodit server")
if not self._server.start():
msg = "Failed to start kodit server"
raise InstanceRunError(msg)
try:
# Prepare and run comparison
self._operations.prepare(instance)
return self._operations.run_comparison(
instance,
skip_evaluation=skip_evaluation,
)
finally:
# Stop server
self._log.info("Stopping kodit server")
self._server.stop()
class InstanceRunError(Exception):
"""Raised when an instance run fails."""
@dataclass(frozen=True)
class BatchResult:
"""Result of running benchmarks for multiple instances."""
total: int
succeeded: int
failed: int
failures: list[str]
def as_dict(self) -> dict:
"""Convert to dictionary for JSON serialization."""
return {
"total": self.total,
"succeeded": self.succeeded,
"failed": self.failed,
"failures": self.failures,
}
@dataclass(frozen=True)
class BatchConfig:
"""Configuration for batch benchmark runs."""
repos_dir: Path
results_dir: Path
model: str
api_key: str
top_k: int
skip_evaluation: bool
host: str
port: int
db_port: int
enrichment_base_url: str
enrichment_model: str
enrichment_parallel_tasks: int
enrichment_timeout: int
embedding_base_url: str
embedding_model: str
embedding_parallel_tasks: int
embedding_timeout: int
llm_timeout: int = 60
class BatchRunner:
"""Runs benchmarks for multiple SWE-bench instances sequentially."""
def __init__(self, config: BatchConfig) -> None:
"""Initialize with configuration."""
self._config = config
self._log = structlog.get_logger(__name__)
def run(
self,
instances: list[SWEBenchInstance],
on_instance_complete: Callable[[str, bool, int, int], None] | None = None,
) -> BatchResult:
"""Run benchmarks for all instances sequentially."""
total = len(instances)
self._log.info("Starting batch benchmark run", total=total)
succeeded = 0
failed = 0
failures: list[str] = []
for i, instance in enumerate(instances, start=1):
instance_id = instance.instance_id
self._log.info(
"Running instance",
progress=f"{i}/{total}",
instance_id=instance_id,
)
success = self._run_single(instance)
if success:
succeeded += 1
else:
failed += 1
failures.append(instance_id)
if on_instance_complete:
on_instance_complete(instance_id, success, i, total)
self._log.info(
"Batch benchmark complete",
total=total,
succeeded=succeeded,
failed=failed,
)
return BatchResult(
total=total,
succeeded=succeeded,
failed=failed,
failures=failures,
)
def _run_single(self, instance: SWEBenchInstance) -> bool:
"""Run benchmark for a single instance. Returns True on success."""
server = ServerProcess(
host=self._config.host,
port=self._config.port,
db_port=self._config.db_port,
enrichment_base_url=self._config.enrichment_base_url,
enrichment_model=self._config.enrichment_model,
enrichment_api_key=self._config.api_key,
enrichment_parallel_tasks=self._config.enrichment_parallel_tasks,
enrichment_timeout=self._config.enrichment_timeout,
embedding_base_url=self._config.embedding_base_url,
embedding_model=self._config.embedding_model,
embedding_api_key=self._config.api_key,
embedding_parallel_tasks=self._config.embedding_parallel_tasks,
embedding_timeout=self._config.embedding_timeout,
)
operations = BenchmarkOperations(
kodit_base_url=server.base_url,
repos_dir=self._config.repos_dir,
results_dir=self._config.results_dir,
model=self._config.model,
api_key=self._config.api_key,
top_k=self._config.top_k,
llm_timeout=self._config.llm_timeout,
)
runner = InstanceRunner(server=server, operations=operations)
try:
result = runner.run(instance, skip_evaluation=self._config.skip_evaluation)
except InstanceRunError as e:
self._log.error(
"Benchmark run failed",
instance_id=instance.instance_id,
error=str(e),
)
return False
except TimeoutError as e:
self._log.error(
"Benchmark run timed out",
instance_id=instance.instance_id,
error=str(e),
)
return False
except Exception as e: # noqa: BLE001
self._log.error(
"Unexpected error during benchmark run",
instance_id=instance.instance_id,
error=str(e),
error_type=type(e).__name__,
)
return False
self._write_comparison(instance.instance_id, result)
return True
def _write_comparison(self, instance_id: str, result: ComparisonResult) -> None:
"""Write comparison result to file."""
import json
comparison_path = self._config.results_dir / f"{instance_id}.comparison.json"
comparison_path.parent.mkdir(parents=True, exist_ok=True)
with comparison_path.open("w") as f:
json.dump(result.as_dict(), f, indent=2)
self._log.info(
"Comparison written",
instance_id=instance_id,
path=str(comparison_path),
)
