chore: import upstream snapshot with attribution
Deploy Documentation / deploy (push) Has been cancelled
CPU Test / Test (Utilities, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (LLM proxy, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Others, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Store, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Utilities, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Weave, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (AgentOps, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (LLM proxy, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Others, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Weave, latest, Python 3.13) (push) Has been cancelled
Dashboard / Chromatic (push) Has been cancelled
CPU Test / Lint - fast (push) Has been cancelled
CPU Test / Lint - next (push) Has been cancelled
CPU Test / Lint - slow (push) Has been cancelled
CPU Test / Lint - JavaScript (push) Has been cancelled
CPU Test / Build documentation (push) Has been cancelled
CPU Test / Test (AgentOps, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (LLM proxy, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (Others, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (Store, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (Weave, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (AgentOps, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Store, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Utilities, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Weave, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (AgentOps, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (LLM proxy, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (Others, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (Store, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (Utilities, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (JavaScript) (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 12:44:17 +08:00
commit 85742ab165
588 changed files with 320176 additions and 0 deletions
+112
View File
@@ -0,0 +1,112 @@
# Training Claude Code with Agent-lightning
[![claude-code CI status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-claude-code.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-claude-code.yml)
This example shows how to wrap Anthropic's Claude Code experience with Agent-lightning instrumentation to solve SWE-bench tasks, collect spans/logs, and optionally convert those traces into HuggingFace datasets.
**NOTE:** This example only shows how to integrate Claude Code as an agent in Agent-lightning. The training part is still under development and welcoming contributions!
## Overview
`claude_code_agent.py` spins up a Lightning Store, an LLM proxy, and the Claude Code controller. Each SWE-bench instance is executed inside the official container image so you can either prompt-tune against Anthropic's hosted models or point Claude Code at a self-hosted OpenAI-compatible backend such as vLLM. When a backend surfaces token IDs/logprobs (e.g., vLLM), the traces are turned into triplets that downstream fine-tuning pipelines can consume.
## Requirements
First, install Agent-lightning following the [installation guide](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/). Then install the SWE-bench harness plus utilities used by this example:
```bash
(uv) pip install swebench transformers datasets python-dotenv
```
Docker must be available because each SWE-bench instance is executed in a container via `swebench_utils`.
Finally, set API credentials depending on backend:
- `ANTHROPIC_API_KEY` for the official Claude Code path.
- `OPENAI_API_KEY` (or another OpenAI-compatible key) for the `openai` backend.
- A running OpenAI-compatible server (e.g., vLLM) when using the `vllm` backend.
## Dataset
`swebench_samples.jsonl` contains a handful of SWE-bench issues for smoke testing. For full-scale benchmarks load `princeton-nlp/SWE-bench` via `load_swebench_dataset` or point `--dataset-path` to your own JSONL file.
## Included Files
| File/Directory | Description |
|----------------|-------------|
| `claude_code_agent.py` | CLI entry point that launches the Lightning store, LLM proxy, and Claude Code agent |
| `claude_code_controller.py` | Manages the SWE-bench Docker runtime and translates model outputs into git patches |
| `extended_adapter.py` | Adapter that converts LLM proxy spans into triplets with token IDs, logprobs, and chat history |
| `swebench_samples.jsonl` | Mini SWE-bench subset for quick validation |
| `swebench_utils/` | Utilities for running/evaluating SWE-bench instances inside containers |
| `templates/handle_hook.template.sh` | Helper script injected into containers for hook handling |
| `templates/settings.template.json` | Base configuration consumed by Claude Code CLI |
## Running the Example
All commands are issued from `examples/claude_code`. Inspect the module-level docstring in `claude_code_agent.py` for the full CLI reference.
### Hosted vLLM (open-source models)
First, launch your model behind an OpenAI-compatible endpoint, for example:
```bash
vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
--max-model-len 131072 \
--enable-auto-tool-choice \
--tool-call-parser qwen3_coder
```
Run the Agent-lightning harness and point it at the server:
```bash
python claude_code_agent.py vllm \
--backend-model-high Qwen/Qwen3-Coder-30B-A3B-Instruct \
--backend-model-low Qwen/Qwen3-Coder-30B-A3B-Instruct \
--frontend-model-high claude-sonnet-4-5-20250929 \
--frontend-model-low claude-haiku-4-5-20251001 \
--base-url http://localhost:8000/v1 \
--dataset-path swebench_samples.jsonl \
--output-dir data_debug \
--max-turns 5 \
--limit 2
```
The backend model names must match what the server exposes. Because this mode surfaces token IDs/logprobs, the script saves both raw span logs and HuggingFace datasets per instance.
### Official Claude Code (Anthropic API)
```bash
export ANTHROPIC_API_KEY=sk-...
python claude_code_agent.py anthropic \
--dataset-path swebench_samples.jsonl \
--output-dir data_anthropic \
--frontend-model-high claude-sonnet-4-5-20250929 \
--frontend-model-low claude-haiku-4-5-20251001
```
Backend model flags are optional here because the Anthropic API strings match the frontend names. This path is ideal for validating prompts against the hosted experience (trace outputs do not contain token IDs or logprobs).
### OpenAI-Compatible Providers
```bash
export OPENAI_API_KEY=sk-...
python claude_code_agent.py openai \
--backend-model-high gpt-4.1 \
--backend-model-low gpt-4o-mini \
--dataset-path swebench_samples.jsonl \
--output-dir data_openai
```
Use this mode whenever Claude Code should talk to Azure OpenAI, OpenAI, or another compatible provider. `--base-url` is optional—pass it if your endpoint differs from the public OpenAI URL.
Adjust `--max-turns`, `--cooldown-seconds`, and `--limit` to control runtime and rate limits regardless of backend.
## Outputs and Trace Collection
- `output_dir/stream_<instance_id>.json` contains the complete span stream captured from the Lightning Store for each rollout.
- When running with `backend_type=vllm`, `output_dir/dataset-<instance_id>/` stores a HuggingFace dataset with token IDs, logprobs, prompts, and metadata produced by `ExtendedLlmProxyTraceToTriplet`.
- `logs/<instance_id>/` is created by the SWE-bench runtime and mirrors the console output from the container.
- Return values from the agent are also evaluated via `swebench_utils.evaluation.evaluate`, so `data_debug` (or your chosen folder) will contain evaluation reports alongside traces.
Use these artifacts to fine-tune models, debug Claude Code behavior, or replay rollouts in downstream Agent-lightning workflows.
+540
View File
@@ -0,0 +1,540 @@
# Copyright (c) Microsoft. All rights reserved.
"""Instrumented driver for running Claude Code on SWE-bench with Agent-lightning.
This script wires together the Lightning Store, LLM proxy, and Claude Code controller so
that every SWE-bench instance is executed inside the official Claude container while
capturing full Agent-lightning traces. It supports three backend modes:
- `vllm`: wrap an OpenAI-compatible endpoint (e.g., vLLM) for hosted OSS models while
collecting prompt/response token ids and logprobs.
- `anthropic`: call the official Claude Code API via `ANTHROPIC_API_KEY` for prompt
tuning. Backend model defaults to the provided frontend names.
- `openai`: route through any OpenAI-compatible provider using `OPENAI_API_KEY`.
Typical usage: hosted vLLM (requires model paths and --base-url)
```bash
# Run vLLM in background
vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
--max-model-len 131072 \
--enable-auto-tool-choice \
--tool-call-parser qwen3_coder \
--port 45993 &
python claude_code_agent.py vllm \
--backend-model-high Qwen/Qwen3-Coder-30B-A3B-Instruct \
--backend-model-low Qwen/Qwen3-Coder-30B-A3B-Instruct \
--base-url http://localhost:45993/v1 \
--dataset-path swebench_samples.jsonl \
```
Official Claude Code via Anthropic:
```bash
export ANTHROPIC_API_KEY=sk-...
python claude_code_agent.py anthropic \
--dataset-path swebench_samples.jsonl \
--output-dir data_anthropic
```
Any OpenAI-compatible backend:
```bash
export OPENAI_API_KEY=sk-...
python claude_code_agent.py openai \
--backend-model-high gpt-5.1-codex-mini \
--backend-model-low gpt-4.1-mini \
--dataset-path swebench_samples.jsonl
```
Use `--debug` to enable debug loggings.
"""
import asyncio
import json
import logging
import os
import resource
from argparse import ArgumentParser
from typing import Any, Dict, List, Literal, Optional, Sequence, cast
from claude_code_controller import ClaudeController
from datasets import Dataset
from extended_adapter import ExtendedLlmProxyTraceToTriplet
from swebench.harness.constants import SWEbenchInstance
from swebench.harness.utils import load_swebench_dataset # pyright: ignore[reportUnknownVariableType]
from swebench_utils.evaluation import evaluate
from swebench_utils.logging import log_for_evaluation
from transformers import AutoTokenizer, PreTrainedTokenizerBase
from agentlightning import (
InMemoryLightningStore,
LightningStoreServer,
LitAgentRunner,
OtelTracer,
setup_logging,
setup_module_logging,
)
from agentlightning.litagent import LitAgent
from agentlightning.llm_proxy import LLMProxy, ModelConfig
from agentlightning.store import LightningStore
from agentlightning.types import AttemptedRollout, NamedResources, ProxyLLM, Rollout, RolloutRawResult, Span
logger = logging.getLogger("claude_code_agent")
def _load_dataset(path: str, epoch: int = 0, limit: Optional[int] = None) -> List[SWEbenchInstance]:
instances: List[SWEbenchInstance] = []
with open(path) as f:
for line in f:
instance = json.loads(line)
instance["epoch"] = epoch
instances.append(instance)
if limit is not None:
instances = instances[:limit]
return instances
def _flatten_messages(messages: List[Any]) -> List[Dict[str, str]]:
flattened: List[Dict[str, str]] = []
for msg in messages:
if msg["role"] in ["system", "user"] and isinstance(msg["content"], list):
msg_content: List[str] = []
for content in msg["content"]:
msg_content.append(content["text"])
msg["content"] = "".join(msg_content)
elif msg["role"] == "assistant" and "tool_calls" in msg:
# NOTE:
# Tool calls are list of dict, though in most case only one tool call is made per call
# We serialize it as json string here to avoid nested structure
msg["tool_calls"] = json.dumps(msg["tool_calls"])
for k in msg:
assert isinstance(msg[k], str), f"\n>>> {msg}"
flattened.append(msg)
return flattened
class ClaudeCodeAgent(LitAgent[SWEbenchInstance]):
"""Claude Code Agent implementation.
This agent is a wrapper of the Claude Code controller,
and it should be used to run the Claude Code agent on SWE-bench datasets.
"""
def __init__(
self,
namespace: Literal["swebench", "starryzhang"] = "swebench",
max_turns: int = 5,
run_method: Literal["python", "cli"] = "cli",
open_file_limit: int = 4096,
cache_level: str = "env", # ["none", "base", "env", "instance"]
clean: bool = False,
force_rebuild: bool = False,
timeout: int = 1_800, # in sec
instance_image_tag: str = "latest",
rewrite_reports: bool = False,
swebench_full_dataset: Optional[List[SWEbenchInstance]] = None,
) -> None:
super().__init__()
self.namespace = namespace
self.max_turns = max_turns
self.run_method = run_method
self.cache_level = cache_level
self.clean = clean
self.force_rebuild = force_rebuild
self.timeout = timeout
self.instance_image_tag = instance_image_tag
self.rewrite_reports = rewrite_reports
self.swebench_full_dataset = (
{each["instance_id"]: each for each in swebench_full_dataset} if swebench_full_dataset is not None else {}
)
# Set the maximum number of open files to the specified limit.
resource.setrlimit(resource.RLIMIT_NOFILE, (open_file_limit, open_file_limit))
async def rollout_async(
self, task: SWEbenchInstance, resources: NamedResources, rollout: Rollout
) -> RolloutRawResult:
if not isinstance(rollout, AttemptedRollout):
# Technically, rollout should be an AttemptedRollout here.
# but the API is not stabilized yet.
raise ValueError("Rollout is not an AttemptedRollout.")
run_id = f"epoch_{task.get('epoch', 0)}"
image = f"{self.namespace}/sweb.eval.x86_64.{task['instance_id'].lower()}".replace("__", "_1776_")
llm = cast(ProxyLLM, resources["llm"])
try:
# 1. init container
controller = ClaudeController(
image,
task,
run_id,
llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
llm.api_key or os.environ.get("ANTHROPIC_AUTH_TOKEN", "dummy"),
)
# 2. execute task
prediction = controller.run_instance(
task, max_turns=self.max_turns, run_method=cast(Literal["python", "cli"], self.run_method)
)
del controller
except Exception as e:
log_for_evaluation(run_id, task["instance_id"], f"Exception during rollout: {e}")
return 0.0
# 3. obtain rewards (evaluation result)
reward = 0.0
# empty patch
if prediction["model_patch"] in ["", None]:
return reward
instance_id = prediction["instance_id"]
result = evaluate(
cast(Any, prediction),
self.swebench_full_dataset[instance_id],
self.cache_level,
self.clean,
self.force_rebuild,
run_id,
self.timeout,
namespace=self.namespace,
instance_image_tag=self.instance_image_tag,
rewrite_reports=self.rewrite_reports,
)
# error patch
if result is None:
return reward
report = result[1]
# resolved/unresolved patch
if report[instance_id]["resolved"]:
reward = 1.0
return reward
def sanity_check_spans(spans: Sequence[Span]) -> None:
assert len(spans) > 1, f"At least two spans are expected for a valid rollout. Found {len(spans)} spans."
assert any(span.name == "raw_gen_ai_request" for span in spans), "raw_gen_ai_request span not found"
assert any(span.name == "agentlightning.annotation" for span in spans), "agentlightning.annotation span not found"
async def run_instance_async(
instance: SWEbenchInstance,
agent: ClaudeCodeAgent,
runner: LitAgentRunner[SWEbenchInstance],
store: LightningStore,
output_dir: Optional[str],
adapter: Optional[ExtendedLlmProxyTraceToTriplet],
tokenizer: Optional[PreTrainedTokenizerBase],
) -> None:
"""Runs the agent on a specific SWE-bench instance.
Running on specific SWE-bench instance and queries the traced spans.
It then extracts the triplets and saves the dataset.
"""
instance_id = instance["instance_id"]
logger.info(f"Starting to run instance: {instance_id}")
# Run the agent and query the traced spans.
with runner.run_context(agent=agent, store=store):
rollout = await runner.step(instance)
logger.info(f"Finished running instance: {instance_id}")
spans = await store.query_spans(rollout.rollout_id)
if output_dir is None:
logger.info(f"Generated {len(spans)} spans for {instance_id}")
return
# 1. Dump raw spans (Common for both types)
raw_path = os.path.join(output_dir, f"stream_{instance_id}.json")
with open(raw_path, "w") as f:
for span in spans:
f.write(json.dumps(span.model_dump()) + "\n")
logger.info(f"Dumped {len(spans)} spans to {raw_path}")
# 2. Extract Triplets and Save Dataset (vLLM specific)
if adapter is not None and tokenizer is not None:
try:
triplets = adapter.adapt(cast(List[Span], spans))
logger.info(f"Extracted {len(triplets)} triplets for {instance_id}")
all_triplets: List[Dict[str, Any]] = []
recent_reward: Optional[float] = None
# Process in reverse to propagate rewards if necessary/logic dictates
for triplet in reversed(triplets):
if triplet.reward is not None:
recent_reward = triplet.reward
prompt_text = tokenizer.decode(triplet.prompt["token_ids"]) # type: ignore
all_triplets.append(
{
"repo": instance.get("repo", ""),
"instance_id": instance_id,
"turn": triplet.metadata["sequence_id"],
"prompt_ids": triplet.prompt["token_ids"],
"gold_completion_ids": triplet.response["token_ids"],
"logprobs": triplet.response["logprobs"],
"reward": recent_reward,
"prompt": prompt_text,
"messages": _flatten_messages(triplet.metadata["messages"]),
}
)
if all_triplets:
ds = Dataset.from_list(all_triplets) # type: ignore
save_path = os.path.join(output_dir, f"dataset-{instance_id}")
ds.save_to_disk(save_path) # type: ignore
logger.info(f"Saved HuggingFace dataset to {save_path}")
except Exception as e:
logger.error(f"Failed to extract triplets for {instance_id}: {e}")
logger.info(f"Finished extracting spans and traces for instance: {instance_id}")
# Quickly sanity check the spans
sanity_check_spans(spans)
logger.info(f"Sanity check passed for instance: {instance_id}")
async def dry_run_claude_code(
*,
dataset_path: str,
haiku_frontend_name: str,
haiku_backend_name: str,
sonnet_frontend_name: str,
sonnet_backend_name: str,
backend_type: Literal["vllm", "anthropic", "openai"],
api_base_url: Optional[str],
output_dir: Optional[str],
max_turns: int,
limit: Optional[int],
cooldown_seconds: float,
) -> None:
"""Executes a dry run of the Claude Code agent on a dataset.
This function handles both 'official' runs (interacting with Anthropic APIs)
and 'hosted' runs (interacting with vLLM or compatible servers). It manages
initialization of the Lightning Store, LLM Proxy, and the execution loop.
If running in 'vllm' mode, it will also attempt to extract triplets using
the provided backend name as the tokenizer path and save a HuggingFace Dataset.
Args:
dataset_path: Path to the JSONL dataset file.
haiku_frontend_name: The model name used in the code to request the 'fast' model.
haiku_backend_name: The actual model name/path on the backend.
sonnet_frontend_name: The model name used in the code to request the 'strong' model.
sonnet_backend_name: The actual model name/path on the backend.
backend_type: The type of backend to configure ("vllm", "anthropic" or "openai").
api_base_url: Base URL for the API. Required for "vllm" or "openai".
output_dir: Directory to save logs, spans, and datasets.
max_turns: Maximum number of steps the agent can take per instance.
limit: Optional limit on the number of instances to process.
"""
dataset = _load_dataset(dataset_path, limit=limit)
# Initialize Infrastructure
tracer = OtelTracer()
runner = LitAgentRunner[SWEbenchInstance](tracer)
store = LightningStoreServer(InMemoryLightningStore(), host="0.0.0.0", port=7654)
await store.start()
# Enable callbacks for training data extraction if using vLLM
callbacks = ["return_token_ids", "opentelemetry", "logprobs"] if backend_type == "vllm" else ["opentelemetry"]
llm_proxy = LLMProxy(port=12358, store=store, callbacks=callbacks)
# Configure Models based on backend type
model_configs: List[ModelConfig] = []
model_params: Dict[str, Any] = {}
if backend_type == "vllm":
model_namespace = "hosted_vllm"
if api_base_url:
model_params["api_base"] = api_base_url
else:
raise ValueError("api_base_url is required for vllm backend")
elif backend_type == "anthropic":
model_namespace = "anthropic"
model_params["api_key"] = "os.environ/ANTHROPIC_API_KEY"
if api_base_url:
model_params["api_base"] = api_base_url
elif backend_type == "openai":
model_namespace = "openai"
model_params["api_key"] = "os.environ/OPENAI_API_KEY"
if api_base_url:
# Users can still override this via environment variables,
# even if they don't pass it in as an argument.
model_params["api_base"] = api_base_url
model_configs.extend(
[
ModelConfig(
model_name=sonnet_frontend_name,
litellm_params={
"model": f"{model_namespace}/{sonnet_backend_name}",
**model_params,
},
),
ModelConfig(
model_name=haiku_frontend_name,
litellm_params={
"model": f"{model_namespace}/{haiku_backend_name}",
**model_params,
},
),
]
)
logger.info(f"Updating model list: {model_configs}")
llm_proxy.update_model_list(model_configs)
await llm_proxy.start()
try:
# Add the LLM proxy as a resource to the store
await store.add_resources({"llm": llm_proxy.as_resource(model="local")})
# Prepare for triplet extraction if vllm
adapter = ExtendedLlmProxyTraceToTriplet() if backend_type == "vllm" else None
tokenizer = None
if backend_type == "vllm":
try:
tokenizer = AutoTokenizer.from_pretrained(sonnet_backend_name) # type: ignore
except Exception as e:
logger.warning(f"Could not load tokenizer for {sonnet_backend_name}: {e}")
# Load full swebench dataset. Mainly for evaluation purposes.
swebench_full_dataset = load_swebench_dataset("princeton-nlp/SWE-bench", split="test")
# Initialize Claude Code Agent
claude_code_agent = ClaudeCodeAgent(swebench_full_dataset=swebench_full_dataset, max_turns=max_turns)
# Execution Loop
for instance in dataset:
await run_instance_async(
instance,
claude_code_agent,
runner,
store,
output_dir,
adapter,
cast(PreTrainedTokenizerBase, tokenizer),
)
# Basic sleep to allow resource cleanup or rate limit cooling
await asyncio.sleep(cooldown_seconds)
finally:
await llm_proxy.stop()
await store.stop()
if __name__ == "__main__":
parser = ArgumentParser(description="Run Claude Code Agent experiments.")
# Backend Selection
parser.add_argument(
"backend_type",
type=str,
choices=["vllm", "anthropic", "openai"],
help="Backend type: 'vllm' for hosted models, 'anthropic' for official API, 'openai' for OpenAI API.",
)
# Model Configuration
parser.add_argument(
"--backend-model-high",
type=str,
default=None,
help="Backend model path/name for expensive model usages (used as vLLM model name / OpenAI model name).",
)
parser.add_argument(
"--backend-model-low",
type=str,
default=None,
help="Backend model path/name for low-price model usages (used as vLLM model name / OpenAI model name).",
)
parser.add_argument(
"--base-url", type=str, default="http://localhost:8000/v1", help="LLM server address (required for vllm)."
)
# Frontend/Agent Configuration
parser.add_argument(
"--frontend-model-high",
type=str,
default="claude-sonnet-4-5-20250929",
help="The frontend high-price model name provided to Claude Code.",
)
parser.add_argument(
"--frontend-model-low",
type=str,
default="claude-haiku-4-5-20251001",
help="The frontend low-price model name provided to Claude Code.",
)
# Execution Configuration
parser.add_argument("--dataset-path", type=str, default="swebench_samples.jsonl", help="Path to the dataset.")
parser.add_argument("--max-turns", type=int, default=5, help="Maximum turns per instance.")
parser.add_argument("--output-dir", type=str, default="data", help="Directory to save output logs.")
parser.add_argument("--limit", type=int, default=None, help="Limit the number of instances to run (for debugging).")
parser.add_argument("--cooldown-seconds", type=float, default=2.0, help="Cooldown seconds between instances.")
parser.add_argument("--debug", action="store_true", help="Enable debug loggings.")
args = parser.parse_args()
if args.output_dir is not None:
os.makedirs(args.output_dir, exist_ok=True)
if args.debug:
setup_logging()
setup_module_logging("DEBUG", name="claude_code_agent")
else:
setup_logging(apply_to=[logger.name])
# Map backend_type to the appropriate args
backend_mode = cast(Literal["vllm", "anthropic", "openai"], args.backend_type)
# If using anthropic, the backend name usually matches the frontend or is specific API string.
# Otherwise, the backend name is the model name/path (e.g., Qwen/...) and must be provided.
if args.backend_model_high is None:
if args.backend_type == "anthropic":
backend_model_high = args.frontend_model_high
else:
raise ValueError("--backend-model-high is required for non-anthropic backends")
else:
backend_model_high = args.backend_model_high
if args.backend_model_low is None:
if args.backend_type == "anthropic":
backend_model_low = args.frontend_model_low
else:
raise ValueError("--backend-model-low is required for non-anthropic backends")
else:
backend_model_low = args.backend_model_low
asyncio.run(
dry_run_claude_code(
dataset_path=args.dataset_path,
haiku_frontend_name=args.frontend_model_low,
haiku_backend_name=backend_model_low,
sonnet_frontend_name=args.frontend_model_high,
sonnet_backend_name=backend_model_high,
backend_type=backend_mode,
api_base_url=args.base_url if backend_mode == "vllm" else None,
output_dir=args.output_dir,
max_turns=args.max_turns,
limit=args.limit,
cooldown_seconds=args.cooldown_seconds,
)
)
@@ -0,0 +1,227 @@
# Copyright (c) Microsoft. All rights reserved.
"""Controller module for managing Claude Code executions in containerized environments.
This module provides the ClaudeController class that manages the execution of Claude Code
within Docker containers. It handles container initialization, command execution, and
patch application for SWE-bench evaluation tasks.
"""
import logging
from functools import partial
from typing import Literal, TypedDict
import dotenv
from swebench.harness.constants import SWEbenchInstance
from swebench_utils.docker_runtime import Runtime
from swebench_utils.logging import log_for_evaluation
SWEBENCH_EXTRA_SYSTEM_PROMPT = """
You are an expert software engineer solving swebench bug fixing tasks.
"""
SWEBENCH_USER_PROMPT = """
You are given a code repository in the current directory (/testbed).
The bug description is:
{description}
=================================================
You task is to fix the bug with the following steps:
(1) write test cases to reproduce the bug.
(2) explore the source codes to locate the bug.
(3) edit the source codes to fix the bug.
(4) rerun your written test cases to validate that the bug is fixed. If not, go back to explore the source codes and fix the codes again.
(5) remember to delete the test cases you write at last.
Please do not commit your edits. We will do it later.
"""
logger = logging.getLogger("claude_code_agent")
class RunInstanceResult(TypedDict):
instance_id: str
model_patch: str
model_name_or_path: str
class ClaudeController:
"""Manages the execution of Claude Code within a Docker runtime.
This controller handles the lifecycle of a SWE-bench task execution, including
environment setup, tool installation, agent execution (via CLI or Python SDK),
and result extraction.
Attributes:
container: The active Docker runtime session.
"""
def __init__(self, image: str, instance: SWEbenchInstance, run_id: str, endpoint: str, api_key: str) -> None:
"""Initialize the ClaudeController.
Args:
image: The Docker image tag.
instance: The dataset instance containing the problem statement and ID.
run_id: The identifier for the evaluation run.
endpoint: The API endpoint URL.
api_key: The API authentication key.
"""
self.image = image
self.instance = instance
self.run_id = run_id
self.endpoint = endpoint
self.api_key = api_key
self.container: Runtime = self.init_container(self.image, self.instance)
def init_container(self, image: str, instance: SWEbenchInstance) -> Runtime:
"""Initializes the Docker container and sets up the Claude Code environment.
This method starts the container session, installs the Claude CLI,
configures environment variables for authentication and sandbox mode.
Args:
image: The Docker image tag to start.
instance: The dataset instance to load into the environment.
Returns:
An initialized and configured Docker runtime object.
"""
container = Runtime.start_session(
image,
instance,
log_function=partial(log_for_evaluation, run_id=self.run_id, instance_id=instance["instance_id"]),
)
# Install Claude CLI
container.send_command("curl -fsSL https://claude.ai/install.sh | bash")
container.send_command('alias claude="$HOME/.local/bin/claude"')
# Configure Environment
dotenv.load_dotenv()
container.send_command(f"export ANTHROPIC_BASE_URL={self.endpoint}")
container.send_command(f"export ANTHROPIC_AUTH_TOKEN={self.api_key}")
container.send_command("export IS_SANDBOX=1")
return container
def _run_cli(self, instance: SWEbenchInstance, max_turns: int, time_limit: int) -> None:
"""Executes Claude Code using the Command Line Interface.
Constructs a safe heredoc for the prompt to avoid shell interpolation issues
and executes the `claude` binary directly.
Args:
instance: The problem instance containing the problem statement.
max_turns: The maximum number of interaction turns allowed.
time_limit: The execution time limit in minutes.
"""
# Prepare prompt safely: write it to a file inside the container using a single-quoted heredoc
# directly applying prompt for heredoc may raise error for windows line ending \r\n
prompt_text = SWEBENCH_USER_PROMPT.format(description=instance["problem_statement"].replace('"""', "'''"))
# Choose a simple filename and a heredoc delimiter unlikely to collide
heredoc_cmd = "cat > /tmp/cc_prompt.txt <<'CC_PROMPT'\n" + prompt_text + "\nCC_PROMPT\n"
self.container.send_command(heredoc_cmd)
# Run claude reading the prompt from the file
claude_cmd = (
f'claude -p "$(cat /tmp/cc_prompt.txt)" '
f'--append-system-prompt "{SWEBENCH_EXTRA_SYSTEM_PROMPT}" '
f"--max-turns {max_turns} "
f"--dangerously-skip-permissions "
f"--output-format json --verbose"
)
logger.info(f"Running Claude Code CLI command: {claude_cmd}")
self.container.send_command(claude_cmd, time_limit * 60)
logger.info(f"Claude Code CLI command completed")
def _run_python_sdk(self, instance: SWEbenchInstance, max_turns: int, time_limit: int) -> None:
"""Executes Claude Code using the Python SDK wrapper.
Installs the Python SDK if necessary, hydrates a template script with the
problem prompt, and executes the generated Python script.
Note:
This path is still under development and not yet stable.
Args:
instance: The problem instance containing the problem statement.
max_turns: The maximum number of interaction turns allowed.
time_limit: The execution time limit in minutes.
"""
# Ensure Python 3.12 is available
self.container.send_command(
f"""
if ! command -v python3 &> /dev/null; then
echo "Python is not installed. Installing Python 3.12..."
sudo apt-get update -qq && sudo apt-get install -y -qq python3.12
else
echo "Python is already installed."
fi
"""
)
self.container.send_command("python3 -m pip install claude-code-sdk")
# Load and fill the execution template
with open("src/agent/cc/claude_code_main.py.template") as f:
entrance_template = f.read()
script_content = (
entrance_template.replace("SYS_PROMPT", SWEBENCH_EXTRA_SYSTEM_PROMPT)
.replace(
"PROMPT", SWEBENCH_USER_PROMPT.format(description=instance["problem_statement"].replace('"""', "'''"))
)
.replace("MAX_STEP", str(max_turns))
)
# Write the script to the container and execute
self.container.send_command(f"cat > /tmp/claude_code_main.py <<'CC_MAIN'\n{script_content}\nCC_MAIN\n")
self.container.send_command("python3 /tmp/claude_code_main.py", time_limit * 60)
return
def run_instance(
self,
instance: SWEbenchInstance,
max_turns: int = 40,
time_limit: int = 30,
run_method: Literal["python", "cli"] = "python",
) -> RunInstanceResult:
"""Runs the agent on a specific SWE-bench instance.
This method orchestrates the agent execution via the specified method (CLI or Python),
and extracts the generated git diff (patch) upon completion.
Args:
instance: The dataset instance dictionary.
max_turns: Maximum conversation turns allowed for the agent. Defaults to 40.
time_limit: Time limit for the execution in minutes. Defaults to 30.
run_method: The execution method, either "python" (SDK) or "cli". Defaults to "python".
Returns:
A dictionary containing the result:
- instance_id: The ID of the processed instance.
- model_patch: The git diff generated by the agent.
- model_name_or_path: Hardcoded to "cc" (Claude Code).
Raises:
ValueError: If `run_method` is not "python" or "cli".
"""
if run_method == "python":
logger.warning("Running Claude Code using Python SDK is still under development and not yet stable.")
self._run_python_sdk(instance, max_turns, time_limit)
elif run_method == "cli":
self._run_cli(instance, max_turns, time_limit)
else:
raise ValueError(f"Wrong run_method '{run_method}', run_method should be in ['python', 'cli']")
result = self.container.send_command("git --no-pager diff HEAD")
git_diff = result.output.replace("git --no-pager diff HEAD\n", "")
return {
"instance_id": instance["instance_id"],
"model_patch": git_diff,
"model_name_or_path": "cc",
}
def __del__(self) -> None:
"""Destructor to ensure container resources are cleaned up."""
if hasattr(self, "container"):
self.container.cleanup()
+163
View File
@@ -0,0 +1,163 @@
# Copyright (c) Microsoft. All rights reserved.
"""Custom adapter module for converting LLM proxy traces to augmented trajectories.
This module provides an augmented LlmProxyTraceToTriplet adapter that converts
LLM proxy spans into augmented trajectories for analysis and evaluation.
It extends the base LlmProxyTraceToTriplet to include additional metadata like chat messages,
log probabilities, and sequence IDs.
"""
import logging
from typing import Any, Dict, List, Optional, Tuple, cast
from agentlightning.adapter.triplet import LlmProxyTraceToTriplet
from agentlightning.types import Span, Triplet
logger = logging.getLogger(__name__)
class ExtendedLlmProxyTraceToTriplet(LlmProxyTraceToTriplet):
"""Convert LLM Proxy spans into trajectories with logprobs and customized metadata.
Augmented fields include:
- chat messages history from [`llm.hosted_vllm.messages`], saved to `Triplet.metadata['messages']`
- logprobs from [`llm.hosted_vllm.choices`], saved to `Triplet.response['logprobs']`
- sequence_id from [`Span.sequence_id`] to locate the order of the span (conversation turn), saved to `Triplet.metadata['sequence_id']`
"""
def _extract_tokens_from_raw(self, attrs: Dict[str, Any]) -> Tuple[List[int], List[int], List[float]]: # type: ignore
"""Extract token ids from raw_gen_ai_request attributes.
- llm.hosted_vllm.prompt_token_ids: string -> List[int]
- llm.hosted_vllm.choices: string -> [{'token_ids': [...]}] -> take first
"""
prompt_ids: List[int] = []
resp_ids: List[int] = []
logprobs: List[float] = []
# prompt
p = attrs.get("llm.hosted_vllm.prompt_token_ids")
p = self._literal_eval_maybe(p)
if isinstance(p, list) and all(isinstance(x, int) for x in p): # type: ignore
prompt_ids = cast(List[int], p)
choices = attrs.get("llm.hosted_vllm.choices")
choices = self._literal_eval_maybe(choices)
if isinstance(choices, list) and choices:
cand = cast(Any, choices[0])
if isinstance(cand, dict):
tids = cast(Dict[str, Any], cand).get("token_ids")
if isinstance(tids, list) and all(isinstance(x, int) for x in tids): # type: ignore
resp_ids = cast(List[int], tids)
if "logprobs" in cand:
logprobs_dict = cast(Dict[str, Any], cand).get("logprobs")
if isinstance(logprobs_dict, dict) and "content" in logprobs_dict:
content = cast(List[Dict[str, Any]], logprobs_dict["content"])
logprobs = [float(item["logprob"]) for item in content if "logprob" in item]
return prompt_ids, resp_ids, logprobs
def adapt(self, source: List[Span], /) -> List[Triplet]: # type: ignore
"""Convert LLM Proxy spans into [`Triplet`][agentlightning.Triplet] trajectories.
Args:
source: Spans emitted by the LLM Proxy containing prompt, response, and reward data.
Returns:
Ordered trajectory transitions matched purely by `sequence_id`.
"""
# 1) Sort deterministically by (sequence_id, start_time).
spans = sorted(
source,
key=lambda s: (s.sequence_id, s.start_time),
)
# 2) Collect LLM calls
llm_items: List[Dict[str, Any]] = []
seen_request_ids: set[str] = set()
for s in spans:
attrs = s.attributes or {}
prompt_ids: List[int] = []
resp_ids: List[int] = []
logprobs: List[float] = []
if s.name == "raw_gen_ai_request":
prompt_ids, resp_ids, logprobs = self._extract_tokens_from_raw(attrs)
if len(prompt_ids) == 0 or len(resp_ids) == 0:
logger.warning(
f"Span {s.span_id} is missing prompt (len={len(prompt_ids)}) or response (len={len(resp_ids)}) token ids. Ignoring this span."
)
continue
elif len(logprobs) == 0:
logger.warning(f"Span {s.span_id} is missing logprobs. Ignoring logprobs for this span.")
continue
elif len(resp_ids) != len(logprobs):
logger.warning(
f"Span {s.span_id} has mismatched response ids and logprobs lengths: "
f"{len(resp_ids)} vs {len(logprobs)}. Ignoring this span."
)
continue
if prompt_ids and resp_ids and logprobs:
rid = self._request_id_from_attrs(attrs)
if rid:
# Duplicated request ID. This request is already handled.
if rid in seen_request_ids:
continue
seen_request_ids.add(rid)
llm_items.append(
dict(
span=s,
seq=s.sequence_id,
response_ids=resp_ids,
prompt_ids=prompt_ids,
request_id=rid,
logprobs=logprobs,
)
)
# Order LLM items by sequence only.
llm_items.sort(key=lambda x: x["seq"])
# Collect rewards by sequence only.
rewards: List[Tuple[int, Optional[float]]] = []
for s in spans:
val = self._maybe_reward_value(s)
if val is not None:
rewards.append((s.sequence_id, val))
# First-occurrence matching by sequence_id only:
# For reward at sequence R, assign to the most recent unmatched LLM with seq < R.
assigned: Dict[str, Optional[float]] = {}
for r_seq, r_val in sorted(rewards, key=lambda x: x[0]):
for item in reversed(llm_items):
sid = item["span"].span_id
if sid in assigned:
continue
if item["seq"] < r_seq:
assigned[sid] = r_val
break
# Build triplets in LLM sequence order.
triplets: List[Triplet] = []
for item in llm_items:
s = item["span"]
triplets.append(
Triplet(
prompt={"token_ids": item["prompt_ids"]},
response={"token_ids": item["response_ids"], "logprobs": item["logprobs"]},
reward=assigned.get(s.span_id, None),
metadata=dict(
# This is called response_id to align with the other adapters.
response_id=item["request_id"],
sequence_id=item["seq"],
messages=self._literal_eval_maybe(s.attributes.get("llm.hosted_vllm.messages")),
),
)
)
return triplets
File diff suppressed because one or more lines are too long
@@ -0,0 +1 @@
# Copyright (c) Microsoft. All rights reserved.
@@ -0,0 +1,430 @@
# Copyright (c) Microsoft. All rights reserved.
"""Docker runtime management for repository setup and command execution.
Provides containerized environment for repository testing with command execution,
file operations, and state management capabilities.
"""
from __future__ import annotations
import json
import logging
import queue
import re
import threading
import time
import uuid
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional
from docker.errors import DockerException, ImageNotFound
from docker.models.containers import Container
from swebench.harness.constants import SWEbenchInstance
from typing_extensions import Self
import docker
# This will log to the console for debugging purposes.
claude_code_logger = logging.getLogger("claude_code_agent.docker_runtime")
CMD_OUTPUT_PS1_BEGIN = "\n###PS1JSON###\n"
CMD_OUTPUT_PS1_END = "\n###PS1END###"
CMD_OUTPUT_METADATA_PS1_REGEX = re.compile(
r"(?m)^\s*" + re.escape(CMD_OUTPUT_PS1_BEGIN.strip()) + r"\s*(.*?)\s*" + re.escape(CMD_OUTPUT_PS1_END.strip()),
re.DOTALL,
)
ANSI_ESCAPE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
TIMEOUT_EXIT_CODE = 124
MEM_LIMIT = "8g"
CPU_CORES = 4
VAR_PATTERNS = {
"exit_code": re.compile(r'"exit_code":\s*(-?\d+)\s*(?:,|\})'),
"username": re.compile(r'"username":\s*"([^"]*)"'),
"hostname": re.compile(r'"hostname":\s*"([^"]*)"'),
"working_dir": re.compile(r'"working_dir":\s*"([^"]*)"'),
"py_interpreter_path": re.compile(r'"py_interpreter_path":\s*"([^"]*)"'),
}
@dataclass
class CmdOutputMetadata:
"""
Additional metadata captured from PS1 shell prompt.
Provides context about command execution environment including
exit codes, user info, working directory, and Python interpreter.
"""
exit_code: int = -1
username: str | None = None
hostname: str | None = None
working_dir: str | None = None
py_interpreter_path: str | None = None
@classmethod
def matches_ps1_metadata(cls, output: str) -> List[re.Match[str]]:
matches: List[re.Match[str]] = []
for match in CMD_OUTPUT_METADATA_PS1_REGEX.finditer(output):
scope = match.group(1).strip()
try:
d = json.loads(scope) # Try to parse as JSON
matches.append(match)
except json.JSONDecodeError:
d = cls.best_effort_match(scope)
if len(d) > 0:
matches.append(match)
return matches
@classmethod
def best_effort_match(cls, scope: str) -> Dict[str, Any]:
out: Dict[str, str] = {}
for field, pattern in VAR_PATTERNS.items():
m = pattern.search(scope)
if m:
out[field] = m.group(1)
else:
out[field] = ""
return out
@classmethod
def from_ps1_match(cls, match: re.Match[str]) -> Self:
"""
Extract metadata from a PS1 prompt regex match.
Args:
match (re.Match[str]): Regex match containing JSON metadata
Returns:
Self: CmdOutputMetadata instance with parsed values
"""
try:
metadata = json.loads(match.group(1))
except:
metadata = cls.best_effort_match(match.group(1))
# Create a copy of metadata to avoid modifying the original
processed = metadata.copy()
# Convert numeric fields
if "exit_code" in metadata:
try:
processed["exit_code"] = int(float(str(metadata["exit_code"])))
except (ValueError, TypeError):
processed["exit_code"] = -1
return cls(**processed)
@dataclass
class CommandResult:
"""
Result of a command execution with output and metadata.
Attributes:
output (str): Command output text
metadata (Optional[CmdOutputMetadata]): Execution context metadata
"""
output: str
metadata: Optional[CmdOutputMetadata]
def to_observation(self, strip: bool = True) -> str:
"""
Convert command result to formatted observation string.
Args:
strip (bool): Whether to truncate long output
Returns:
str: Formatted observation with output and context
"""
# compile regex once for efficiency
ANSI_ESCAPE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
output = ANSI_ESCAPE.sub("", self.output).replace("\r", "")
if len(output) > 1024 * 8 and strip:
output = output[: 1024 * 4] + "....stripped due to length....\n" + output[-1024 * 4 :]
if self.metadata is None:
return f"\n{output}\n"
return f"""{output}
{self.metadata.username}@{self.metadata.hostname}:{self.metadata.working_dir} $
exit code: {self.metadata.exit_code}
"""
class Runtime:
"""
Docker container runtime for repository setup and testing.
Manages a Docker container with persistent bash session, command execution,
file operations, and container lifecycle management.
"""
def __init__(self, container: Container, log_function: Callable[..., None]) -> None:
"""
Initialize runtime with an existing Docker container.
Args:
container (Container): Docker container instance to manage
"""
self.container = container
self.logger = log_function # Set logger early so it's available even if init fails later
self.sock: Any = self.container.attach_socket(params={"stdin": 1, "stdout": 1, "stderr": 1, "stream": 1}) # type: ignore
self.output_queue: queue.Queue[bytes] = queue.Queue()
self._start_output_thread()
self._clear_initial_prompt()
json_str = json.dumps(
{
"exit_code": "$?",
"username": r"\u",
"hostname": r"\h",
"working_dir": r"$(pwd)",
"py_interpreter_path": r'$(which python 2>/dev/null || echo "")',
},
indent=2,
).replace('"', r"\"")
ps1 = CMD_OUTPUT_PS1_BEGIN + json_str + CMD_OUTPUT_PS1_END + "\n"
self.send_command(f'export PROMPT_COMMAND=\'export PS1="{ps1}"\'; export PS2=""')
self.send_command("apt update -qq && apt install -y -qq git")
self.stopped = False
def _stream_output(self):
while True:
try:
output = self._recv_bytes(4096)
if not output:
break
self.output_queue.put(output)
except (OSError, ConnectionError) as e:
print(f"Connection error in _stream_output: {e}")
break
except Exception as e:
# print(f"Unexpected error in _stream_output: {e}")
break
def _start_output_thread(self):
self.output_thread = threading.Thread(target=self._stream_output, daemon=True)
self.output_thread.start()
# TODO: kill the thread if main thread is stopped
def _clear_initial_prompt(self):
time.sleep(0.5)
while not self.output_queue.empty():
self.output_queue.get()
def _read_raw_output(self, timeout: float = 30) -> tuple[str, Optional[CmdOutputMetadata]]:
accumulated_output = ""
start_time = time.time()
while time.time() - start_time < timeout:
try:
chunk = self.output_queue.get(timeout=0.1)
accumulated_output += chunk.decode("utf-8", errors="ignore")
# PSReadLine injects ANSI + cursor control; normalize before matching
accumulated_clean = ANSI_ESCAPE.sub("", accumulated_output).replace("\r", "")
ps1_matches = CmdOutputMetadata.matches_ps1_metadata(accumulated_clean)
if ps1_matches:
break
except queue.Empty:
continue
accumulated_output = ANSI_ESCAPE.sub("", accumulated_output).replace("\r", "")
ps1_matches = CmdOutputMetadata.matches_ps1_metadata(accumulated_output)
metadata = CmdOutputMetadata.from_ps1_match(ps1_matches[-1]) if ps1_matches else None
output = self._combine_outputs_between_matches(
accumulated_output,
ps1_matches,
)
return output, metadata
def _combine_outputs_between_matches(self, pane_content: str, ps1_matches: list[re.Match[str]]) -> str:
if len(ps1_matches) == 1:
return pane_content[: ps1_matches[0].start()]
elif len(ps1_matches) == 0:
return pane_content
output_segments: List[str] = []
for i in range(len(ps1_matches) - 1):
output_segment = pane_content[ps1_matches[i].end() + 1 : ps1_matches[i + 1].start()]
output_segments.append(output_segment)
return "\n".join(output_segments) + "\n" if output_segments else ""
def _recv_bytes(self, n: int = 4096) -> bytes:
# Prefer the public API on whatever object the SDK returns
for m in ("recv", "read"):
if hasattr(self.sock, m):
return getattr(self.sock, m)(n)
# Last-resort fallback for odd wrappers that still expose ._sock
if hasattr(self.sock, "_sock"):
for m in ("recv", "read"):
if hasattr(self.sock._sock, m):
return getattr(self.sock._sock, m)(n)
raise TypeError(f"Don't know how to read from {type(self.sock).__name__}")
def _send_bytes(self, data: bytes) -> None:
if hasattr(self.sock, "_sock"):
for m in ("send", "sendall", "write"):
if hasattr(self.sock._sock, m):
getattr(self.sock._sock, m)(data)
return
for m in ("send", "sendall", "write"):
if hasattr(self.sock, m):
getattr(self.sock, m)(data)
return
raise TypeError(f"Don't know how to write to {type(self.sock).__name__}")
def _log_command_result(self, result: CommandResult) -> None:
claude_code_logger.debug("Docker runtime command finished with metadata: %s", result.metadata)
if len(result.output) > 2048:
logged_output = result.output[:1024] + "\n(... stripped due to length ...)\n" + result.output[-1024:]
else:
logged_output = result.output
claude_code_logger.debug(
"Docker runtime command finished with output (length = %d):\n%s", len(result.output), logged_output
)
# Output to the evaluation logger simultaneously
self.logger(text=logged_output)
def send_command(self, command: str, timeout: float = 20 * 60) -> CommandResult:
# Redact sensitive API keys from the command before logging
redacted_command = command
for sensitive_var in ["ANTHROPIC_AUTH_TOKEN", "API_KEY", "SECRET_KEY"]:
pattern = rf"(export (.*?){re.escape(sensitive_var)}(.*?)=)[^\s]+"
redacted_command = re.sub(pattern, rf"\1****REDACTED****", redacted_command)
claude_code_logger.info("Docker runtime receiving command: %s", redacted_command)
# Normalize newline semantics for interactive shells
if not command.endswith("\n"):
command += "\n"
while not self.output_queue.empty():
self.output_queue.get()
self._send_bytes(command.encode())
output, metadata = self._read_raw_output(timeout=timeout)
# TODO: Check exit code of the command (claude code download fail will not be caught by this)
if metadata is not None:
result = CommandResult(output=output, metadata=metadata)
self._log_command_result(result)
return result
# handle timeout
self._send_bytes(b"\x03")
kill_timeout = 5.0
kill_output, kill_metadata = self._read_raw_output(timeout=kill_timeout)
output = output + kill_output + "\n**Exited due to timeout**\n"
if kill_metadata is not None:
kill_metadata.exit_code = TIMEOUT_EXIT_CODE
result = CommandResult(output=output, metadata=kill_metadata)
self._log_command_result(result)
return result
fallback_metadata = CmdOutputMetadata(
exit_code=TIMEOUT_EXIT_CODE,
)
result = CommandResult(output=output, metadata=fallback_metadata)
self._log_command_result(result)
return result
def cleanup(self) -> None:
if self.stopped:
return
try:
claude_code_logger.info(f"Stopping container: {self.container.id}")
self.container.stop()
claude_code_logger.info(f"Removing container: {self.container.id}")
self.container.remove(force=True)
claude_code_logger.info(f"Container removed: {self.container.id}")
self.stopped = True
except Exception as e:
print(f"Failed to stop container: {e}")
def __del__(self):
self.cleanup()
@staticmethod
def pull_image(image_name: str) -> bool:
"""
Pull Docker image from registry.
Args:
image_name (str): Name of the Docker image to pull
Returns:
bool: True if successful, False if image not found
"""
client = docker.from_env()
try:
client.images.pull(image_name)
return True
except ImageNotFound:
return False
@classmethod
def start_session(
cls,
image_name: str,
instance: SWEbenchInstance,
log_function: Callable[..., None] = lambda: None,
) -> Runtime:
"""
Start a Docker container session for repository testing.
Args:
image_name (str): Base Docker image name
instance (dict): SWE-bench instance data with repo info
Returns:
SetupRuntime: Configured runtime session ready for command execution
Raises:
RuntimeError: If Docker is not available
"""
try:
docker.from_env().ping() # type: ignore
except DockerException:
raise RuntimeError("Docker is not installed or not running.")
_ = cls.pull_image(image_name)
client = docker.from_env(timeout=600)
container_id = instance["instance_id"]
container_name = f"git-launch-{container_id}-{str(uuid.uuid4())[:4]}"
info: Dict[str, str] = client.version() # type: ignore
engine_os: str = (info.get("Os") or info.get("OSType") or "").lower() # type: ignore
# which operating system this code is running on, note windows can run linux containers, so engine_os != (container) platform
extra_hosts = {"host.docker.internal": "host-gateway"} if "linux" in engine_os else None
shell_command = "/bin/bash"
working_dir = "/testbed"
claude_code_logger.info(
f"Starting container {container_name} with image {image_name}. Shell command: {shell_command}"
)
container = client.containers.run(
image_name,
name=container_name,
command=shell_command,
stdin_open=True,
tty=True,
detach=True,
environment={
"TERM": "xterm-mono",
},
working_dir=working_dir,
extra_hosts=extra_hosts,
network_mode="host",
cpu_quota=int(CPU_CORES * 100000),
mem_limit=MEM_LIMIT,
)
claude_code_logger.info(f"Container {container_name} started with ID: {container.id}")
session = cls(container, log_function=log_function)
return session
@@ -0,0 +1,258 @@
# Copyright (c) Microsoft. All rights reserved.
"""Evaluation module for SWE-bench instance testing and grading.
This module provides core functionality for evaluating model predictions on SWE-bench
instances. It handles containerized execution of test scripts, patch application,
and generation of evaluation reports. The module orchestrates the complete evaluation
process including container management, patch application, test execution, and result
grading.
"""
import json
import traceback
from pathlib import Path, PurePosixPath
from typing import Any, Dict, Optional
from docker.models.containers import ExecResult
from swebench.harness.constants import (
APPLY_PATCH_FAIL,
APPLY_PATCH_PASS,
DOCKER_PATCH,
DOCKER_USER,
DOCKER_WORKDIR,
INSTANCE_IMAGE_BUILD_DIR,
KEY_MODEL,
KEY_PREDICTION,
LOG_INSTANCE,
LOG_REPORT,
LOG_TEST_OUTPUT,
RUN_EVALUATION_LOG_DIR,
UTF8,
SWEbenchInstance,
)
from swebench.harness.docker_build import close_logger # type: ignore
from swebench.harness.docker_build import (
BuildImageError,
build_container,
setup_logger,
)
from swebench.harness.docker_utils import cleanup_container # type: ignore
from swebench.harness.docker_utils import exec_run_with_timeout # type: ignore
from swebench.harness.docker_utils import remove_image # type: ignore
from swebench.harness.docker_utils import should_remove # type: ignore
from swebench.harness.docker_utils import (
copy_to_container,
)
from swebench.harness.grading import get_eval_report
from swebench.harness.test_spec.test_spec import TestSpec, make_test_spec
from swebench.harness.utils import EvaluationError
import docker
GIT_APPLY_CMDS = [
"git apply --verbose",
"git apply --verbose --reject",
"patch --batch --fuzz=5 -p1 -i",
]
def run_instance(
test_spec: TestSpec,
pred: Dict[str, Any],
rm_image: bool,
force_rebuild: bool,
client: docker.DockerClient,
run_id: str,
timeout: int | None = None,
rewrite_reports: bool = False,
):
"""
Run a single instance with the given prediction.
Args:
test_spec (TestSpec): TestSpec instance
pred (dict): Prediction w/ model_name_or_path, model_patch, instance_id
rm_image (bool): Whether to remove the image after running
force_rebuild (bool): Whether to force rebuild the image
client (docker.DockerClient): Docker client
run_id (str): Run ID
timeout (int): Timeout for running tests
rewrite_reports (bool): True if eval run is just to reformat existing report
"""
# Set up logging directory
instance_id = test_spec.instance_id
model_name_or_path = pred.get(KEY_MODEL, "None").replace("/", "__")
log_dir = RUN_EVALUATION_LOG_DIR / run_id / model_name_or_path / instance_id
# Set up report file
report_path = log_dir / LOG_REPORT
if rewrite_reports:
test_output_path = log_dir / LOG_TEST_OUTPUT
if not test_output_path.exists():
raise ValueError(f"Test output file {test_output_path} does not exist")
report = get_eval_report(
test_spec=test_spec,
prediction=pred,
test_log_path=test_output_path,
include_tests_status=True,
)
# Write report to report.json
with open(report_path, "w") as f:
f.write(json.dumps(report, indent=4))
return instance_id, report
if report_path.exists():
return instance_id, json.loads(report_path.read_text())
if not test_spec.is_remote_image:
# Link the image build dir in the log dir
build_dir = INSTANCE_IMAGE_BUILD_DIR / test_spec.instance_image_key.replace(":", "__")
image_build_link = log_dir / "image_build_dir"
if not image_build_link.exists():
try:
# link the image build dir in the log dir
image_build_link.symlink_to(build_dir.absolute(), target_is_directory=True)
except:
# some error, idk why
pass
# Set up logger
log_dir.mkdir(parents=True, exist_ok=True)
log_file = log_dir / LOG_INSTANCE
logger = setup_logger(instance_id, log_file)
# Run the instance
container = None
try:
# Build + start instance container (instance image should already be built)
container = build_container(test_spec, client, run_id, logger, rm_image, force_rebuild)
container.start()
logger.info(f"Container for {instance_id} started: {container.id}")
# Copy model prediction as patch file to container
patch_file = Path(log_dir / "patch.diff")
patch_file.write_text(pred[KEY_PREDICTION] or "")
logger.info(f"Intermediate patch for {instance_id} written to {patch_file}, now applying to container...")
copy_to_container(container, patch_file, PurePosixPath(DOCKER_PATCH)) # type: ignore
# Attempt to apply patch to container (TODO: FIX THIS)
val: Optional[ExecResult] = None
for git_apply_cmd in GIT_APPLY_CMDS:
val = container.exec_run( # type: ignore
f"{git_apply_cmd} {DOCKER_PATCH}",
workdir=DOCKER_WORKDIR,
user=DOCKER_USER,
)
if val.exit_code == 0:
logger.info(f"{APPLY_PATCH_PASS}:\n{val.output.decode(UTF8)}")
break
else:
logger.info(f"Failed to apply patch to container: {git_apply_cmd}")
if val is not None:
logger.info(f"{APPLY_PATCH_FAIL}:\n{val.output.decode(UTF8)}")
raise EvaluationError(
instance_id,
f"{APPLY_PATCH_FAIL}:\n{val.output.decode(UTF8)}",
logger,
)
# Get git diff before running eval script
git_diff_output_before = (
container.exec_run("git -c core.fileMode=false diff", workdir=DOCKER_WORKDIR).output.decode(UTF8).strip() # type: ignore
)
logger.info(f"Git diff before:\n{git_diff_output_before}")
eval_file = Path(log_dir / "eval.sh")
eval_file.write_text(test_spec.eval_script)
logger.info(f"Eval script for {instance_id} written to {eval_file}; copying to container...")
copy_to_container(container, eval_file, PurePosixPath("/eval.sh")) # type: ignore
# Run eval script, write output to logs
test_output, timed_out, total_runtime = exec_run_with_timeout(container, "/bin/bash /eval.sh", timeout)
test_output_path = log_dir / LOG_TEST_OUTPUT
logger.info(f"Test runtime: {total_runtime:_.2f} seconds")
with open(test_output_path, "w") as f:
f.write(test_output)
logger.info(f"Test output for {instance_id} written to {test_output_path}")
if timed_out:
f.write(f"\n\nTimeout error: {timeout} seconds exceeded.")
raise EvaluationError(
instance_id,
f"Test timed out after {timeout} seconds.",
logger,
)
# Get git diff after running eval script (ignore permission changes)
git_diff_output_after = (
container.exec_run("git -c core.fileMode=false diff", workdir=str(DOCKER_WORKDIR)).output.decode(UTF8).strip() # type: ignore
)
# Check if git diff changed after running eval script
logger.info(f"Git diff after:\n{git_diff_output_after}")
if git_diff_output_after != git_diff_output_before:
logger.info("Git diff changed after running eval script")
# Get report from test output
logger.info(f"Grading answer for {instance_id}...")
report = get_eval_report(
test_spec=test_spec,
prediction=pred,
test_log_path=test_output_path,
include_tests_status=True,
)
logger.info(f"report: {report}\n" f"Result for {instance_id}: resolved: {report[instance_id]['resolved']}")
# Write report to report.json
with open(report_path, "w") as f:
f.write(json.dumps(report, indent=4))
return instance_id, report
except EvaluationError as e:
error_msg = traceback.format_exc()
logger.info(error_msg)
print(e)
except BuildImageError as e:
error_msg = traceback.format_exc()
logger.info(error_msg)
print(e)
except Exception as e:
error_msg = f"Error in evaluating model for {instance_id}: {e}\n" f"{traceback.format_exc()}"
logger.error(error_msg)
finally:
# Remove instance container + image, close logger
cleanup_container(client, container, logger)
if rm_image:
remove_image(client, test_spec.instance_image_key, logger)
close_logger(logger)
return
def evaluate(
prediction: Dict[str, Any],
instance: SWEbenchInstance,
cache_level: str,
clean: bool,
force_rebuild: bool,
run_id: str,
timeout: Optional[int],
namespace: Optional[str],
instance_image_tag: str,
rewrite_reports: bool,
):
client = docker.from_env()
test_spec = make_test_spec(instance, namespace=namespace, instance_image_tag=instance_image_tag)
instance_image_ids = {
test_spec.instance_image_key,
}
existing_images = {tag for i in client.images.list(all=True) for tag in i.tags if tag in instance_image_ids}
return run_instance(
test_spec,
prediction,
should_remove(test_spec.instance_image_key, cache_level, clean, existing_images),
force_rebuild,
client,
run_id,
timeout,
rewrite_reports,
)
@@ -0,0 +1,26 @@
# Copyright (c) Microsoft. All rights reserved.
"""Logging utility module for SWE-bench evaluation runs.
This module provides a simple logging utility function that writes evaluation
results and logs to timestamped files organized by run ID and instance ID.
"""
import datetime
import os
def log_for_evaluation(run_id: str, instance_id: str, text: str) -> None:
"""Log a message for evaluation purposes of SWE-Bench.
The format follows the SWE-Bench evaluation framework.
Args:
run_id: The run ID of the evaluation.
instance_id: The instance ID of the evaluation.
text: The text to log.
"""
os.makedirs(f"./logs/{run_id}", exist_ok=True)
current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
with open(f"./logs/{run_id}/{instance_id}", mode="a") as f:
print(f"\n\n{current_time}\n{text}\n", file=f)
+13
View File
@@ -0,0 +1,13 @@
#!/bin/bash
# Read from stdin
input=$(cat)
# Define output file
output_file="/tmp/hook.out"
# Append input followed by two newlines
echo -e "${input}\n\n" >> "$output_file"
# Exit with status 0
exit 0
@@ -0,0 +1,94 @@
{
"hooks": {
"PreToolUse": [
{
"hooks": [
{
"type": "command",
"command": "/tmp/handle_hook.sh"
}
]
}
],
"PostToolUse": [
{
"hooks": [
{
"type": "command",
"command": "/tmp/handle_hook.sh"
}
]
}
],
"Notification": [
{
"hooks": [
{
"type": "command",
"command": "/tmp/handle_hook.sh"
}
]
}
],
"UserPromptSubmit": [
{
"hooks": [
{
"type": "command",
"command": "/tmp/handle_hook.sh"
}
]
}
],
"Stop": [
{
"hooks": [
{
"type": "command",
"command": "/tmp/handle_hook.sh"
}
]
}
],
"SubagentStop": [
{
"hooks": [
{
"type": "command",
"command": "/tmp/handle_hook.sh"
}
]
}
],
"PreCompact": [
{
"hooks": [
{
"type": "command",
"command": "/tmp/handle_hook.sh"
}
]
}
],
"SessionStart": [
{
"hooks": [
{
"type": "command",
"command": "/tmp/handle_hook.sh"
}
]
}
],
"SessionEnd": [
{
"hooks": [
{
"type": "command",
"command": "/tmp/handle_hook.sh"
}
]
}
]
}
}