[None][feat] Enable agent serving evaluation via trace-replay on Scaffolding

examples/scaffolding/benchmarks/__main__.py

@@ -14,13 +14,15 @@
 """
 
 import argparse
+import os
 import sys
 
 from tensorrt_llm.scaffolding import TaskMetricsCollector
 
 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)
 

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)