chore: import upstream snapshot with attribution

tests/benchmark/benchmark_store.py

@@ -0,0 +1,417 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Benchmarking store performance by writing and querying spans from the store."""
+
+import argparse
+import asyncio
+import os
+import random
+import sys
+import threading
+import time
+from typing import Any, Dict, List, Literal, Optional, Sequence, Set, Tuple, cast
+
+from rich.console import Console
+
+import agentlightning as agl
+from agentlightning.utils.otel import get_tracer
+
+from .utils import flatten_dict, random_dict
+
+console = Console(width=200)
+
+# Minus 10 to leave time for setting up env.
+MAX_RUNTIME_SECONDS = (int(os.getenv("GITHUB_ACTIONS_TIMEOUT_MINUTES", "30")) - 10) * 60
+MAX_STALE_SECONDS = 300
+
+
+class RolloutProgressTracker:
+    """Helper for tracking rollout progress and surfacing stale worker states."""
+
+    def __init__(self, max_stale_seconds: float = MAX_STALE_SECONDS) -> None:
+        self._max_stale_seconds = max_stale_seconds
+        self._last_progress = time.perf_counter()
+
+    def record_progress(self) -> None:
+        self._last_progress = time.perf_counter()
+
+    async def handle_progress(
+        self,
+        *,
+        progress_made: bool,
+        pending_rollout_ids: Sequence[str],
+        store: agl.LightningStore,
+    ) -> None:
+        if progress_made:
+            self.record_progress()
+            return
+        await self._check_for_stale(pending_rollout_ids=pending_rollout_ids, store=store)
+
+    async def _check_for_stale(self, *, pending_rollout_ids: Sequence[str], store: agl.LightningStore) -> None:
+        if not pending_rollout_ids:
+            return
+        elapsed = time.perf_counter() - self._last_progress
+        if elapsed <= self._max_stale_seconds / 2:
+            return
+        console.print(f"Stale rollouts: {pending_rollout_ids}")
+        if elapsed > self._max_stale_seconds:
+            current_workers = await store.query_workers()
+            console.print("Stalled. Current worker status shown below:")
+            for worker in current_workers:
+                console.print(f"  Worker: {worker}", no_wrap=True, overflow="ignore", crop=False)
+            raise RuntimeError("Rollout progress has stalled for too long")
+
+
+def _abort_due_to_timeout() -> None:
+    sys.stderr.write(f"[benchmark] Exiting after exceeding the {MAX_RUNTIME_SECONDS // 60} minute timeout.\n")
+    sys.stderr.flush()
+    os._exit(1)
+
+
+def _start_timeout_guard(timeout_seconds: float) -> threading.Timer:
+    timer = threading.Timer(timeout_seconds, _abort_due_to_timeout)
+    timer.daemon = True
+    timer.start()
+    return timer
+
+
+def generate_attributes() -> Dict[str, Any]:
+    return flatten_dict(
+        random_dict(
+            depth=(1, 3),
+            breadth=(2, 6),
+            key_length=(3, 20),
+            value_length=(5, 300),
+        )
+    )
+
+
+def make_agent(max_rounds: int, sleep_seconds: float) -> agl.LitAgent[str]:
+    @agl.rollout
+    async def agent(task: str, llm: agl.LLM):
+        tracer = get_tracer()
+        rounds = random.randint(1, max_rounds)
+        selected_round = random.randint(0, rounds - 1)
+
+        for i in range(rounds):
+            with tracer.start_as_current_span(f"agent{i}") as span:
+                # Nested Span
+                with tracer.start_as_current_span(f"round{i}_1") as span:
+                    await asyncio.sleep(random.uniform(0.0, sleep_seconds))
+                    span.set_attributes(generate_attributes())
+                    if i == selected_round:
+                        span.set_attribute("task", task)
+
+                # Nested Span
+                with tracer.start_as_current_span(f"round{i}_2") as span:
+                    await asyncio.sleep(random.uniform(0.0, sleep_seconds))
+                    span.set_attributes(generate_attributes())
+
+            if random.uniform(0, 1) < 0.5:
+                agl.emit_reward(random.uniform(0.0, 1.0))
+
+        # Final Span
+        with tracer.start_as_current_span("final") as span:
+            await asyncio.sleep(random.uniform(0.0, sleep_seconds))
+            span.set_attributes(generate_attributes())
+
+        agl.emit_reward(random.uniform(1.0, 2.0))
+
+    return agent
+
+
+def check_spans(spans: Sequence[agl.Span], task: str) -> None:
+    """Check if the spans contain the task."""
+    found_task = any(span.attributes.get("task") == task for span in spans)
+
+    final_reward = agl.find_final_reward(spans)
+    if final_reward is None:
+        raise ValueError("Final reward is not found")
+    if not (final_reward >= 1 and final_reward <= 2):
+        raise ValueError(f"Final reward {final_reward} is not in the range of 1 to 2")
+    if not found_task:
+        raise ValueError(f"Task {task} is not found in the spans")
+
+
+class AlgorithmBatch(agl.Algorithm):
+    def __init__(
+        self,
+        mode: Literal["batch", "batch_partial", "single"],
+        total_tasks: int,
+        batch_size: Optional[int] = None,
+        remaining_tasks: Optional[int] = None,
+        concurrency: Optional[int] = None,
+    ):
+        self.mode = mode
+        self.total_tasks = total_tasks
+        self.batch_size = batch_size
+        self.remaining_tasks = remaining_tasks
+        self.concurrency = concurrency
+
+    async def run(
+        self, train_dataset: Optional[agl.Dataset[Any]] = None, val_dataset: Optional[agl.Dataset[Any]] = None
+    ):
+        if self.mode == "batch":
+            assert self.batch_size is not None
+            await self.algorithm_batch(self.total_tasks, self.batch_size)
+        elif self.mode == "batch_partial":
+            assert self.batch_size is not None
+            assert self.remaining_tasks is not None
+            await self.algorithm_batch_with_completion_threshold(
+                self.total_tasks, self.batch_size, self.remaining_tasks
+            )
+        elif self.mode == "single":
+            assert self.concurrency is not None
+            await self.algorithm_batch_single(self.total_tasks, self.concurrency)
+        else:
+            raise ValueError(f"Invalid mode: {self.mode}")
+
+    async def algorithm_batch(self, total_tasks: int, batch_size: int):
+        """
+        At each time, the algorithm will enqueue a batch of rollouts of size `batch_size`.
+        The algorithm will use wait_for_rollouts to wait for all rollouts to complete.
+        It then checks whether all rollouts are successful and check the spans to ensure the task is found
+        and the last reward is in the range of 1 to 2.
+        After that, the algorithm will enqueue a new batch of new tasks, until the total number of tasks is reached.
+        """
+        store = self.get_store()
+        tracker = RolloutProgressTracker()
+        submitted = 0
+
+        while submitted < total_tasks:
+            print(f"Submitting batch {submitted} of {total_tasks}")
+            batch_count = min(batch_size, total_tasks - submitted)
+            batch_rollouts: List[Tuple[str, str]] = []
+            await store.add_resources(
+                {
+                    "llm": agl.LLM(
+                        endpoint=f"http://localhost:{submitted}/v1",
+                        model=f"test-model-{submitted}",
+                    )
+                }
+            )
+            for _ in range(batch_count):
+                task_name = f"task-{submitted}-generated"
+                rollout = await store.enqueue_rollout(input=task_name, mode="train")
+                batch_rollouts.append((rollout.rollout_id, task_name))
+                submitted += 1
+
+            pending = {rollout_id: task_name for rollout_id, task_name in batch_rollouts}
+            completed_ids: Set[str] = set()
+            tracker.record_progress()
+            while len(completed_ids) < len(batch_rollouts):
+                finished_rollouts = await store.wait_for_rollouts(
+                    rollout_ids=[rollout_id for rollout_id, _ in batch_rollouts],
+                    timeout=0.0,
+                )
+                complete_ids_updated: bool = False
+                for rollout in finished_rollouts:
+                    rollout_id = rollout.rollout_id
+                    if rollout_id in completed_ids:
+                        continue
+                    if rollout.status != "succeeded":
+                        raise RuntimeError(f"Rollout {rollout_id} finished with status {rollout.status}")
+                    spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
+                    check_spans(spans, pending[rollout_id])
+                    completed_ids.add(rollout_id)
+                    complete_ids_updated = True
+
+                unfinished_ids = [rollout_id for rollout_id, _ in batch_rollouts if rollout_id not in completed_ids]
+                await tracker.handle_progress(
+                    progress_made=complete_ids_updated,
+                    pending_rollout_ids=unfinished_ids,
+                    store=store,
+                )
+
+                await asyncio.sleep(5.0)
+
+    async def algorithm_batch_with_completion_threshold(self, total_tasks: int, batch_size: int, remaining_tasks: int):
+        """Different from `algorithm_batch`, this algorithm will use query_rollouts to get rollouts' status.
+        It will enqueue a new batch of new tasks when the number of running rollouts is less than the remaining tasks threshold.
+        """
+        store = self.get_store()
+        tracker = RolloutProgressTracker()
+        submitted = 0
+        completed = 0
+        active_rollouts: Dict[str, str] = {}
+
+        while completed < total_tasks:
+            console.print(f"Completed {completed} of {total_tasks} rollouts")
+            if submitted < total_tasks and len(active_rollouts) < remaining_tasks:
+                batch_count = min(batch_size, total_tasks - submitted)
+                await store.add_resources(
+                    {
+                        "llm": agl.LLM(
+                            endpoint=f"http://localhost:{submitted}/v1",
+                            model=f"test-model-{submitted}",
+                        )
+                    }
+                )
+                for _ in range(batch_count):
+                    task_name = f"task-{submitted}"
+                    rollout = await store.enqueue_rollout(input=task_name, mode="train")
+                    active_rollouts[rollout.rollout_id] = task_name
+                    submitted += 1
+                continue
+
+            if not active_rollouts:
+                await asyncio.sleep(0.01)
+                continue
+
+            rollouts = await store.query_rollouts(rollout_id_in=list(active_rollouts.keys()))
+            newly_completed = 0
+            for rollout in rollouts:
+                rollout_id = rollout.rollout_id
+                if rollout_id not in active_rollouts:
+                    continue
+                if rollout.status in ("queuing", "preparing", "running", "requeuing"):
+                    continue
+                if rollout.status != "succeeded":
+                    raise RuntimeError(f"Rollout {rollout_id} finished with status {rollout.status}")
+                spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
+                check_spans(spans, active_rollouts.pop(rollout_id))
+                completed += 1
+                newly_completed += 1
+
+            await tracker.handle_progress(
+                progress_made=newly_completed > 0,
+                pending_rollout_ids=list(active_rollouts.keys()),
+                store=store,
+            )
+
+            if newly_completed == 0:
+                await asyncio.sleep(5.0)
+
+    async def algorithm_batch_single(self, total_tasks: int, concurrency: int):
+        """Different from `algorithm_batch`, this algorithm will use one async function to enqueue one rollout at a time.
+        The function only cares about the rollout it's currently processing.
+        It waits for the rollouts with `get_rollout_by_id` and check the spans to ensure the rollout is successful.
+        The concurrency is managed via a asyncio semaphore.
+        """
+        store = self.get_store()
+        semaphore = asyncio.Semaphore(concurrency)
+        tracker = RolloutProgressTracker()
+        active_rollouts: Set[str] = set()
+        active_lock = asyncio.Lock()
+
+        async def emit_progress(progress_made: bool) -> None:
+            if progress_made:
+                async with active_lock:
+                    pending_ids = list(active_rollouts)
+                await tracker.handle_progress(progress_made=True, pending_rollout_ids=pending_ids, store=store)
+                return
+            async with active_lock:
+                pending_ids = list(active_rollouts)
+            await tracker.handle_progress(progress_made=False, pending_rollout_ids=pending_ids, store=store)
+
+        async def handle_single(task_index: int) -> None:
+            task_name = f"task-{task_index}"
+            async with semaphore:
+                console.print(f"Submitting task {task_index} of {total_tasks}")
+                await store.add_resources(
+                    {
+                        "llm": agl.LLM(
+                            endpoint=f"http://localhost:{task_index}/v1",
+                            model=f"test-model-{task_index}",
+                        )
+                    }
+                )
+                rollout = await store.enqueue_rollout(input=task_name, mode="train")
+                rollout_id = rollout.rollout_id
+                async with active_lock:
+                    active_rollouts.add(rollout_id)
+                try:
+                    while True:
+                        current = await store.get_rollout_by_id(rollout_id)
+                        if current is not None and current.status in ("failed", "succeeded", "cancelled"):
+                            if current.status != "succeeded":
+                                raise RuntimeError(f"Rollout {rollout_id} finished with status {current.status}")
+                            break
+                        await emit_progress(progress_made=False)
+                        await asyncio.sleep(5.0)
+                    spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
+                    check_spans(spans, task_name)
+                    await emit_progress(progress_made=True)
+                finally:
+                    async with active_lock:
+                        active_rollouts.discard(rollout_id)
+
+        all_tasks = [handle_single(i) for i in range(total_tasks)]
+        await asyncio.gather(*all_tasks)
+
+
+def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Benchmark LightningStore implementations with synthetic rollouts.")
+    parser.add_argument("--store-url", default="http://localhost:4747", help="Lightning Store endpoint base URL.")
+    parser.add_argument(
+        "--mode",
+        choices=("batch", "batch_partial", "single"),
+        default="batch",
+        help="Algorithm mode to exercise different submission patterns.",
+    )
+    parser.add_argument("--total-tasks", type=int, default=128 * 128, help="Total number of rollouts to submit.")
+    parser.add_argument("--batch-size", type=int, default=128, help="Batch size for batch-style modes.")
+    parser.add_argument(
+        "--remaining-tasks",
+        type=int,
+        default=512,
+        help="Target number of in-flight rollouts before submitting more (batch_partial mode).",
+    )
+    parser.add_argument("--concurrency", type=int, default=32, help="Maximum concurrent rollouts for single mode.")
+    parser.add_argument("--n-runners", type=int, default=32, help="Number of runner processes to launch.")
+    parser.add_argument("--max-rounds", type=int, default=10, help="Maximum number of rounds for each rollout.")
+    parser.add_argument("--sleep-seconds", type=float, default=1.0, help="Sleep seconds for each rollout.")
+    parser.add_argument("--debug", action="store_true", help="Enable verbose debug logging.")
+    parser.add_argument("--debug-otel", action="store_true", help="Enable verbose debug logging for OTel.")
+    args = parser.parse_args(argv)
+
+    if args.total_tasks <= 0:
+        parser.error("--total-tasks must be positive")
+    if args.n_runners <= 0:
+        parser.error("--n-runners must be positive")
+    if args.mode in {"batch", "batch_partial"} and (args.batch_size is None or args.batch_size <= 0):
+        parser.error("--batch-size must be positive for batch modes")
+    if args.mode == "batch_partial" and (args.remaining_tasks is None or args.remaining_tasks <= 0):
+        parser.error("--remaining-tasks must be positive for batch_partial mode")
+    if args.mode == "single" and (args.concurrency is None or args.concurrency <= 0):
+        parser.error("--concurrency must be positive for single mode")
+    if args.max_rounds <= 0:
+        parser.error("--max-rounds must be positive")
+    if args.sleep_seconds <= 0:
+        parser.error("--sleep-seconds must be positive")
+
+    return args
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    args = parse_args(argv)
+    agl.setup_logging(
+        "DEBUG" if args.debug else "INFO",
+        submodule_levels={"agentlightning.utils.otel": "DEBUG" if args.debug_otel else "INFO"},
+    )
+    store = agl.LightningStoreClient(args.store_url)
+    timeout_guard = _start_timeout_guard(MAX_RUNTIME_SECONDS)
+    try:
+        trainer = agl.Trainer(
+            store=store,
+            algorithm=AlgorithmBatch(
+                mode=cast(Literal["batch", "batch_partial", "single"], args.mode),
+                total_tasks=args.total_tasks,
+                batch_size=args.batch_size,
+                remaining_tasks=args.remaining_tasks,
+                concurrency=args.concurrency,
+            ),
+            n_runners=args.n_runners,
+            strategy={
+                "type": "cs",
+                "managed_store": False,
+            },
+        )
+        trainer.fit(make_agent(max_rounds=args.max_rounds, sleep_seconds=args.sleep_seconds))
+    finally:
+        timeout_guard.cancel()
+        asyncio.run(store.close())
+
+
+if __name__ == "__main__":
+    main()

tests/benchmark/collection_benchmark.py

@@ -0,0 +1,507 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Collection-level contention benchmarks for Agent Lightning."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import math
+import multiprocessing as mp
+import random
+import threading
+import time
+import uuid
+from contextlib import asynccontextmanager
+from dataclasses import asdict, dataclass
+from multiprocessing.process import BaseProcess
+from pathlib import Path
+from queue import Empty, Queue
+from typing import Any, AsyncContextManager, Callable, Dict, List, Mapping, Sequence
+
+from pymongo import AsyncMongoClient
+from rich.console import Console
+from rich.table import Table
+
+from agentlightning.store.collection.base import LightningCollections
+from agentlightning.store.collection.memory import InMemoryLightningCollections
+from agentlightning.store.collection.mongo import MongoClientPool, MongoLightningCollections
+from agentlightning.types import Rollout, RolloutConfig
+
+console = Console()
+
+DEFAULT_TOTAL_TASKS = 100_000
+DEFAULT_CONCURRENCY = 1_024
+DEFAULT_TASK_PREFIX = "collection-bench"
+MONGO_DEFAULT_DB = "agentlightning_collection_bench"
+
+
+@dataclass
+class WorkerResult:
+    durations: List[float]
+    failures: int
+
+
+@dataclass
+class BenchmarkResult:
+    backend: str
+    name: str
+    total_tasks: int
+    concurrency: int
+    successes: int
+    failures: int
+    duration: float
+    throughput: float
+    avg_latency: float
+    p50_latency: float
+    p95_latency: float
+    p99_latency: float
+    min_latency: float
+    max_latency: float
+    success_rate: float
+    ops_per_worker: float
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Benchmark LightningStore collections without the store server.")
+    parser.add_argument("benchmark", choices=("insert", "dequeue"), help="Benchmarks to run.")
+    parser.add_argument("--backend", choices=("memory", "mongo"), default="memory", help="Collection backend to test.")
+    parser.add_argument("--total-tasks", type=int, default=DEFAULT_TOTAL_TASKS, help="Total operations to run.")
+    parser.add_argument("--concurrency", type=int, default=DEFAULT_CONCURRENCY, help="Number of concurrent workers.")
+    parser.add_argument("--task-prefix", default=DEFAULT_TASK_PREFIX, help="Base prefix for generated workload IDs.")
+    parser.add_argument("--summary-file", help="Optional newline-delimited JSON summary output.")
+    parser.add_argument(
+        "--mongo-uri", default="mongodb://localhost:27017/?replicaSet=rs0", help="Mongo connection URI."
+    )
+    parser.add_argument("--mongo-database", default=MONGO_DEFAULT_DB, help="Mongo database for benchmark artifacts.")
+    return parser.parse_args(argv)
+
+
+def _percentile(values: Sequence[float], percentile: float) -> float:
+    if not values:
+        return 0.0
+    if len(values) == 1:
+        return values[0]
+    rank = (len(values) - 1) * percentile
+    lower = math.floor(rank)
+    upper = math.ceil(rank)
+    if lower == upper:
+        return values[int(rank)]
+    return values[lower] * (upper - rank) + values[upper] * (rank - lower)
+
+
+def _aggregate_results(
+    *,
+    backend: str,
+    name: str,
+    results: Sequence[WorkerResult],
+    concurrency: int,
+    total_tasks: int,
+    duration: float,
+) -> BenchmarkResult:
+    successes = sum(len(result.durations) for result in results)
+    failures = sum(result.failures for result in results)
+    latencies = [lat for result in results for lat in result.durations]
+    throughput = successes / duration if duration > 0 else 0.0
+    avg_latency = (sum(latencies) / len(latencies)) if latencies else 0.0
+    sorted_latencies = sorted(latencies)
+    return BenchmarkResult(
+        backend=backend,
+        name=name,
+        total_tasks=total_tasks,
+        concurrency=concurrency,
+        successes=successes,
+        failures=failures,
+        duration=duration,
+        throughput=throughput,
+        avg_latency=avg_latency,
+        p50_latency=_percentile(sorted_latencies, 0.50),
+        p95_latency=_percentile(sorted_latencies, 0.95),
+        p99_latency=_percentile(sorted_latencies, 0.99),
+        min_latency=sorted_latencies[0] if sorted_latencies else 0.0,
+        max_latency=sorted_latencies[-1] if sorted_latencies else 0.0,
+        success_rate=(successes / (successes + failures)) if (successes + failures) else 0.0,
+        ops_per_worker=(successes / concurrency) if concurrency else 0.0,
+    )
+
+
+def _render_results(results: Sequence[BenchmarkResult]) -> None:
+    if not results:
+        console.print("[yellow]No benchmark results to display.[/yellow]")
+        return
+    table = Table(title="Collection Benchmarks", show_lines=False)
+    table.add_column("Backend")
+    table.add_column("Benchmark")
+    table.add_column("Successes", justify="right")
+    table.add_column("Failures", justify="right")
+    table.add_column("Throughput (req/s)", justify="right")
+    table.add_column("Avg Latency (ms)", justify="right")
+    table.add_column("P95 (ms)", justify="right")
+    table.add_column("P99 (ms)", justify="right")
+    table.add_column("Success Rate", justify="right")
+    for result in results:
+        table.add_row(
+            result.backend,
+            result.name,
+            f"{result.successes:,}",
+            f"{result.failures:,}",
+            f"{result.throughput:,.2f}",
+            f"{result.avg_latency * 1e3:,.2f}",
+            f"{result.p95_latency * 1e3:,.2f}",
+            f"{result.p99_latency * 1e3:,.2f}",
+            f"{result.success_rate * 100:,.2f}%",
+        )
+    console.print(table)
+
+
+def _write_summary(results: Sequence[BenchmarkResult], file_path: Path) -> None:
+    file_path.parent.mkdir(parents=True, exist_ok=True)
+    with file_path.open("a", encoding="utf-8") as handle:
+        for result in results:
+            handle.write(json.dumps(result.to_dict()) + "\n")
+
+
+def _make_rollout(worker_index: int, sequence: int, task_prefix: str) -> Rollout:
+    rollout_id = f"{task_prefix}-ro-{worker_index}-{sequence}-{uuid.uuid4().hex}"
+    current_time = time.time()
+    return Rollout(
+        rollout_id=rollout_id,
+        input={"task": rollout_id},
+        start_time=current_time,
+        end_time=None,
+        mode="train",
+        resources_id=None,
+        status="queuing",
+        config=RolloutConfig(),
+        metadata={},
+    )
+
+
+async def _preload_queue(collections: LightningCollections, total_tasks: int, task_prefix: str) -> None:
+    batch: List[str] = []
+    for idx in range(total_tasks):
+        batch.append(f"{task_prefix}-queue-{idx}")
+        if len(batch) >= 512:
+            async with collections.atomic(mode="rw", labels=["rollout_queue"]) as collections_atomic:
+                await collections_atomic.rollout_queue.enqueue(batch)
+            batch.clear()
+    if batch:
+        async with collections.atomic(mode="rw", labels=["rollout_queue"]) as collections_atomic:
+            await collections_atomic.rollout_queue.enqueue(batch)
+
+
+async def _reset_mongo_database(uri: str, database: str) -> None:
+    client = AsyncMongoClient[Mapping[str, Any]](uri)
+    try:
+        await client.drop_database(database)
+    finally:
+        await client.close()
+
+
+class BaseBenchmark:
+    """Shared control flow for collection benchmarks across backends."""
+
+    def __init__(
+        self, *, backend: str, total_tasks: int, concurrency: int, task_prefix: str, name: str, kind: str
+    ) -> None:
+        self.backend = backend
+        self.total_tasks = total_tasks
+        self.concurrency = concurrency
+        self.task_prefix = task_prefix
+        self.name = name
+        self.kind = kind
+
+    def run(self) -> BenchmarkResult:
+        asyncio.run(self.setup())
+        start = time.perf_counter()
+
+        results = self.spawn_workers(worker_fn=self.worker_entrypoint)
+        duration = time.perf_counter() - start
+        return _aggregate_results(
+            backend=self.backend,
+            name=self.name,
+            results=results,
+            concurrency=self.concurrency,
+            total_tasks=self.total_tasks,
+            duration=duration,
+        )
+
+    def spawn_workers(
+        self,
+        worker_fn: Callable[[int, Any, Any], WorkerResult],
+    ) -> List[WorkerResult]:
+        raise NotImplementedError()
+
+    def worker_entrypoint(self, worker_index: int, task_queue: Any, start_barrier: Any) -> WorkerResult:
+        start_barrier.wait()
+        console.print(f"Worker {worker_index} starting")
+
+        async def _runner() -> WorkerResult:
+            async with self.worker_context() as collections:
+                if self.kind == "insert":
+                    return await insert_worker_async(
+                        collections,
+                        worker_index=worker_index,
+                        task_queue=task_queue,
+                        task_prefix=self.task_prefix,
+                    )
+                if self.kind == "dequeue":
+                    return await dequeue_worker_async(
+                        collections,
+                        worker_index=worker_index,
+                        task_queue=task_queue,
+                    )
+                raise ValueError(f"Unknown benchmark kind: {self.kind}")
+
+        return asyncio.run(_runner())
+
+    def worker_context(self, *args: Any, **kwargs: Any) -> AsyncContextManager[LightningCollections]:
+        """Provide the execution context for the benchmark workers."""
+        raise NotImplementedError()
+
+    async def setup(self) -> None:
+        """Prepare backend-specific state before running workers."""
+        if self.kind == "dequeue":
+            async with self.worker_context() as collections:
+                await _preload_queue(collections, self.total_tasks, self.task_prefix)
+
+
+class MemoryBenchmark(BaseBenchmark):
+
+    def __init__(
+        self,
+        *,
+        total_tasks: int,
+        concurrency: int,
+        task_prefix: str,
+        kind: str,
+    ) -> None:
+        super().__init__(
+            total_tasks=total_tasks,
+            concurrency=concurrency,
+            task_prefix=task_prefix,
+            name=f"collection-{kind}",
+            backend="memory",
+            kind=kind,
+        )
+        self.collections = InMemoryLightningCollections(lock_type="thread")
+
+    def spawn_workers(
+        self,
+        worker_fn: Callable[[int, Any, Any], WorkerResult],
+    ) -> List[WorkerResult]:
+        task_queue: Queue[int] = Queue()
+        for task_id in range(self.total_tasks):
+            task_queue.put(task_id)
+        start_barrier = threading.Barrier(self.concurrency)
+        results: List[WorkerResult | None] = [None] * self.concurrency
+
+        def _thread_target(worker_index: int) -> None:
+            results[worker_index] = worker_fn(worker_index, task_queue, start_barrier)
+
+        threads: List[threading.Thread] = []
+        for worker_index in range(self.concurrency):
+            thread = threading.Thread(target=_thread_target, args=(worker_index,))
+            thread.start()
+            threads.append(thread)
+        for thread in threads:
+            thread.join()
+
+        return [result for result in results if result is not None]
+
+    @asynccontextmanager
+    async def worker_context(self, *args: Any, **kwargs: Any):
+        yield self.collections
+
+
+class MongoBenchmark(BaseBenchmark):
+    def __init__(
+        self,
+        *,
+        total_tasks: int,
+        concurrency: int,
+        task_prefix: str,
+        kind: str,
+        mongo_uri: str,
+        mongo_database: str,
+    ) -> None:
+        super().__init__(
+            total_tasks=total_tasks,
+            concurrency=concurrency,
+            task_prefix=task_prefix,
+            name=f"collection-{kind}",
+            backend="mongo",
+            kind=kind,
+        )
+        self.mongo_uri = mongo_uri
+        self.mongo_database = mongo_database
+        self.partition_id = f"partition-{uuid.uuid4().hex}"
+
+    async def setup(self) -> None:
+        await _reset_mongo_database(self.mongo_uri, self.mongo_database)
+        return await super().setup()
+
+    @asynccontextmanager
+    async def worker_context(self):
+        pool = MongoClientPool[Mapping[str, Any]](mongo_uri=self.mongo_uri)
+        collections = MongoLightningCollections(
+            client_pool=pool,
+            database_name=self.mongo_database,
+            partition_id=self.partition_id,
+            tracker=None,
+        )
+
+        try:
+            yield collections
+        finally:
+            await pool.close()
+
+    def spawn_workers(
+        self,
+        worker_fn: Callable[[int, Any, Any], WorkerResult],
+    ) -> List[WorkerResult]:
+        ctx = mp.get_context("fork")
+        task_queue = ctx.Queue()
+        for task_id in range(self.total_tasks):
+            task_queue.put(task_id)
+        start_barrier = ctx.Barrier(self.concurrency)
+        result_queue = ctx.Queue()
+
+        processes: List[BaseProcess] = []
+        for worker_index in range(self.concurrency):
+            process = ctx.Process(
+                target=_process_worker_target,
+                args=(self, worker_index, task_queue, start_barrier, result_queue),
+            )
+            process.start()
+            processes.append(process)
+
+        collected: List[WorkerResult] = []
+        errors: List[Exception] = []
+        for _ in range(self.concurrency):
+            item = result_queue.get()
+            if isinstance(item, Exception):
+                errors.append(item)
+            else:
+                collected.append(item)
+
+        for process in processes:
+            process.join()
+
+        if errors:
+            raise RuntimeError("One or more worker processes failed") from errors[0]
+
+        return collected
+
+
+def _process_worker_target(
+    benchmark: BaseBenchmark,
+    worker_index: int,
+    task_queue: Any,
+    start_barrier: Any,
+    result_queue: Any,
+) -> None:
+    try:
+        result = benchmark.worker_entrypoint(worker_index, task_queue, start_barrier)
+    except Exception as exc:
+        result_queue.put(exc)
+        raise
+    else:
+        result_queue.put(result)
+
+
+async def insert_worker_async(
+    collections: LightningCollections,
+    *,
+    worker_index: int,
+    task_queue: Any,
+    task_prefix: str,
+) -> WorkerResult:
+    durations: List[float] = []
+    failures = 0
+    while True:
+        try:
+            sequence = task_queue.get_nowait()
+        except Empty:
+            break
+        rollout = _make_rollout(worker_index, sequence, task_prefix)
+        req_start = time.perf_counter()
+        try:
+            async with collections.atomic(mode="rw", labels=["rollouts"]) as collections_atomic:
+                if random.uniform(0, 1) < 0.01:
+                    console.print("Inserting rollout:", rollout.rollout_id)
+                await collections_atomic.rollouts.insert([rollout])
+            durations.append(time.perf_counter() - req_start)
+        except Exception:
+            failures += 1
+    return WorkerResult(durations=durations, failures=failures)
+
+
+async def dequeue_worker_async(
+    collections: LightningCollections,
+    *,
+    worker_index: int,
+    task_queue: Any,
+) -> WorkerResult:
+    del worker_index  # unused but kept for symmetry
+    durations: List[float] = []
+    failures = 0
+    while True:
+        try:
+            task_queue.get_nowait()
+        except Empty:
+            break
+        req_start = time.perf_counter()
+        try:
+            async with collections.atomic(mode="rw", labels=["rollout_queue"]) as collections_atomic:
+                items = await collections_atomic.rollout_queue.dequeue(limit=1)
+                if items and random.uniform(0, 1) < 0.01:
+                    console.print("Dequeued items:", items[0])
+        except Exception:
+            failures += 1
+            continue
+        if not items:
+            break
+        durations.append(time.perf_counter() - req_start)
+    return WorkerResult(durations=durations, failures=failures)
+
+
+def run_benchmark(args: argparse.Namespace, benchmark_kind: str) -> BenchmarkResult:
+    params = {
+        "total_tasks": args.total_tasks,
+        "concurrency": args.concurrency,
+        "task_prefix": args.task_prefix,
+    }
+    if args.backend == "memory":
+        return MemoryBenchmark(kind=benchmark_kind, **params).run()
+
+    mongo_params = {
+        **params,
+        "mongo_uri": args.mongo_uri,
+        "mongo_database": args.mongo_database,
+    }
+    return MongoBenchmark(kind=benchmark_kind, **mongo_params).run()
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    args = parse_args(argv)
+    if args.total_tasks <= 0:
+        raise ValueError("total-tasks must be positive")
+    if args.concurrency <= 0:
+        raise ValueError("concurrency must be positive")
+
+    results: List[BenchmarkResult] = []
+    results.append(run_benchmark(args, args.benchmark))
+
+    _render_results(results)
+
+    if args.summary_file:
+        _write_summary(results, Path(args.summary_file))
+
+
+if __name__ == "__main__":  # pragma: no cover - manual execution
+    main()

tests/benchmark/micro_benchmark.py

@@ -0,0 +1,457 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Micro benchmarks for the store."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import multiprocessing
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional, Sequence
+
+from rich.console import Console
+
+import agentlightning as agl
+from agentlightning.types import EnqueueRolloutRequest, OtelResource, Span, SpanContext, TraceStatus
+from agentlightning.utils.metrics import ConsoleMetricsBackend, MultiMetricsBackend
+from agentlightning.utils.system_snapshot import system_snapshot
+
+from .utils import flatten_dict, random_dict
+
+console = Console()
+
+
+async def _enqueue_rollouts_for_benchmark(store_url: str, *, total_rollouts: int, task_prefix: str) -> None:
+    """Utility that enqueues a fixed number of rollouts for a benchmark."""
+    store = agl.LightningStoreClient(store_url)
+    console.print(f"Enqueuing {total_rollouts} rollouts for {task_prefix} benchmark")
+    try:
+        await store.enqueue_many_rollouts(
+            [EnqueueRolloutRequest(input={"task": f"{task_prefix}-Task-{i}"}) for i in range(total_rollouts)]
+        )
+    finally:
+        await store.close()
+
+
+def _close_store_client(store: agl.LightningStoreClient) -> None:
+    try:
+        asyncio.run(store.close())
+    except Exception:
+        pass
+
+
+def _make_span(rollout_id: str, attempt_id: str, sequence_id: int, name: str, attribute_size: int) -> Span:
+    trace_hex = f"{sequence_id:032x}"
+    span_hex = f"{sequence_id:016x}"
+    return Span(
+        rollout_id=rollout_id,
+        attempt_id=attempt_id,
+        sequence_id=sequence_id,
+        trace_id=trace_hex,
+        span_id=span_hex,
+        parent_id=None,
+        name=name,
+        status=TraceStatus(status_code="OK"),
+        attributes=flatten_dict(
+            random_dict(
+                depth=1,
+                breadth=attribute_size,
+                key_length=(3, 20),
+                value_length=(5, 300),
+            )
+        ),
+        events=[],
+        links=[],
+        start_time=None,
+        end_time=None,
+        context=SpanContext(trace_id=trace_hex, span_id=span_hex, is_remote=False, trace_state={}),
+        parent=None,
+        resource=OtelResource(attributes={}, schema_url=""),
+    )
+
+
+@dataclass
+class BenchmarkSummary:
+    mode: str
+    total_tasks: int
+    successes: int
+    duration: float
+
+    @property
+    def success_rate(self) -> float:
+        if self.total_tasks == 0:
+            return 0.0
+        return self.successes / self.total_tasks
+
+    @property
+    def throughput(self) -> float:
+        if self.duration <= 0:
+            return 0.0
+        return self.successes / self.duration
+
+
+def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Micro benchmarks for the store.")
+    parser.add_argument("--store-url", default="http://localhost:4747", help="Lightning Store endpoint base URL.")
+    parser.add_argument("--summary-file", help="File to append final benchmark summary.")
+    parser.add_argument(
+        "mode",
+        choices=("worker", "dequeue-empty", "dequeue-only", "rollout", "dequeue-update-attempt", "metrics"),
+        help="Mode to exercise different operations (metrics targets MultiMetricsBackend fan-out).",
+    )
+    args = parser.parse_args(argv)
+    return args
+
+
+def _update_worker_task(args: tuple[str, str, str]) -> bool:
+    store_url, worker_id, task_id = args
+    console.print(f"Updating worker {worker_id} for task {task_id}")
+    store = agl.LightningStoreClient(store_url)
+    try:
+        asyncio.run(store.update_worker(worker_id, system_snapshot()))
+        return True
+    except Exception as e:
+        console.print(f"Error updating worker {worker_id} for task {task_id}: {e}")
+        return False
+    finally:
+        _close_store_client(store)
+
+
+def simulate_many_update_workers(store_url: str) -> BenchmarkSummary:
+    """Simulate many update workers."""
+
+    start_time = time.time()
+
+    # Use a multiprocessing pool to update workers.
+    worker_ids = [(f"Worker-{i % 1024}", f"Task-{j}") for i in range(1024) for j in range(10)]
+    with multiprocessing.get_context("fork").Pool(processes=1024) as pool:
+        successful_tasks = pool.map(_update_worker_task, [(store_url, *worker_id) for worker_id in worker_ids])
+
+    end_time = time.time()
+    successes = sum(successful_tasks)
+    duration = end_time - start_time
+    throughput = successes / duration if duration > 0 else 0.0
+    console.print(f"Success rate: {successes / len(worker_ids):.3f}")
+    console.print(f"Time taken: {duration:.3f} seconds")
+    console.print(f"Throughput: {throughput:.3f} workers/second")
+    return BenchmarkSummary(mode="worker", total_tasks=len(worker_ids), successes=successes, duration=duration)
+
+
+def _dequeue_empty_and_update_workers_task(args: tuple[str, str, str]) -> bool:
+    store_url, worker_id, task_id = args
+    console.print(f"Dequeueing empty and updating worker {worker_id} for task {task_id}")
+    store = agl.LightningStoreClient(store_url)
+
+    async def _async_task() -> None:
+        await store.dequeue_rollout(worker_id=worker_id)
+        await store.update_worker(worker_id, system_snapshot())
+
+    try:
+        asyncio.run(_async_task())
+        return True
+    except Exception as e:
+        console.print(f"Error dequeueing empty and updating worker {worker_id} for task {task_id}: {e}")
+        return False
+    finally:
+        _close_store_client(store)
+
+
+def simulate_dequeue_empty_and_update_workers(store_url: str) -> BenchmarkSummary:
+    """Simulate dequeue empty and update workers."""
+    start_time = time.time()
+
+    worker_ids = [(f"Worker-{i % 1024}", f"Task-{j}") for i in range(1024) for j in range(10)]
+    with multiprocessing.get_context("fork").Pool(processes=1024) as pool:
+        successful_tasks = pool.map(
+            _dequeue_empty_and_update_workers_task, [(store_url, *worker_id) for worker_id in worker_ids]
+        )
+
+    end_time = time.time()
+    successes = sum(successful_tasks)
+    duration = end_time - start_time
+    throughput = successes / duration if duration > 0 else 0.0
+    console.print(f"Success rate: {successes / len(worker_ids):.3f}")
+    console.print(f"Time taken: {duration:.3f} seconds")
+    console.print(f"Throughput: {throughput:.3f} workers/second")
+    return BenchmarkSummary(mode="dequeue-empty", total_tasks=len(worker_ids), successes=successes, duration=duration)
+
+
+def _rollout_flow_task(args: tuple[str, int, int]) -> bool:
+    store_url, task_id, spans_per_attempt = args
+    store = agl.LightningStoreClient(store_url)
+
+    async def _async_task() -> None:
+        console.print(f"Starting rollout for task {task_id} with {spans_per_attempt} spans")
+        attempted = await store.start_rollout(input={"task": task_id})
+        rollout_id = attempted.rollout_id
+        attempt_id = attempted.attempt.attempt_id
+        for seq in range(1, spans_per_attempt + 1):
+            console.print(f"Adding span {seq} for task {task_id} with {spans_per_attempt} spans")
+            span = _make_span(
+                rollout_id,
+                attempt_id,
+                task_id * spans_per_attempt + seq,
+                f"micro-span-{seq}",
+                attribute_size=1,
+            )
+            await store.add_span(span)
+        console.print(f"Updating attempt {attempt_id} for task {task_id} with {spans_per_attempt} spans")
+        await store.update_attempt(rollout_id, attempt_id, status="succeeded")
+
+    try:
+        asyncio.run(_async_task())
+        return True
+    except Exception as e:
+        console.print(f"Error running rollout task {task_id}: {e}")
+        return False
+    finally:
+        _close_store_client(store)
+
+
+def simulate_rollout_with_spans(store_url: str, spans_per_attempt: int = 4) -> BenchmarkSummary:
+    """Simulate full rollout lifecycle with spans."""
+    start_time = time.time()
+    task_ids = list(range(1024 * 4))
+    with multiprocessing.get_context("fork").Pool(processes=256) as pool:
+        successful_tasks = pool.map(
+            _rollout_flow_task, [(store_url, task_id, spans_per_attempt) for task_id in task_ids]
+        )
+
+    end_time = time.time()
+    successes = sum(successful_tasks)
+    duration = end_time - start_time
+    throughput = successes / duration if duration > 0 else 0.0
+    console.print(f"Rollout success rate: {successes / len(task_ids):.3f}")
+    console.print(f"Time taken: {duration:.3f} seconds")
+    console.print(f"Throughput: {throughput:.3f} rollouts/second")
+    return BenchmarkSummary(mode="rollout", total_tasks=len(task_ids), successes=successes, duration=duration)
+
+
+def _dequeue_only_task(args: tuple[str, str, str]) -> bool:
+    store_url, worker_id, task_id = args
+    console.print(f"[Dequeue-Only Task {task_id}] Dequeueing rollout for worker {worker_id}")
+    store = agl.LightningStoreClient(store_url)
+
+    async def _async_task() -> bool:
+        attempted = await store.dequeue_rollout()  # no worker_id
+        if attempted is None:
+            console.print(f"[Dequeue-Only Task {task_id}] No rollout available to dequeue")
+            return False
+        return True
+
+    try:
+        return asyncio.run(_async_task())
+    except Exception as e:
+        console.print(f"Error dequeueing only worker {worker_id} for task {task_id}: {e}")
+        return False
+    finally:
+        _close_store_client(store)
+
+
+def dequeue_rollouts(store_url: str) -> BenchmarkSummary:
+    """Benchmark simple dequeues without any additional mutations."""
+    start_time = time.time()
+    total_workers = 512
+    attempts_per_worker = 16
+    total_rollouts = total_workers * attempts_per_worker
+
+    asyncio.run(_enqueue_rollouts_for_benchmark(store_url, total_rollouts=total_rollouts, task_prefix="DequeueOnly"))
+
+    worker_jobs = [
+        (f"Worker-{worker_idx}-Attempt-{attempt_idx}", f"Task-{attempt_idx * total_workers + worker_idx}")
+        for worker_idx in range(total_workers)
+        for attempt_idx in range(attempts_per_worker)
+    ]
+    with multiprocessing.get_context("fork").Pool(processes=total_workers) as pool:
+        successful_tasks = pool.map(
+            _dequeue_only_task, [(store_url, worker_id, task_id) for worker_id, task_id in worker_jobs]
+        )
+
+    async def _query_remaining_rollouts() -> List[str]:
+        store = agl.LightningStoreClient(store_url)
+        try:
+            remaining_rollouts = await store.query_rollouts(status_in=["queuing"])
+            return [item.rollout_id for item in remaining_rollouts]
+        finally:
+            await store.close()
+
+    end_time = time.time()
+    remaining_rollouts = asyncio.run(_query_remaining_rollouts())
+    successes = sum(successful_tasks)
+    duration = end_time - start_time
+    throughput = successes / duration if duration > 0 else 0.0
+    console.print(f"Remaining rollouts: {remaining_rollouts}")
+    console.print(f"Remaining rollouts count: {len(remaining_rollouts)}")
+    console.print(f"Dequeue-only success rate: {successes / len(worker_jobs):.3f}")
+    console.print(f"Time taken: {duration:.3f} seconds")
+    console.print(f"Throughput: {throughput:.3f} rollouts/second")
+    return BenchmarkSummary(mode="dequeue-only", total_tasks=len(worker_jobs), successes=successes, duration=duration)
+
+
+def _dequeue_and_update_attempt_task(args: tuple[str, str, str, int]) -> bool:
+    store_url, worker_id, task_id, spans_per_attempt = args
+    console.print(f"Dequeueing and update attempt with worker {worker_id} for task {task_id}")
+    store = agl.LightningStoreClient(store_url)
+
+    async def _async_task() -> bool:
+        console.print(f"[Task {task_id}] Dequeueing rollout")
+        attempted = await store.dequeue_rollout(worker_id=worker_id)
+        if attempted is None:
+            console.print(f"[Task {task_id}] No rollout available to dequeue")
+            return False
+        console.print(f"[Task {task_id}] Retrieving span sequence IDs")
+        sequence_ids = await store.get_many_span_sequence_ids(
+            [(attempted.rollout_id, attempted.attempt.attempt_id) for _ in range(spans_per_attempt)]
+        )
+        if len(sequence_ids) != spans_per_attempt:
+            console.print(
+                f"[Task {task_id}] Unable to retrieve enough span sequence IDs: "
+                f"expected={spans_per_attempt} got={len(sequence_ids)}"
+            )
+            return False
+        console.print(f"[Task {task_id}] Adding {spans_per_attempt} spans")
+        spans = [
+            _make_span(
+                attempted.rollout_id,
+                attempted.attempt.attempt_id,
+                sequence_id,
+                f"micro-span-{sequence_id}",
+                attribute_size=32,
+            )
+            for sequence_id in sequence_ids
+        ]
+        stored_spans = await store.add_many_spans(spans)
+        if len(stored_spans) != len(spans):
+            console.print(
+                f"[Task {task_id}] Only stored {len(stored_spans)}/{len(spans)} spans for "
+                f"rollout_id={attempted.rollout_id} attempt_id={attempted.attempt.attempt_id}"
+            )
+            return False
+        console.print(
+            f"[Task {task_id}] Updating attempt to succeeded: rollout_id={attempted.rollout_id} "
+            f"attempt_id={attempted.attempt.attempt_id}"
+        )
+        await store.update_attempt(attempted.rollout_id, attempted.attempt.attempt_id, status="succeeded")
+        return True
+
+    try:
+        return asyncio.run(_async_task())
+    except Exception as e:
+        console.print(f"Error dequeueing and updating worker {worker_id} for task {task_id}: {e}")
+        return False
+    finally:
+        _close_store_client(store)
+
+
+def dequeue_and_update_attempts(store_url: str, spans_per_attempt: int = 4) -> BenchmarkSummary:
+    """Simulate dequeueing rollouts and updating attempts with spans."""
+    start_time = time.time()
+    total_workers = 512
+    attempts_per_worker = 16
+    total_rollouts = total_workers * attempts_per_worker
+
+    asyncio.run(_enqueue_rollouts_for_benchmark(store_url, total_rollouts=total_rollouts, task_prefix="Dequeue"))
+
+    worker_jobs = [
+        (f"Worker-{worker_idx}-Attempt-{attempt_idx}", f"Task-{attempt_idx * total_workers + worker_idx}")
+        for worker_idx in range(total_workers)
+        for attempt_idx in range(attempts_per_worker)
+    ]
+    with multiprocessing.get_context("fork").Pool(processes=total_workers) as pool:
+        successful_tasks = pool.map(
+            _dequeue_and_update_attempt_task,
+            [(store_url, worker_id, task_id, spans_per_attempt) for worker_id, task_id in worker_jobs],
+        )
+
+    end_time = time.time()
+    successes = sum(successful_tasks)
+    duration = end_time - start_time
+    throughput = successes / duration if duration > 0 else 0.0
+    console.print(f"Dequeue and update attempt success rate: {successes / len(worker_jobs):.3f}")
+    console.print(f"Time taken: {duration:.3f} seconds")
+    console.print(f"Throughput: {throughput:.3f} rollouts/second")
+    return BenchmarkSummary(
+        mode="dequeue-update-attempt", total_tasks=len(worker_jobs), successes=successes, duration=duration
+    )
+
+
+def benchmark_multi_metrics_backend(iterations: int = 10_000_000) -> BenchmarkSummary:
+    """Benchmark MultiMetricsBackend fan-out cost."""
+
+    console.print(f"Benchmarking MultiMetricsBackend for {iterations} iterations (2 metric ops per iteration)")
+
+    agl.setup_logging()
+
+    console_backend = ConsoleMetricsBackend(window_seconds=0.5, log_interval_seconds=0.1, group_level=None)
+    console_backend_secondary = ConsoleMetricsBackend(
+        window_seconds=None, log_interval_seconds=1_000_000.0, group_level=None
+    )
+    backend = MultiMetricsBackend([console_backend, console_backend_secondary])
+
+    backend.register_counter("benchmark.metrics.counter", label_names=["worker"])
+    backend.register_histogram(
+        "benchmark.metrics.latency",
+        label_names=["worker"],
+        buckets=(0.001, 0.005, 0.05, 0.5, 1.0),
+    )
+    labels = {"worker": "benchmark"}
+
+    async def _exercise_metrics() -> None:
+        for i in range(iterations):
+            await backend.inc_counter("benchmark.metrics.counter", labels=labels)
+            await backend.observe_histogram(
+                "benchmark.metrics.latency",
+                value=(i % 100) / 100.0,
+                labels=labels,
+            )
+
+    start_time = time.time()
+    asyncio.run(_exercise_metrics())
+    duration = time.time() - start_time
+    total_ops = iterations * 2
+    throughput = total_ops / duration if duration > 0 else 0.0
+
+    console.print(f"Executed {total_ops} metric updates in {duration:.3f}s ({throughput:.1f} ops/s)")
+    return BenchmarkSummary(mode="metrics", total_tasks=total_ops, successes=total_ops, duration=duration)
+
+
+def record_summary(summary: BenchmarkSummary, summary_file: Optional[str]) -> None:
+    message = (
+        f"[summary] mode={summary.mode} success_rate={summary.success_rate:.3f} "
+        f"throughput={summary.throughput:.3f} ops/s duration={summary.duration:.3f}s "
+        f"success={summary.successes}/{summary.total_tasks}"
+    )
+    console.print(message)
+    if summary_file:
+        path = Path(summary_file)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with path.open("a", encoding="utf-8") as fh:
+            fh.write(message + "\n")
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    args = parse_args(argv)
+    if args.mode == "worker":
+        summary = simulate_many_update_workers(args.store_url)
+    elif args.mode == "dequeue-empty":
+        summary = simulate_dequeue_empty_and_update_workers(args.store_url)
+    elif args.mode == "dequeue-only":
+        summary = dequeue_rollouts(args.store_url)
+    elif args.mode == "rollout":
+        summary = simulate_rollout_with_spans(args.store_url)
+    elif args.mode == "dequeue-update-attempt":
+        summary = dequeue_and_update_attempts(args.store_url)
+    elif args.mode == "metrics":
+        summary = benchmark_multi_metrics_backend()
+    else:
+        raise ValueError(f"Invalid mode: {args.mode}")
+    record_summary(summary, args.summary_file)
+    if summary.success_rate < 1.0:
+        raise ValueError(f"Benchmark failed with success rate {summary.success_rate:.3f}")
+
+
+if __name__ == "__main__":
+    main()

tests/benchmark/utils.py

@@ -0,0 +1,133 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Generating random test data for benchmarking."""
+
+import random
+import string
+from typing import Any, Callable, Dict, Optional, Tuple, Union, cast
+
+
+def random_string(length: int, *, alphabet: Optional[str] = None) -> str:
+    """
+    Generate a random string of fixed length.
+
+    Args:
+        length: Length of the generated string.
+        alphabet: Optional character set to draw from. If None, uses [A-Za-z0-9].
+    """
+    if length < 0:
+        raise ValueError("String length cannot be negative.")
+
+    alphabet = alphabet or (string.ascii_letters + string.digits)
+    return "".join(random.choices(alphabet, k=length))
+
+
+def _resolve_param(value: Union[int, Tuple[int, int]], name: str) -> int:
+    """
+    Convert parameter into a concrete integer.
+    If value is an int, return it.
+    If value is a tuple, interpret it as (low, high) and sample uniformly.
+    """
+    if isinstance(value, int):
+        if value < 0:
+            raise ValueError(f"{name} cannot be negative.")
+        return value
+
+    if (
+        isinstance(value, tuple)  # type: ignore
+        and len(value) == 2
+        and isinstance(value[0], int)  # type: ignore
+        and isinstance(value[1], int)  # type: ignore
+    ):
+        low, high = value
+        if low < 0 or high < 0:
+            raise ValueError(f"{name} range cannot contain negative values.")
+        if low > high:
+            raise ValueError(f"{name} tuple must be (low, high) with low <= high.")
+        return random.randint(low, high)
+
+    raise TypeError(f"{name} must be an int or a 2-tuple of ints.")
+
+
+def default_value_factory(length: int) -> str:
+    """Default value factory for generating string payloads."""
+    return random_string(length)
+
+
+def random_dict(
+    *,
+    depth: Union[int, Tuple[int, int]],
+    breadth: Union[int, Tuple[int, int]],
+    key_length: Union[int, Tuple[int, int]],
+    value_length: Union[int, Tuple[int, int]],
+    value_factory: Optional[Callable[[int], Any]] = None,
+) -> Dict[str, Any]:
+    """
+    Generate a nested dictionary with configurable depth, breadth, and
+    value sizes. Integer or (low, high) tuples are supported for
+    all structural parameters.
+
+    Args:
+        depth: Number of nested levels or a tuple specifying a range.
+        breadth: Number of keys per level or a tuple range.
+        key_length: Length of each key or a tuple range.
+        value_length: Length of each value or a tuple range.
+        value_factory: Function mapping `value_length` → value.
+
+    Returns:
+        A nested dictionary of arbitrary size.
+    """
+    # Default factory
+    if value_factory is None:
+        value_factory = random_string
+
+    def build(level: int) -> Dict[str, Any]:
+        # For each level, breadth/key/value lengths may vary, so draw fresh each time
+        current_breadth = _resolve_param(breadth, "breadth")
+
+        if current_breadth < 0:
+            raise ValueError("Breadth cannot be negative.")
+
+        target_depth = depth if isinstance(depth, int) else _resolve_param((level, depth[1]), "depth")
+
+        if level == target_depth:
+            # leaf nodes
+            return {
+                random_string(_resolve_param(key_length, "key_length")): value_factory(
+                    _resolve_param(value_length, "value_length")
+                )
+                for _ in range(current_breadth)
+            }
+
+        # nested nodes
+        return {
+            random_string(_resolve_param(key_length, "key_length")): build(level + 1) for _ in range(current_breadth)
+        }
+
+    return build(1)
+
+
+def flatten_dict(d: Dict[str, Any], prefix: str = "") -> Dict[str, Any]:
+    """Flatten a nested dictionary into a single level dictionary. Keys are joined by dots."""
+
+    result: Dict[str, Any] = {}
+    for key, value in d.items():
+        if isinstance(value, dict):
+            result.update(flatten_dict(cast(Dict[str, Any], value), f"{prefix}.{key}" if prefix else key))
+        else:
+            result[f"{prefix}.{key}" if prefix else key] = value
+    return result
+
+
+if __name__ == "__main__":
+    # Example usage
+    import json
+
+    structured_dict = random_dict(
+        depth=(1, 3),
+        breadth=(2, 6),
+        key_length=(3, 20),
+        value_length=(5, 300),
+    )
+
+    print(json.dumps(flatten_dict(structured_dict), indent=2))