.

.dockerignore

@@ -0,0 +1,14 @@
+.venv
+**/.venv
+__pycache__
+.git
+.gitignore
+**/node_modules
+dist
+build
+.env
+docker
+.pytest_cache
+.vscode
+**/*.log
+examples/**/data

.github/workflows/tests-full.yml

@@ -82,34 +82,50 @@ jobs:
       - name: Build dashboard
         run: cd dashboard && npm run build
 
-      - name: Start MongoDB container
+      - name: Setup Docker environments
         run: |
           set -euo pipefail
-          cat /etc/security/limits.conf
-          docker run -d \
-            --name mongodb-test \
-            --ulimit nofile=65535:65535 \
-            -p 27017:27017 \
-            mongo:8.2 \
-            --replSet test-rs
 
-          # Wait for mongod to come up
-          for i in $(seq 1 30); do
-            if docker exec mongodb-test mongosh --quiet --eval 'db.runCommand({ ping: 1 })' >/dev/null 2>&1; then
-              echo "Mongo is up"
-              break
+          cd docker
+
+          # Setup data directories
+          ./setup.sh
+
+          # Start Dockers
+          docker compose -f compose.mongo.yml up -d
+
+          SERVICE_NAME=mongo
+          TIMEOUT=60   # seconds
+          SLEEP=2
+
+          cid="$(docker compose -f compose.mongo.yml ps -q "$SERVICE_NAME")"
+          if [ -z "$cid" ]; then
+            echo "Service $SERVICE_NAME is not running"
+            exit 1
+          fi
+
+          echo "Waiting for $SERVICE_NAME to become healthy..."
+          end=$((SECONDS + TIMEOUT))
+
+          while [ "$SECONDS" -lt "$end" ]; do
+            status="$(docker inspect -f '{{.State.Health.Status}}' "$cid")"
+            echo "Current status: $status"
+
+            if [ "$status" = "healthy" ]; then
+              echo "$SERVICE_NAME is healthy ✅"
+              exit 0
+            elif [ "$status" = "unhealthy" ]; then
+              echo "$SERVICE_NAME is unhealthy ❌"
+              docker logs "$cid" || true
+              exit 1
             fi
-            echo "Waiting for Mongo..."
-            sleep 2
+
+            sleep "$SLEEP"
           done
 
-          # Init replica set (simple single-node)
-          docker exec mongodb-test mongosh --quiet --eval '
-            rs.initiate({
-              _id: "test-rs",
-              members: [{ _id: 0, host: "localhost:27017" }]
-            })
-          '
+          echo "Timed out waiting for $SERVICE_NAME to become healthy after ${TIMEOUT}s"
+          docker logs "$cid" || true
+          exit 1
         shell: bash
 
       - name: Launch LiteLLM Proxy
@@ -126,7 +142,7 @@ jobs:
           PYTEST_ADDOPTS: "--color=yes"
           OPENAI_BASE_URL: http://localhost:12306/
           OPENAI_API_KEY: dummy
-          AGL_TEST_MONGO_URI: mongodb://localhost:27017/?replicaSet=test-rs
+          AGL_TEST_MONGO_URI: mongodb://localhost:27017/?replicaSet=rs0
 
 
   minimal-examples:

.gitignore

@@ -213,3 +213,6 @@ agentlightning/dashboard/**/*.css
 agentlightning/dashboard/**/*.js
 agentlightning/dashboard/**/*.html
 agentlightning/dashboard/**/*.svg
+
+# Docker data
+docker/data/

agentlightning/cli/store.py

@@ -18,6 +18,7 @@ logger = logging.getLogger(__name__)
 
 def main(argv: Iterable[str] | None = None) -> int:
     parser = argparse.ArgumentParser(description="Run a LightningStore server")
+    parser.add_argument("--host", default="0.0.0.0", help="Host to bind the server to")
     parser.add_argument("--port", type=int, default=4747, help="Port to run the server on")
     parser.add_argument(
         "--cors-origin",
@@ -31,17 +32,60 @@ def main(argv: Iterable[str] | None = None) -> int:
         choices=["DEBUG", "INFO", "WARNING", "ERROR"],
         help="Configure the logging level for the store.",
     )
+    parser.add_argument(
+        "--prometheus",
+        action="store_true",
+        help="Enable Prometheus metrics.",
+    )
+    parser.add_argument(
+        "--n-workers",
+        default=1,
+        type=int,
+        help=(
+            "Number of workers to run in the server. When it's greater than 1, the server will be run using `mp` launch mode. "
+            "Only applicable for zero-copy stores such as MongoDB backend."
+        ),
+    )
+
+    parser.add_argument(
+        "--backend",
+        choices=["memory", "mongo"],
+        default="memory",
+        help="Backend to use for the store.",
+    )
+    parser.add_argument(
+        "--mongo-uri",
+        default="mongodb://localhost:27017/?replicaSet=rs0",
+        help="MongoDB URI to use for the store. Applicable only if --backend is 'mongo'.",
+    )
+
     args = parser.parse_args(list(argv) if argv is not None else None)
 
     setup_logging(args.log_level)
 
-    store = InMemoryLightningStore()
+    if args.backend == "memory":
+        store = InMemoryLightningStore()
+    elif args.backend == "mongo":
+        from agentlightning.store.mongo import MongoLightningStore
+
+        store = MongoLightningStore(client=args.mongo_uri)
+    else:
+        raise ValueError(f"Invalid backend: {args.backend}")
+
+    if args.n_workers > 1:
+        logger.info(f"Running the server using `mp` launch mode with {args.n_workers} workers.")
+        launch_mode = "mp"
+    else:
+        logger.info("Running the server using `asyncio` launch mode.")
+        launch_mode = "asyncio"
     server = LightningStoreServer(
         store,
-        host="0.0.0.0",
+        host=args.host,
         port=args.port,
         cors_allow_origins=args.cors_origins,
-        launch_mode="asyncio",
+        launch_mode=launch_mode,
+        prometheus=args.prometheus,
+        n_workers=args.n_workers,
     )
     try:
         asyncio.run(server.run_forever())

agentlightning/store/client_server.py

@@ -219,6 +219,8 @@ class LightningStoreServer(LightningStore):
             which runs the server in a separate thread.
         launcher_args: The arguments to use for the server launcher.
             It's not allowed to set `host`, `port`, `launch_mode` together with `launcher_args`.
+        n_workers: The number of workers to run in the server. Only applicable for `mp` launch mode.
+        prometheus: Whether to enable Prometheus metrics.
     """
 
     def __init__(
@@ -229,6 +231,8 @@ class LightningStoreServer(LightningStore):
         cors_allow_origins: Sequence[str] | str | None = None,
         launch_mode: LaunchMode = "thread",
         launcher_args: PythonServerLauncherArgs | None = None,
+        n_workers: int = 1,
+        prometheus: bool = False,
     ):
         super().__init__()
         self.store = store
@@ -265,6 +269,7 @@ class LightningStoreServer(LightningStore):
             app=self.app,
             args=self.launcher_args,
         )
+        self._prometheus = prometheus
 
         self._lock: threading.Lock = threading.Lock()
         self._cors_allow_origins = self._normalize_cors_origins(cors_allow_origins)
@@ -309,6 +314,7 @@ class LightningStoreServer(LightningStore):
         return {
             "launcher_args": self.launcher_args,
             "server_launcher": self.server_launcher,
+            "_prometheus": self._prometheus,
             "_owner_pid": self._owner_pid,
         }
 
@@ -326,6 +332,7 @@ class LightningStoreServer(LightningStore):
         self.store = None
         self.launcher_args = state["launcher_args"]
         self.server_launcher = state["server_launcher"]
+        self._prometheus = state["_prometheus"]
         self._owner_pid = state["_owner_pid"]
         self._cors_allow_origins = state.get("_cors_allow_origins")
         self._client = None
@@ -407,6 +414,11 @@ class LightningStoreServer(LightningStore):
     def _setup_routes(self):
         """Set up FastAPI routes for all store operations."""
         assert self.app is not None
+        api = APIRouter(prefix=API_V1_PREFIX)
+
+        # The outermost-layer of monitoring
+        if self._prometheus:
+            self._setup_prometheus(api=api, app=self.app)
 
         @self.app.middleware("http")
         async def _app_exception_handler(  # pyright: ignore[reportUnusedFunction]
@@ -461,8 +473,6 @@ class LightningStoreServer(LightningStore):
             )
             return response
 
-        api = APIRouter(prefix=API_V1_PREFIX)
-
         def _validate_paginated_request(
             request: Union[
                 QueryRolloutsRequest,
@@ -720,6 +730,58 @@ class LightningStoreServer(LightningStore):
         # Finally, mount the dashboard assets
         self._setup_dashboard()
 
+    def _setup_prometheus(self, api: APIRouter, app: FastAPI):
+        """Setup Prometheus metrics endpoints."""
+        try:
+            from prometheus_client import (
+                CONTENT_TYPE_LATEST,
+                Counter,
+                Histogram,
+                generate_latest,
+            )
+        except ImportError:
+            raise ImportError(
+                "Prometheus client is not installed. Please either install it or set prometheus to False."
+            )
+
+        HTTP_REQUESTS = Counter(
+            "http_requests_total",
+            "Total HTTP requests",
+            ["method", "path", "status_code"],
+        )
+
+        # TODO: For multi-process scenarios, should use prometheus_client.multiprocess mode.
+        HTTP_LATENCY = Histogram(
+            "http_request_duration_seconds",
+            "Latency of HTTP requests",
+            ["method", "path"],
+            buckets=[0.001, 0.005, 0.01, 0.05, 0.1, 0.25, 0.5, 1, 2, 5, 10],
+        )
+
+        @app.middleware("http")
+        async def prometheus_http_middleware(  # pyright: ignore[reportUnusedFunction]
+            request: Request, call_next: Callable[[Request], Awaitable[Response]]
+        ) -> Response:
+            start = time.perf_counter()
+            response = await call_next(request)
+            elapsed = time.perf_counter() - start
+
+            path = request.url.path
+            method = request.method
+            status = response.status_code
+
+            HTTP_REQUESTS.labels(method, path, status).inc()
+            HTTP_LATENCY.labels(method, path).observe(elapsed)
+
+            return response
+
+        @api.get("/prometheus")
+        async def prometheus_metrics():  # pyright: ignore[reportUnusedFunction]
+            return Response(
+                content=generate_latest(),
+                media_type=CONTENT_TYPE_LATEST,
+            )
+
     def _setup_otlp(self, api: APIRouter):
         """Setup OTLP endpoints."""
 
@@ -794,6 +856,10 @@ class LightningStoreServer(LightningStore):
         - In the owner process: delegate to the in-process store.
         - In a different process: delegate to a HTTP client talking to the server.
         """
+        # If the store is zero-copy, we can just call the method directly.
+        if self.store is not None and self.store.capabilities.get("zero_copy", False):
+            return await getattr(self.store, method_name)(*args, **kwargs)
+
         if os.getpid() == self._owner_pid:
             if method_name == "wait_for_rollouts":
                 # wait_for_rollouts can block for a long time; avoid holding the lock

docker/Dockerfile.dev

@@ -0,0 +1,39 @@
+# Use a Python image with uv pre-installed
+FROM ghcr.io/astral-sh/uv:python3.12-bookworm
+
+# Setup a non-root user
+RUN groupadd --system --gid 999 nonroot \
+ && useradd --system --gid 999 --uid 999 --create-home nonroot
+
+# Install the project into `/app`
+WORKDIR /app
+
+# Enable bytecode compilation
+ENV UV_COMPILE_BYTECODE=1
+
+# Copy from the cache instead of linking since it's a mounted volume
+ENV UV_LINK_MODE=copy
+
+# Ensure installed tools can be executed out of the box
+ENV UV_TOOL_BIN_DIR=/usr/local/bin
+
+# Install the project's dependencies using the lockfile and settings
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,source=uv.lock,target=uv.lock \
+    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+    uv sync --locked --no-install-project --group dev --extra mongo --group core-stable
+
+# Then, add the rest of the project source code and install it
+# Installing separately from its dependencies allows optimal layer caching
+COPY . /app
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv sync --locked --group dev --extra mongo --group core-stable
+
+# Place executables in the environment at the front of the path
+ENV PATH="/app/.venv/bin:$PATH"
+
+# Reset the entrypoint, don't invoke `uv`
+ENTRYPOINT []
+
+# Use the non-root user to run our application
+USER nonroot

docker/compose.mongo.yml

@@ -0,0 +1,23 @@
+# Docker-compose file to launch a MongoDB server for development.
+# It's used to test the MongoDB store implementation.
+
+services:
+
+  mongo:
+    image: mongo:latest
+    ulimits:
+      nofile:
+        soft: 65535
+        hard: 65535
+    ports:
+      - "27017:27017"
+    command: ["mongod", "--bind_ip_all", "--replSet", "rs0"]
+    volumes:
+      - ../scripts/mongodb_init_rs_host.js:/docker-entrypoint-initdb.d/init-rs.js:ro
+      - ./data/mongo-host:/data/db
+    healthcheck:
+      test: ["CMD", "mongosh", "--eval", "db.adminCommand('ping')"]
+      interval: 10s
+      timeout: 5s
+      retries: 5
+      start_period: 30s

docker/compose.prometheus-memory-store.yml

@@ -0,0 +1,33 @@
+services:
+
+  app:
+    extends:
+      file: compose.store.yml
+      service: app
+
+    command: agl store --host 0.0.0.0 --port 4747 --prometheus --backend memory
+
+  node-exporter:
+    image: prom/node-exporter:latest
+    # In CI you might not have full /proc, but this is OK for container-level stats
+    pid: "host"
+    network_mode: "service:app" # share network with app for simplicity
+    command:
+      - '--path.rootfs=/host'
+    volumes:
+      - '/:/host:ro,rslave'
+
+  prometheus:
+    image: prom/prometheus:latest
+    command:
+      - '--config.file=/etc/prometheus/prometheus.yml'
+      - '--storage.tsdb.path=/prometheus'
+      - '--storage.tsdb.retention.time=1h'
+    volumes:
+      - ./prometheus.memory-store.yml:/etc/prometheus/prometheus.yml:ro
+      - ./data/prometheus:/prometheus
+    depends_on:
+      - app
+      - node-exporter
+    ports:
+      - "9090:9090"

docker/compose.prometheus-mongo-store.yml

@@ -0,0 +1,62 @@
+services:
+
+  mongo:
+    extends:
+      file: compose.mongo.yml
+      service: mongo
+
+    hostname: mongo
+    volumes:
+      - ./data/mongo-container:/data/db
+      - ../scripts/mongodb_init_rs_profiling.js:/docker-entrypoint-initdb.d/init-rs.js:ro
+
+    # This forces "mongo" to resolve to localhost ONLY inside this container.
+    # This allows rs.initiate to succeed using the hostname "mongo".
+    extra_hosts:
+      - "mongo:127.0.0.1"
+
+  app:
+    extends:
+      file: compose.store.yml
+      service: app
+
+    depends_on:
+      - mongo
+
+    command: agl store --host 0.0.0.0 --port 4747 --prometheus --backend mongo --mongo-uri mongodb://mongo:27017/?replicaSet=rs0 --n-workers 4
+
+  mongodb-exporter:
+    image: percona/mongodb_exporter:0.47.1
+    command:
+      - '--mongodb.uri=mongodb://mongo:27017/'
+      - '--collect-all'
+      - '--mongodb.collstats-colls=agentlightning.rollouts,agentlightning.attempts,agentlightning.spans,agentlightning.resources,agentlightning.workers,agentlightning.rollout_queue,agentlightning.span_sequence_ids'
+    depends_on:
+      - mongo
+    ports:
+      - "9216:9216"
+
+  node-exporter:
+    image: prom/node-exporter:latest
+    # In CI you might not have full /proc, but this is OK for container-level stats
+    pid: "host"
+    command:
+      - '--path.rootfs=/host'
+    volumes:
+      - '/:/host:ro,rslave'
+
+  prometheus:
+    image: prom/prometheus:latest
+    command:
+      - '--config.file=/etc/prometheus/prometheus.yml'
+      - '--storage.tsdb.path=/prometheus'
+      - '--storage.tsdb.retention.time=1h'
+    volumes:
+      - ./prometheus.mongo-store.yml:/etc/prometheus/prometheus.yml:ro
+      - ./data/prometheus:/prometheus
+    depends_on:
+      - app
+      - mongodb-exporter
+      - node-exporter
+    ports:
+      - "9090:9090"

docker/compose.store.yml

@@ -0,0 +1,26 @@
+services:
+
+  app:
+    build:
+      context: ../
+      dockerfile: docker/Dockerfile.dev
+
+    ports:
+      - "4747:4747"
+
+    command: agl store --host 0.0.0.0 --port 4747
+
+    develop:
+      watch:
+        # Sync the working directory with the `/app` directory in the container
+        - action: sync
+          path: ..
+          target: /app
+          # Exclude the project virtual environment — it could be for a
+          # different platform in the container
+          ignore:
+            - .venv/
+
+        # Rebuild the image if dependencies change by checking uv.lock
+        - action: rebuild
+          path: ../uv.lock

docker/prometheus.memory-store.yml

@@ -0,0 +1,17 @@
+global:
+  scrape_interval: 2s
+  evaluation_interval: 2s
+
+scrape_configs:
+  - job_name: app
+    static_configs:
+      - targets: ["app:4747"]
+    metrics_path: /v1/prometheus
+
+  - job_name: mongodb
+    static_configs:
+      - targets: ["mongodb-exporter:9216"]
+
+  - job_name: node
+    static_configs:
+      - targets: ["node-exporter:9100"]

docker/prometheus.mongo-store.yml

@@ -0,0 +1,13 @@
+global:
+  scrape_interval: 2s
+  evaluation_interval: 2s
+
+scrape_configs:
+  - job_name: app
+    static_configs:
+      - targets: ["app:4747"]
+    metrics_path: /v1/prometheus
+
+  - job_name: node
+    static_configs:
+      - targets: ["node-exporter:9100"]

docker/setup.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+set -euo pipefail
+
+# Create data directories
+mkdir -p data/prometheus data/mongo-container data/mongo-host
+
+# Change permissions
+chmod 777 data/prometheus data/mongo-container data/mongo-host

examples/minimal/write_traces.py

@@ -102,6 +102,7 @@ async def send_traces_via_agentops(use_client: bool = False):
                     {"role": "user", "content": "Hello, what's your name?"},
                 ],
             )
+            console.print(response)
             assert response.choices[0].message.content is not None
             assert "chatgpt" in response.choices[0].message.content.lower()
 

pyproject.toml

@@ -65,6 +65,7 @@ dev = [
   "mkdocs-git-authors-plugin",
   "mkdocs-macros-plugin",
   "mkdocs-autorefs",
+  "prometheus-client",
 ]
 experiment = [
   "random-word",

scripts/mongodb_init_rs_host.js

@@ -0,0 +1,9 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+// MongoDB replica set initialization script.
+// Use this if you are accessing MongoDB from the **host**.
+
+rs.initiate({
+  _id: "rs0",
+  members: [{ _id: 0, host: "localhost:27017" }],
+});

scripts/mongodb_init_rs_profiling.js

@@ -0,0 +1,12 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+// MongoDB replica set initialization script.
+// Use this if you are accessing MongoDB from another **container**.
+// `mongodb_init_rs_host.js` is the counterpart if accessing from the host.
+
+rs.initiate({
+  _id: "rs0",
+  members: [{ _id: 0, host: "mongo:27017" }],
+});
+
+db.setProfilingLevel(2);

tests/benchmark/benchmark_store.py

@@ -0,0 +1,262 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Benchmarking store performance by writing and querying spans from the store."""
+
+import asyncio
+import random
+from typing import Any, Dict, List, Literal, Optional, Sequence, Set, Tuple
+
+import agentlightning as agl
+from agentlightning.emitter.utils import get_tracer
+
+from .utils import flatten_dict, random_dict
+
+
+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),
+        )
+    )
+
+
+@agl.rollout
+async def agent(task: str, llm: agl.LLM):
+    tracer = get_tracer()
+    rounds = random.randint(1, 10)
+    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, 1.0))
+                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, 1.0))
+                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, 1.0))
+        span.set_attributes(generate_attributes())
+
+    agl.emit_reward(random.uniform(1.0, 2.0))
+
+
+def check_spans(spans: Sequence[agl.Span], task: str) -> None:
+    """Check if the spans contain the task."""
+    found_task = False
+    last_reward_in_12 = None
+    for span in spans:
+        if span.attributes.get("task") == task:
+            found_task = True
+        if span.name == agl.SpanNames.REWARD.value:
+            if span.attributes.get("reward") is None:
+                raise ValueError("Reward is not set for a reward span")
+            rew = float(span.attributes.get("reward"))  # type: ignore
+            if rew >= 1 and rew <= 2:
+                last_reward_in_12 = True
+            else:
+                last_reward_in_12 = False
+    if not found_task:
+        raise ValueError(f"Task {task} is not found in the spans")
+    if last_reward_in_12 is None:
+        raise ValueError("Last reward is not found")
+    elif not last_reward_in_12:
+        raise ValueError("Last reward is not in the range of 1 to 2")
+
+
+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()
+        submitted = 0
+
+        while submitted < 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()
+            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,
+                )
+                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)
+                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()
+        submitted = 0
+        completed = 0
+        active_rollouts: Dict[str, str] = {}
+
+        while completed < total_tasks:
+            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
+
+            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)
+
+        async def handle_single(task_index: int) -> None:
+            task_name = f"task-{task_index}"
+            async with semaphore:
+                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
+                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 asyncio.sleep(5.0)
+                spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
+                check_spans(spans, task_name)
+
+        all_tasks = [handle_single(i) for i in range(total_tasks)]
+        await asyncio.gather(*all_tasks)
+
+
+def main() -> None:
+    store = agl.LightningStoreClient("http://localhost:4747")
+    try:
+        trainer = agl.Trainer(
+            store=store,
+            algorithm=AlgorithmBatch(mode="batch", total_tasks=1024, batch_size=128),
+            n_runners=32,
+            strategy={
+                "type": "cs",
+                "managed_store": False,
+            },
+        )
+        trainer.fit(agent)
+    finally:
+        asyncio.run(store.close())
+
+
+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))

tests/store/test_implementation.py

@@ -2014,8 +2014,12 @@ async def test_wait_nonexistent_rollout_with_finite_timeout(store_fixture: Light
     completed = await store_fixture.wait_for_rollouts(rollout_ids=["nonexistent"], timeout=0.1)
     elapsed = time.time() - start
 
-    # Should timeout quickly (not wait indefinitely)
-    assert elapsed < 1.0
+    if isinstance(store_fixture, InMemoryLightningStore):
+        # Should timeout quickly (not wait indefinitely)
+        assert elapsed < 0.2
+    else:
+        # Should be slower, but not too slow
+        assert elapsed < 2.0
     assert len(completed) == 0
 
 

uv.lock

@@ -220,6 +220,7 @@ dev = [
     { name = "mkdocs-material", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
     { name = "mkdocstrings", extra = ["python"], marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
     { name = "pre-commit", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
+    { name = "prometheus-client", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
     { name = "pyright", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
     { name = "pytest", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
     { name = "pytest-asyncio", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
@@ -419,6 +420,7 @@ dev = [
     { name = "mkdocs-material" },
     { name = "mkdocstrings", extras = ["python"] },
     { name = "pre-commit" },
+    { name = "prometheus-client" },
     { name = "pyright" },
     { name = "pytest" },
     { name = "pytest-asyncio" },