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
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:
@@ -0,0 +1,112 @@
|
||||
# Training Claude Code with Agent-lightning
|
||||
|
||||
[](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.
|
||||
@@ -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()
|
||||
@@ -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
@@ -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"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user