diff --git a/examples/scaffolding/benchmarks/__main__.py b/examples/scaffolding/benchmarks/__main__.py index 6fc89f696244..d5e5941967b1 100644 --- a/examples/scaffolding/benchmarks/__main__.py +++ b/examples/scaffolding/benchmarks/__main__.py @@ -14,6 +14,7 @@ """ import argparse +import os import sys from tensorrt_llm.scaffolding import TaskMetricsCollector @@ -21,6 +22,7 @@ from .agent_benchmark import async_agent_benchmark, async_burst_agent_benchmark from .benchmark_utils import run_benchmark_in_thread from .chat_benchmark import async_chat_benchmark +from .coder_benchmark import async_coder_benchmark from .multiround_chat_benchmark import async_multiround_chat_benchmark @@ -129,6 +131,68 @@ def parse_arguments(): help="[Agent only] Enable query collector for debugging", ) + # Coder agent parameters + parser.add_argument( + "--enable_coder", + action="store_true", + help="Enable Coder agent benchmark (uses Apiary sandboxes via ApiaryMCPWorker)", + ) + parser.add_argument( + "--coder_concurrency", + type=int, + default=32, + help="Concurrency for Coder agent benchmark (default: 32)", + ) + parser.add_argument( + "--coder_prompt_num", + type=int, + default=8, + help="Number of prompts for Coder benchmark (default: 8)", + ) + parser.add_argument( + "--coder_max_iterations", + type=int, + default=50, + help="Max tool-calling iterations per Coder request (default: 50)", + ) + parser.add_argument( + "--coder_max_connections", + type=int, + default=200, + help="Max concurrent Apiary sandbox connections for Coder (default: 200)", + ) + parser.add_argument( + "--mcp_url", + type=str, + default="http://0.0.0.0:8083/sse", + help="Coder Apiary MCP server URL (default: http://0.0.0.0:8083/sse)", + ) + parser.add_argument( + "--coder_image", + type=str, + default="ubuntu:22.04", + help="Docker image name used for Coder benchmark sandboxes (default: ubuntu:22.04)", + ) + parser.add_argument( + "--apiary_url", + type=str, + default=os.getenv("APIARY_URL", "http://127.0.0.1:8080"), + help="[Coder only] Apiary daemon URL used to register the sandbox image " + "(default: $APIARY_URL or http://127.0.0.1:8080)", + ) + parser.add_argument( + "--apiary_token", + type=str, + default=os.getenv("APIARY_API_TOKEN"), + help="[Coder only] Bearer token for the Apiary daemon (default: $APIARY_API_TOKEN)", + ) + parser.add_argument( + "--coder_rate", + type=float, + default=1.0, + help="[Rate mode] Poisson arrival rate (req/s) for Coder benchmark. Default: 1.0.", + ) + # Burst agent parameters parser.add_argument( "--enable_burst_agent", @@ -456,6 +520,7 @@ def parse_arguments(): # Benchmark registry: (async_func, display_name, flag_name) BENCHMARK_REGISTRY = [ (async_agent_benchmark, "Agent-Benchmark", "enable_normal_agent"), + (async_coder_benchmark, "Coder-Benchmark", "enable_coder"), (async_burst_agent_benchmark, "Burst-Agent-Benchmark", "enable_burst_agent"), (async_chat_benchmark, "Chat-Benchmark", "enable_chat"), (async_multiround_chat_benchmark, "Multiround-Chat-Benchmark", "enable_multiround_chat"), @@ -497,8 +562,8 @@ def main(): if not enabled_benchmarks: print( - "No benchmark enabled. Use --enable_normal_agent, --enable_burst_agent, " - "--enable_chat, or --enable_multiround_chat" + "No benchmark enabled. Use --enable_normal_agent, --enable_coder, " + "--enable_burst_agent, --enable_chat, or --enable_multiround_chat" ) sys.exit(1) diff --git a/examples/scaffolding/benchmarks/coder_benchmark.py b/examples/scaffolding/benchmarks/coder_benchmark.py new file mode 100644 index 000000000000..5650c5d0d93a --- /dev/null +++ b/examples/scaffolding/benchmarks/coder_benchmark.py @@ -0,0 +1,203 @@ +"""Coder agent benchmark for scaffolding benchmarks. + +Runs the Coder agent against Apiary sandboxes with configurable concurrency. +Each concurrent request gets its own isolated sandbox via ApiaryMCPWorker. +""" + +import sys + +from apiary_client import AsyncApiary +from openai import AsyncOpenAI + +from tensorrt_llm.scaffolding import ApiaryMCPWorker, QueryCollector, TRTOpenaiWorker +from tensorrt_llm.scaffolding.benchmark import ScaffoldingBenchRequest, async_scaffolding_benchmark +from tensorrt_llm.scaffolding.contrib.Coder import create_coder_scaffolding_llm +from tensorrt_llm.scaffolding.load_generation_strategy import ( + ConcurrentStrategy, + PoissonRateStrategy, + UniformWarmupStrategy, +) + +from .benchmark_utils import print_benchmark_results, print_lock, shutdown_llm + +DEFAULT_CODER_PROMPTS = [ + "Add comprehensive error handling to all public functions in the project.", + "Implement a thread-safe LRU cache with configurable capacity.", + "Write a CLI tool that converts CSV files to JSON with streaming support.", + "Create a retry decorator with exponential backoff and jitter.", + "Implement a simple key-value store with TTL-based expiration.", + "Add input validation and type checking to the API endpoint handlers.", + "Write a log aggregation utility that merges and deduplicates log entries.", + "Create a configuration loader that supports YAML, JSON, and environment variables.", +] + + +def load_coder_prompts(num_prompts: int) -> list[str]: + """Load prompts for the Coder benchmark from built-in coding tasks.""" + prompts = DEFAULT_CODER_PROMPTS.copy() + if len(prompts) < num_prompts: + original = prompts.copy() + repeat = (num_prompts + len(original) - 1) // len(original) + prompts = [] + for i in range(repeat): + for p in original: + tag = f"[{i}]." if i > 0 else "" + prompts.append(f"{tag}{p}") + return prompts[:num_prompts] + + +async def register_coder_image(args) -> None: + """Register the Coder benchmark image with the Apiary daemon. + + The image used by every benchmark request must be registered before the + MCP server creates sandboxes against it (otherwise Apiary returns 404). + Failures are surfaced as a hard error so the benchmark doesn't silently + produce 502s for every iteration. + """ + image = getattr(args, "coder_image", "ubuntu:22.04") + apiary_url = getattr(args, "apiary_url", "http://127.0.0.1:8080") + apiary_token = getattr(args, "apiary_token", None) + + apiary = AsyncApiary( + apiary_url=apiary_url, + apiary_token=apiary_token, + images=[image], + ) + try: + if not await apiary.health_check(retries=10, interval=1.0): + raise RuntimeError( + f"Apiary daemon at {apiary_url} is not reachable. " + "Start it with `apiary init && apiary daemon --bind ...`." + ) + status = await apiary.load() + if status is not None and image in status.failed: + reason = next( + ( + entry.get("reason") + for entry in status.failed_images + if entry.get("name") == image + ), + "unknown", + ) + raise RuntimeError(f"Failed to register image {image!r} with Apiary: {reason}") + finally: + await apiary.close() + + +async def create_coder_resources(args): + """Create isolated resources for a Coder benchmark run. + + Returns: + Tuple of (llm, mcp_worker, generation_worker) for cleanup. + """ + client = AsyncOpenAI(api_key=args.openai_api_key, base_url=args.base_url) + generation_worker = TRTOpenaiWorker( + client, args.model, getattr(args, "kv_cache_hint_agent", False) + ) + + mcp_url = getattr(args, "mcp_url", "http://0.0.0.0:8083/sse") + max_conns = getattr(args, "coder_max_connections", 200) + mcp_worker = ApiaryMCPWorker(mcp_url, max_connections=max_conns) + + llm = create_coder_scaffolding_llm( + generation_worker, + mcp_worker, + max_tokens=getattr(args, "max_tokens_agent", 65536), + max_iterations=getattr(args, "coder_max_iterations", 50), + max_parallel_requests=getattr(args, "max_parallel_requests", 1024), + enable_statistics=getattr(args, "enable_statistics", False), + ) + + return llm, mcp_worker, generation_worker + + +async def cleanup_coder_resources(llm, mcp_worker): + """Cleanup Coder benchmark resources.""" + await mcp_worker.async_shutdown() + await shutdown_llm(llm) + + +async def run_coder_benchmark_core( + llm, prompts, concurrency, benchmark_name, args, use_poisson_arrival=True +): + """Core Coder benchmark logic. + + Args: + llm: The ScaffoldingLlm instance. + prompts: List of prompts to benchmark. + concurrency: Number of concurrent requests. + benchmark_name: Name for the benchmark (used in output). + args: Command line arguments. + + Returns: + Tuple of (results, requests_start_time, requests_execution_time, total_time). + """ + task_collection_types = {} + requests = [ + ScaffoldingBenchRequest( + prompt=prompt, + scope_params={"image": getattr(args, "coder_image", "ubuntu:22.04")}, + ) + for prompt in prompts + ] + + if use_poisson_arrival and getattr(args, "load_mode", "concurrent") == "rate": + strategy = PoissonRateStrategy( + rate=getattr(args, "coder_rate", 1.0), + random_seed=getattr(args, "rate_seed", 42), + ) + elif getattr(args, "warmup_window", None) is not None: + strategy = UniformWarmupStrategy( + num_requests=len(requests), + warmup_window=args.warmup_window, + max_concurrency=concurrency, + ) + else: + strategy = ConcurrentStrategy(concurrency=concurrency) + print(f" Strategy: {strategy}") + + ( + results, + requests_start_time, + requests_execution_time, + total_time, + ) = await async_scaffolding_benchmark(llm, task_collection_types, requests, strategy=strategy) + + print_benchmark_results( + benchmark_name, + results, + requests_start_time, + requests_execution_time, + total_time, + ) + + if getattr(args, "enable_query_collector", False): + QueryCollector.get_global_info() + with print_lock: + print(f"Query info dumped to query_result.json! ({benchmark_name})") + + return results, requests_start_time, requests_execution_time, total_time + + +async def async_coder_benchmark(args): + """Run the Coder agent benchmark. + + Returns: + Tuple of (results, requests_start_time, requests_execution_time, total_time). + """ + concurrency = getattr(args, "coder_concurrency", 32) + num_prompts = getattr(args, "coder_prompt_num", 8) + + await register_coder_image(args) + + llm, mcp_worker, _ = await create_coder_resources(args) + prompts = load_coder_prompts(num_prompts) + + with print_lock: + print(f"\n[Coder] Starting benchmark with {num_prompts} prompts, concurrency={concurrency}") + sys.stdout.flush() + + try: + return await run_coder_benchmark_core(llm, prompts, concurrency, "Coder", args) + finally: + await cleanup_coder_resources(llm, mcp_worker) diff --git a/examples/scaffolding/contrib/Coder/README.md b/examples/scaffolding/contrib/Coder/README.md new file mode 100644 index 000000000000..9593e7114fa1 --- /dev/null +++ b/examples/scaffolding/contrib/Coder/README.md @@ -0,0 +1,260 @@ +# Scaffolding Coder Agent + +Agentic coding system built on the TensorRT-LLM Scaffolding framework. The Coder agent uses an LLM for reasoning and planning, and executes filesystem and shell operations inside isolated Apiary sandboxes through a dedicated MCP server. + +## Architecture + +```text +LLM server <-> ScaffoldingLlm <-> Coder / SWEBenchCoder + | + v + ApiaryMCPWorker + | + v + examples/.../coder_mcp.py + | + v + apiary_client.ApiarySessionMux + | + v + Apiary daemon / sessions +``` + +Three services are involved: + +- Apiary daemon: manages sandbox sessions, image registry (populated at runtime via HTTP), and command execution +- `coder_mcp.py`: exposes the Coder tool surface over MCP SSE +- LLM server: OpenAI-compatible endpoint used by Scaffolding + +## Tooling Model + +The Coder agent expects these MCP tools: + +- `read_file` +- `list_dir` +- `grep_files` +- `exec` +- `shell` +- `update_plan` +- `think` +- `complete_task` + +File edits go through `shell` (e.g. `sed -i`, `cat <<'EOF' > path` heredocs, `tee`). + +`ApiaryMCPWorker` opens one SSE connection per Scaffolding execution scope, so parallel branches naturally get isolated sandboxes. + +## Prerequisites + +```bash +# Apiary Python bindings (shared with coder_mcp, runners, and SWE-bench helpers) +pip install /path/to/apiary/bindings/python +``` + +For HuggingFace SWE-bench datasets, install the `swebench` extra: + +```bash +pip install '/path/to/apiary/bindings/python[swebench]' +``` + +## Start Apiary + + +The recommended deployment method is the Apiary container shipped under `/apiary/docker-compose.yml`. + +### Container (recommended) + +```bash +cd /path/to/apiary + +# Build and start the container; 8080 is exposed on the host by default. +docker compose up -d +``` + +The container's entrypoint runs `apiary init && apiary daemon --bind 0.0.0.0:8080`, leaving you with an empty pool ready to accept image registrations from clients. + +Useful environment variables (see the Apiary README for the full list): + +| Variable | Default | Purpose | +|---|---|---| +| `APIARY_PORT` | `8080` | Host port to publish | +| `APIARY_BIND` | `0.0.0.0:8080` | Bind address inside the container | +| `APIARY_API_TOKEN` | (empty) | Bearer token for API auth (empty disables auth) | +| `APIARY_MAX_SANDBOXES` | `40` | Pool concurrency cap | +| `APIARY_LAYERS_DIR` | `/var/lib/apiary/layers` | Layer cache (named volume) | +| `APIARY_OVERLAY_DIR` | `/var/lib/apiary/overlays` | Overlay scratch (named volume) | + +Verify from the host: + +```bash +curl -s http://172.17.0.1:8080/healthz # {"status":"ok"} +curl -s http://172.17.0.1:8080/api/v1/status # pool counters + registered_images +``` + +### Native install + +Only needed when Docker is unavailable. Requires Linux 5.11+, cgroups v2 with delegation, and the `uidmap` package. + +```bash +cd /path/to/apiary +cargo build --release + +apiary init --max-sandboxes 40 +apiary daemon --bind 0.0.0.0:8080 +``` + +The Coder runners (`run_coder.py`, `run_swebench.py`, the benchmark) then register the images they need via `POST /api/v1/images` before they dispatch any work. + +## Start the MCP Server + +```bash +python examples/scaffolding/mcp/coder/coder_mcp.py \ + --apiary-url http://172.17.0.1:8080 \ + --default-image ubuntu:22.04 \ + --port 8083 +``` + +Key flags: + +- `--apiary-url` — Apiary daemon URL +- `--apiary-token` — bearer token for daemon auth +- `--mcp-token` — bearer token required on the SSE endpoint +- `--default-image` — fallback Docker image for sandbox sessions when an SSE client omits the `image` query parameter (must already be registered with the daemon) +- `--working-dir` — default sandbox working directory +- `--idle-timeout` — idle session reap timeout in seconds + +Per-request image selection works through the SSE `image` query parameter. `ApiaryMCPWorker.set_scope_params(..., image=...)` is how the Scaffolding runners select the correct sandbox image for each request. The image is expected to be registered with the daemon already; the runners listed below do that for you. + +## Run a Single Coder Task + +```bash +python examples/scaffolding/contrib/Coder/run_coder.py \ + --base_url http://localhost:8000/v1 \ + --model Qwen3/Qwen3-30B-A3B \ + --apiary_url http://172.17.0.1:8080 \ + --mcp_url http://127.0.0.1:8083/sse \ + --image ubuntu:22.04 \ + --prompt "Implement a thread-safe LRU cache in Python" \ + --max_iterations 50 \ + --enable_tracing +``` + +The runner registers `--image` with the Apiary daemon on startup (via `AsyncApiary`) and only dispatches the request once the image is loaded. + +Important flags: + +- `--image`: Docker image used for the request's sandbox (auto-registered) +- `--apiary_url` / `--apiary_token`: How to reach the Apiary daemon for image registration (defaults to `$APIARY_URL` / `$APIARY_API_TOKEN`) +- `--mcp_url`: `coder_mcp.py` SSE endpoint +- `--max_mcp_connections`: Max concurrent SSE / sandbox connections + +## Run SWE-bench + +The runner resolves the SWE-bench image set from the dataset, registers all unique images with the Apiary daemon (with per-image progress logging), and only then starts dispatching agent requests. + +```bash +python examples/scaffolding/contrib/Coder/run_swebench.py \ + --dataset lite \ + --split dev \ + --apiary_url http://172.17.0.1:8080 \ + --base_url http://localhost:8000/v1 \ + --model Qwen3/Qwen3-30B-A3B \ + --mcp_url http://0.0.0.0:8083/sse \ + --max_parallel_requests 16 +``` + +Useful flags: + +- `--apiary_url` / `--apiary_token`: Apiary daemon target +- `--apiary_load_timeout`: Bound the wait for image registration (default: no timeout — large SWE-bench splits can take a while on a cold cache) + +If the daemon is unreachable the runner aborts immediately with a clear error before any LLM work is dispatched. Per-image failures are logged and surfaced as warnings; only those instances are affected — the rest of the batch still runs against the images that did load. + +If you want to pre-load the image set out of band (for example to share a warmed-up daemon across multiple runs), use the `apiary-load-swebench` CLI shipped with `apiary-client[swebench]`: + +```bash +apiary-load-swebench --apiary-url http://172.17.0.1:8080 --dataset lite +``` + +The runner is idempotent: already-loaded images are reported as `alreadypresent` and skip the load pipeline. + +Outputs: + +- `swebench_output/---