fix tests

.github/workflows/benchmark.yml

@@ -18,7 +18,7 @@ jobs:
           - id: mongo
             compose_file: compose.prometheus-mongo-store.yml
         workload:
-          - id: minimal-production
+          - id: scenario-minimal-scale
             display: Minimal production scale
             kind: scenario
             store_workers: 4
@@ -33,7 +33,7 @@ jobs:
               --n-runners 32
               --max-rounds 6
               --sleep-seconds 0.5
-          - id: medium-production
+          - id: scenario-medium-scale
             display: Medium production scale
             kind: scenario
             store_workers: 16
@@ -48,7 +48,22 @@ jobs:
               --n-runners 100
               --max-rounds 10
               --sleep-seconds 0.1
-          - id: large-batch
+          - id: scenario-midhigh-scale
+            display: Mid-high production scale
+            kind: scenario
+            store_workers: 24
+            runner:
+              - self-hosted
+              - 1ES.Pool=agl-runner-cpu
+            timeout: 60
+            args: >-
+              --mode batch
+              --total-tasks 20000
+              --batch-size 2048
+              --n-runners 256
+              --max-rounds 8
+              --sleep-seconds 0.1
+          - id: scenario-large-batch
             display: Large batch waves
             kind: scenario
             store_workers: 32
@@ -63,7 +78,7 @@ jobs:
               --n-runners 256
               --max-rounds 6
               --sleep-seconds 0.1
-          - id: long-queues
+          - id: scenario-long-queues
             display: Long rollout queues
             kind: scenario
             store_workers: 32
@@ -79,7 +94,7 @@ jobs:
               --remaining-tasks 4096
               --max-rounds 4
               --sleep-seconds 0.1
-          - id: high-concurrency
+          - id: scenario-high-concurrency
             display: High-throughput concurrent requests
             kind: scenario
             store_workers: 32
@@ -94,7 +109,7 @@ jobs:
               --n-runners 256
               --max-rounds 2
               --sleep-seconds 0.1
-          - id: heavy-traces
+          - id: scenario-heavy-traces
             display: Heavy rollouts with deep traces
             kind: scenario
             store_workers: 64
@@ -149,7 +164,7 @@ jobs:
           - id: micro-metrics
             display: Multi-metric fan-out
             kind: micro
-            store_workers: 1
+            store_workers: 8
             runner: ubuntu-latest
             timeout: 15
             cli: metrics
@@ -166,7 +181,7 @@ jobs:
       ANALYSIS_FILE: ${{ format('analysis-{0}.log', matrix.workload.id) }}
       SUMMARY_FILE: ${{ format('summary-{0}.log', matrix.workload.id) }}
       PROM_ARCHIVE_BASENAME: ${{ format('prometheus-{0}-{1}', matrix.workload.id, matrix.backend.id) }}
-      ARTIFACT_NAME: ${{ format('{0}-{1}-{2}', matrix.workload.kind == 'micro' && 'micro-benchmark' || 'benchmark', matrix.workload.id, matrix.backend.id) }}
+      ARTIFACT_NAME: ${{ format('{0}-{1}', matrix.workload.id, matrix.backend.id) }}
     steps:
       - uses: actions/checkout@v4
 
@@ -200,6 +215,9 @@ jobs:
           set -euo pipefail
           for attempt in {1..60}; do
             if curl -fsS "$STORE_API_URL/health" >/dev/null 2>&1; then
+              sleep 1
+              curl -fsS "$STORE_API_URL/rollouts"  # Warm up the scraper
+              sleep 15  # Allow some time for the baseline metrics to be established
               exit 0
             fi
             sleep 1
@@ -303,3 +321,116 @@ jobs:
           name: ${{ env.ARTIFACT_NAME }}
           path: ${{ env.ARTIFACT_DIR }}
           if-no-files-found: error
+
+  collection-benchmarks:
+    name: collection (${{ matrix.backend.id }}, ${{ matrix.workload.id }})
+    runs-on: ${{ matrix.backend.runner }}
+    timeout-minutes: 15
+    strategy:
+      fail-fast: false
+      matrix:
+        backend:
+          - id: memory
+            needs_mongo: false
+            runner: ubuntu-latest
+          - id: mongo
+            needs_mongo: true
+            runner: ubuntu-latest
+        workload:
+          - id: high-insert
+            total_tasks: 100000
+            concurrency: 2048
+            type: insert
+          - id: medium-insert
+            total_tasks: 100000
+            concurrency: 128
+            type: insert
+          - id: low-insert
+            total_tasks: 100000
+            concurrency: 4
+            type: insert
+          - id: high-dequeue
+            total_tasks: 100000
+            concurrency: 2048
+            type: dequeue
+          - id: medium-dequeue
+            total_tasks: 100000
+            concurrency: 128
+            type: dequeue
+          - id: low-dequeue
+            total_tasks: 100000
+            concurrency: 4
+            type: dequeue
+    env:
+      ARTIFACT_DIR: ${{ format('artifacts/{0}-{1}', matrix.backend.id, matrix.workload.id) }}
+      SUMMARY_FILE: ${{ format('artifacts/{0}-{1}/summary-{0}-{1}.jsonl', matrix.backend.id, matrix.workload.id) }}
+      ARTIFACT_NAME: ${{ format('collections-{0}-{1}', matrix.backend.id, matrix.workload.id) }}
+      MONGO_URI: mongodb://localhost:27017/?replicaSet=rs0
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: '3.12'
+
+      - name: Sync dependencies
+        run: uv sync --frozen --extra mongo --group core-stable --group dev
+
+      - name: Launch MongoDB
+        if: ${{ matrix.backend.needs_mongo }}
+        run: |
+          set -euo pipefail
+          cd docker
+          docker compose -f compose.mongo.yml down -v || true
+          docker compose -f compose.mongo.yml up -d --quiet-pull
+          for attempt in {1..60}; do
+            if docker compose -f compose.mongo.yml exec -T mongo mongosh --quiet --eval 'db.runCommand({ping:1})' >/dev/null 2>&1; then
+              exit 0
+            fi
+            sleep 2
+          done
+          echo "MongoDB did not become ready in time" >&2
+          docker compose -f compose.mongo.yml logs mongo
+          exit 1
+
+      - name: Run collection benchmark
+        run: |
+          set -euo pipefail
+          mkdir -p "$ARTIFACT_DIR"
+          echo "Running collection benchmark (backend=${{ matrix.backend.id }}, workload=${{ matrix.workload.id }})"
+          uv run --locked --no-sync python -m tests.benchmark.collection_benchmark \
+            "${{ matrix.workload.type }}" \
+            --backend "${{ matrix.backend.id }}" \
+            --total-tasks "${{ matrix.workload.total_tasks }}" \
+            --concurrency "${{ matrix.workload.concurrency }}" \
+            --task-prefix "${{ matrix.backend.id }}-${{ matrix.workload.id }}" \
+            --summary-file "$SUMMARY_FILE" \
+            --mongo-uri "$MONGO_URI" \
+            --mongo-database agentlightning_collection_bench
+
+      - name: Show collection benchmark summary
+        if: ${{ always() }}
+        run: |
+          set -euo pipefail
+          if [ -f "$SUMMARY_FILE" ]; then
+            echo "Collection benchmark summary (${{ matrix.backend.id }}):"
+            cat "$SUMMARY_FILE"
+          else
+            echo "Summary file not found: $SUMMARY_FILE"
+          fi
+
+      - name: Stop MongoDB
+        if: ${{ always() && matrix.backend.needs_mongo }}
+        run: |
+          set -euo pipefail
+          cd docker
+          docker compose -f compose.mongo.yml down -v || true
+
+      - name: Upload collection artifacts
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ env.ARTIFACT_NAME }}
+          path: ${{ env.ARTIFACT_DIR }}
+          if-no-files-found: error

AGENTS.md

@@ -21,7 +21,7 @@ Always commit the refreshed `uv.lock` when dependencies shift, and mention optio
 ## Coding Style & Naming Conventions
 - Target `requires-python >= 3.10`, four-space indentation, 120-character lines (though docstrings may run longer), and formatter-owned diffs (Black + isort, `black` profile). Use `snake_case` for modules, functions, and variables; `PascalCase` for classes and React components; lowercase hyphenation for CLI flags, branch names, and TypeScript filenames.
 - Maintain exhaustive type hints (pyright enforces them) and prefer shared dataclasses or Pydantic models from `agentlightning.types`.
-- Author Google-style docstrings for new modules or public methods—succinct descriptions, no redundant type info, no redundant `Key features/components` bullet points, and `[][]` syntax for cross-references.
+- Author Google-style docstrings for new modules or public methods—succinct descriptions, no redundant type info, no redundant `Key features/components` bullet points. Use mkdocs styles: `[][]` syntax for cross-references and single backticks for inline code blocks.
 - Writing logs is encouraged, especially for long functions with multiple steps and try-except blocks that catch all exceptions. Use `logging.getLogger(__name__)` to get loggers. Distinguish between DEBUG, INFO, WARNING, and ERROR logs.
 
 ## Testing Guidelines

agentlightning/cli/__init__.py

@@ -12,6 +12,7 @@ from typing import Dict, Iterable, Tuple
 _SUBCOMMANDS: Dict[str, Tuple[str, str]] = {
     "vllm": ("agentlightning.cli.vllm", "Run the vLLM CLI with Agent Lightning instrumentation."),
     "store": ("agentlightning.cli.store", "Run a LightningStore server."),
+    "prometheus": ("agentlightning.cli.prometheus", "Serve Prometheus metrics from the multiprocess registry."),
     "agentops": ("agentlightning.cli.agentops_server", "Start the AgentOps server manager."),
 }
 

agentlightning/cli/prometheus.py

@@ -0,0 +1,115 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Serve Prometheus metrics from the Agent Lightning multiprocess registry."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+import os
+from pathlib import Path
+from typing import Iterable
+
+from fastapi import FastAPI
+from prometheus_client import make_asgi_app  # pyright: ignore[reportUnknownVariableType]
+
+from agentlightning.logging import setup as setup_logging
+from agentlightning.utils.metrics import get_prometheus_registry
+from agentlightning.utils.server_launcher import PythonServerLauncher, PythonServerLauncherArgs
+
+logger = logging.getLogger(__name__)
+
+
+def ensure_prometheus_dir() -> str:
+    """Ensure PROMETHEUS_MULTIPROC_DIR is set and the directory exists."""
+
+    directory = os.getenv("PROMETHEUS_MULTIPROC_DIR")
+    if directory is None:
+        raise ValueError("PROMETHEUS_MULTIPROC_DIR is not set.")
+
+    Path(directory).mkdir(parents=True, exist_ok=True)
+    logger.info("Serving Prometheus multiprocess metrics from %s", directory)
+    return directory
+
+
+def create_prometheus_app(metrics_path: str = "/v1/prometheus") -> FastAPI:
+    """Create a FastAPI app that exposes Prometheus metrics and a health endpoint.
+
+    Args:
+        metrics_path: URL path to expose the Prometheus metrics endpoint on.
+
+    Returns:
+        A FastAPI application ready to serve metrics.
+    """
+
+    if not metrics_path.startswith("/"):
+        raise ValueError("metrics_path must start with '/'.")
+
+    normalized_path = metrics_path.rstrip("/")
+    if normalized_path in ("", "/"):
+        raise ValueError("metrics_path must not be '/'. Choose a sub-path such as /v1/prometheus.")
+
+    app = FastAPI(title="Agent Lightning Prometheus exporter", docs_url=None, redoc_url=None)
+    metrics_app = make_asgi_app(registry=get_prometheus_registry())  # pyright: ignore[reportUnknownVariableType]
+    app.mount(normalized_path, metrics_app)  # pyright: ignore[reportUnknownArgumentType]
+
+    @app.get("/health")
+    async def healthcheck() -> dict[str, str]:  # pyright: ignore[reportUnusedFunction]
+        return {"status": "ok"}
+
+    return app
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Serve Prometheus metrics outside the LightningStore server.")
+    parser.add_argument("--host", default="0.0.0.0", help="Host to bind the metrics server to.")
+    parser.add_argument("--port", type=int, default=4748, help="Port to expose the Prometheus metrics on.")
+    parser.add_argument(
+        "--metrics-path",
+        default="/v1/prometheus",
+        help="HTTP path used to expose metrics. Must start with '/' and not be the root path.",
+    )
+    parser.add_argument(
+        "--log-level",
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="Configure the logging level for the metrics server.",
+    )
+    parser.add_argument(
+        "--access-log",
+        action="store_true",
+        help="Enable uvicorn access logs. Disabled by default to reduce noise.",
+    )
+    args = parser.parse_args(list(argv) if argv is not None else None)
+
+    setup_logging(args.log_level)
+    ensure_prometheus_dir()
+
+    try:
+        app = create_prometheus_app(args.metrics_path)
+    except ValueError as exc:
+        logger.error("Failed to configure prometheus app: %s", exc)
+        return 1
+
+    launcher_args = PythonServerLauncherArgs(
+        host=args.host,
+        port=args.port,
+        log_level=getattr(logging, args.log_level),
+        access_log=args.access_log,
+        healthcheck_url="/health",
+    )
+    launcher = PythonServerLauncher(app, launcher_args)
+
+    try:
+        asyncio.run(launcher.run_forever())
+    except KeyboardInterrupt:
+        logger.info("Received shutdown signal. Stopping Prometheus server.")
+    except RuntimeError as exc:
+        logger.error("Prometheus server failed to start: %s", exc, exc_info=True)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

agentlightning/cli/store.py

@@ -100,7 +100,7 @@ def main(argv: Iterable[str] | None = None) -> int:
     elif args.backend == "mongo":
         from agentlightning.store.mongo import MongoLightningStore
 
-        store = MongoLightningStore(client=args.mongo_uri, tracker=tracker)
+        store = MongoLightningStore(mongo_uri=args.mongo_uri, tracker=tracker)
     else:
         raise ValueError(f"Invalid backend: {args.backend}")
 

agentlightning/store/client_server.py

@@ -270,6 +270,7 @@ class LightningStoreServer(LightningStore):
                 port=port,
                 launch_mode=launch_mode,
                 healthcheck_url=API_V1_AGL_PREFIX + "/health",
+                n_workers=n_workers,
             )
 
         store_capabilities = self.store.capabilities

agentlightning/store/collection/base.py

@@ -3,9 +3,9 @@
 from __future__ import annotations
 
 import functools
-import inspect
 import time
 from contextlib import asynccontextmanager
+from numbers import Real
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -22,6 +22,7 @@ from typing import (
     Sequence,
     Tuple,
     Type,
+    TypeGuard,
     TypeVar,
     cast,
 )
@@ -32,7 +33,6 @@ from agentlightning.utils.metrics import MetricsBackend
 if TYPE_CHECKING:
     from typing import Self
 
-from agentlightning.store.base import LightningStore
 from agentlightning.types import (
     Attempt,
     FilterField,
@@ -53,41 +53,17 @@ T_callable = TypeVar("T_callable", bound=Callable[..., Any])
 AtomicMode = Literal["r", "w", "rw"]
 """What is expected within the atomic context. Can be "read", "write", or "read-write"."""
 
-AtomicLabels = Literal["rollouts", "attempts", "spans", "resources", "workers", "rollout_queue", "span_sequence_ids"]
+AtomicLabels = Literal[
+    "rollouts", "attempts", "spans", "resources", "workers", "rollout_queue", "span_sequence_ids", "generic"
+]
 """Labels for atomic operations.
 
 These labels are used to identify the collections that are affected by the atomic operation.
+
+The `generic` label is used to identify atomic operations that are not associated with any specific collection.
 """
 
 
-COLLECTION_TRACKING_STORE_METHODS = frozenset(
-    [name for name in LightningStore.__dict__ if not name.startswith("_")] + ["_healthcheck"]
-)
-
-_UNKNOWN_STORE_METHOD = "unknown"
-
-
-def _nearest_lightning_store_method_from_stack() -> str:
-    """Stack introspection so that we capture the nearest public API method from the
-    call stack whenever metrics are recorded."""
-    frame = inspect.currentframe()
-    try:
-        if frame is None:
-            return _UNKNOWN_STORE_METHOD
-        frame = frame.f_back
-        while frame is not None:
-            self_obj = frame.f_locals.get("self")
-            method_name = frame.f_locals.get("method_name")
-            if method_name in COLLECTION_TRACKING_STORE_METHODS and isinstance(self_obj, LightningStore):
-                return method_name
-            frame = frame.f_back
-        return _UNKNOWN_STORE_METHOD
-    except Exception:
-        return _UNKNOWN_STORE_METHOD
-    finally:
-        del frame
-
-
 def resolve_error_type(exc: BaseException | None) -> str:
     if exc is None:
         return "N/A"
@@ -120,6 +96,25 @@ def tracked(operation: str):
     return decorator
 
 
+def ensure_numeric(value: Any, *, description: str) -> TypeGuard[Real]:
+    """Validate that *value* behaves like a real number.
+
+    Returns true or crashes.
+    """
+
+    if isinstance(value, bool):
+        raise TypeError(f"{description} must be numeric; got bool")
+    if not isinstance(value, Real):
+        raise TypeError(f"{description} must be numeric; got {type(value).__name__}")
+    return True
+
+
+class DuplicatedPrimaryKeyError(ValueError):
+    """Error raised when a duplicate key is encountered."""
+
+    pass
+
+
 class TrackedCollection:
     """An object that can be tracked by the metrics backend."""
 
@@ -153,10 +148,12 @@ class TrackedCollection:
             yield
 
         else:
+            from agentlightning.store.collection_based import nearest_lightning_store_method_from_stack
+
             # Enable tracking
             start_time = time.perf_counter()
             status: str = "OK"
-            store_method = _nearest_lightning_store_method_from_stack()
+            public_store_method, private_store_method = nearest_lightning_store_method_from_stack()
             try:
                 yield
             except BaseException as exc:
@@ -167,7 +164,8 @@ class TrackedCollection:
                 await self._tracker.inc_counter(  # pyright: ignore[reportPrivateUsage]
                     "agl.collections.total",
                     labels={
-                        "store_method": store_method,
+                        "store_pubmeth": public_store_method,
+                        "store_privmeth": private_store_method,
                         "operation": operation,
                         "collection": collection,
                         "status": status,
@@ -178,7 +176,8 @@ class TrackedCollection:
                     "agl.collections.latency",
                     value=elapsed,
                     labels={
-                        "store_method": store_method,
+                        "store_pubmeth": public_store_method,
+                        "store_privmeth": private_store_method,
                         "operation": operation,
                         "collection": collection,
                         "status": status,
@@ -382,6 +381,22 @@ class KeyValue(TrackedCollection, Generic[K, V]):
         """Set the value for the given key."""
         raise NotImplementedError()
 
+    async def inc(self, key: K, amount: V) -> V:
+        """Increase the numeric value for the given key by `amount` and return the new value.
+
+        Raises:
+            TypeError: If the existing value or `amount` is not numeric.
+        """
+        raise NotImplementedError()
+
+    async def chmax(self, key: K, value: V) -> V:
+        """Set the value for the given key to the maximum of the current and new value.
+
+        Raises:
+            TypeError: If the existing value or `value` is not numeric.
+        """
+        raise NotImplementedError()
+
     async def pop(self, key: K, default: V | None = None) -> V | None:
         """Pop the value for the given key, or the default value if the key is not found."""
         raise NotImplementedError()
@@ -405,7 +420,7 @@ class LightningCollections(TrackedCollection):
     def register_collection_metrics(self, extra_labels: Optional[Sequence[str]] = None) -> None:
         if self._tracker is None:
             return
-        labels = ["store_method", "operation", "collection", "status"]
+        labels = ["store_pubmeth", "operation", "collection", "store_privmeth", "status"]
         if extra_labels is not None:
             labels.extend(extra_labels)
         self._tracker.register_histogram(

agentlightning/store/collection/memory.py

@@ -23,6 +23,7 @@ from typing import (
     Type,
     TypeVar,
     Union,
+    cast,
 )
 
 import aiologic
@@ -42,12 +43,15 @@ from agentlightning.types import (
 from agentlightning.utils.metrics import MetricsBackend
 
 from .base import (
+    AtomicLabels,
     AtomicMode,
     Collection,
+    DuplicatedPrimaryKeyError,
     FilterMap,
     KeyValue,
     LightningCollections,
     Queue,
+    ensure_numeric,
     normalize_filter_options,
     resolve_sort_options,
     tracked,
@@ -313,7 +317,9 @@ class ListBasedCollection(Collection[T]):
 
             if mode == "insert":
                 if exists:
-                    raise ValueError(f"Item already exists with primary key(s): {self._render_key_values(key_values)}")
+                    raise DuplicatedPrimaryKeyError(
+                        f"Item already exists with primary key(s): {self._render_key_values(key_values)}"
+                    )
                 parent[final_key] = item
                 self._size += 1
             else:  # upsert
@@ -603,7 +609,7 @@ class ListBasedCollection(Collection[T]):
         """Insert the given items.
 
         Raises:
-            ValueError: If any item with the same primary keys already exists.
+            DuplicatedPrimaryKeyError: If any item with the same primary keys already exists.
         """
         seen_keys: set[Tuple[Any, ...]] = set()
         prepared: List[T] = []
@@ -611,8 +617,8 @@ class ListBasedCollection(Collection[T]):
             self._ensure_item_type(item)
             key_values = self._extract_primary_key_values(item)
             if key_values in seen_keys:
-                raise ValueError(
-                    f"Insert payload contains duplicate primary key(s): {self._render_key_values(key_values)}"
+                raise DuplicatedPrimaryKeyError(
+                    f"Insert payload contains duplicated primary key(s): {self._render_key_values(key_values)}"
                 )
             seen_keys.add(key_values)
             prepared.append(item)
@@ -756,6 +762,33 @@ class DictBasedKeyValue(KeyValue[K, V]):
     async def set(self, key: K, value: V) -> None:
         self._values[key] = value
 
+    @tracked("inc")
+    async def inc(self, key: K, amount: V) -> V:
+        assert ensure_numeric(amount, description="amount")
+        if key in self._values:
+            current_value = self._values[key]
+            assert ensure_numeric(current_value, description=f"value for key {key!r}")
+            new_value = cast(V, current_value + amount)
+            self._values[key] = new_value
+        else:
+            new_value = amount
+            self._values[key] = new_value
+        return new_value
+
+    @tracked("chmax")
+    async def chmax(self, key: K, value: V) -> V:
+        assert ensure_numeric(value, description="value")
+        if key in self._values:
+            current_value = self._values[key]
+            assert ensure_numeric(current_value, description=f"value for key {key!r}")
+            if value > current_value:
+                self._values[key] = value
+                return value
+            return current_value
+        else:
+            self._values[key] = value
+            return value
+
     @tracked("pop")
     async def pop(self, key: K, default: V | None = None) -> V | None:
         return self._values.pop(key, default)
@@ -773,7 +806,7 @@ class InMemoryLightningCollections(LightningCollections):
 
     def __init__(self, lock_type: Literal["thread", "asyncio"], tracker: MetricsBackend | None = None):
         super().__init__(tracker=tracker)
-        self._lock = {
+        self._lock: Mapping[AtomicLabels, _LoopAwareAsyncLock | _ThreadSafeAsyncLock] = {
             "rollouts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
             "attempts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
             "spans": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
@@ -781,6 +814,7 @@ class InMemoryLightningCollections(LightningCollections):
             "workers": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
             "rollout_queue": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
             "span_sequence_ids": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "generic": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
         }
         self._rollouts = ListBasedCollection(
             items=[], item_type=Rollout, primary_keys=["rollout_id"], id="rollouts", tracker=tracker
@@ -836,7 +870,12 @@ class InMemoryLightningCollections(LightningCollections):
 
     @asynccontextmanager
     async def atomic(
-        self, *, mode: AtomicMode = "rw", snapshot: bool = False, labels: Optional[Sequence[str]] = None, **kwargs: Any
+        self,
+        *,
+        mode: AtomicMode = "rw",
+        snapshot: bool = False,
+        labels: Optional[Sequence[AtomicLabels]] = None,
+        **kwargs: Any,
     ):
         """In-memory collections apply a lock outside. It doesn't need to manipulate the collections inside.
 

agentlightning/store/collection/mongo.py

@@ -6,7 +6,6 @@ import asyncio
 import logging
 import random
 import re
-import threading
 import time
 from contextlib import asynccontextmanager
 from datetime import datetime
@@ -27,6 +26,8 @@ from typing import (
     cast,
 )
 
+import aiologic
+
 from agentlightning.utils.metrics import MetricsBackend
 
 if TYPE_CHECKING:
@@ -37,7 +38,14 @@ from pymongo import AsyncMongoClient, ReadPreference, ReturnDocument, WriteConce
 from pymongo.asynchronous.client_session import AsyncClientSession
 from pymongo.asynchronous.collection import AsyncCollection
 from pymongo.asynchronous.database import AsyncDatabase
-from pymongo.errors import CollectionInvalid, ConnectionFailure, DuplicateKeyError, OperationFailure, PyMongoError
+from pymongo.errors import (
+    BulkWriteError,
+    CollectionInvalid,
+    ConnectionFailure,
+    DuplicateKeyError,
+    OperationFailure,
+    PyMongoError,
+)
 from pymongo.read_concern import ReadConcern
 
 from agentlightning.types import (
@@ -52,11 +60,14 @@ from agentlightning.types import (
 )
 
 from .base import (
+    AtomicLabels,
     AtomicMode,
     Collection,
+    DuplicatedPrimaryKeyError,
     KeyValue,
     LightningCollections,
     Queue,
+    ensure_numeric,
     normalize_filter_options,
     resolve_sort_options,
     tracked,
@@ -237,23 +248,18 @@ async def _ensure_collection(
 
 
 class MongoClientPool(Generic[T_mapping]):
-    """A pool of MongoDB clients, each binded to a specific event loop.
+    """A pool of MongoDB clients, each bound to a specific event loop.
 
-    This class is to resolve the issue of MongoDB client cannot be shared across event loops:
-
-    ```
-    Cannot use AsyncMongoClient in different event loop. AsyncMongoClient uses low-level asyncio APIs that bind it to the event loop it was created on.
-    ```
-
-    Use the client pool with a context manager to ensure all clients are closed when the context is exited.
+    The pool lazily creates `AsyncMongoClient` instances per event loop using the provided
+    connection parameters, ensuring we never try to reuse a client across loops.
     """
 
-    def __init__(self, client: AsyncMongoClient[T_mapping]):
-        self._lock = threading.Lock()
-
-        self._client_base = client
+    def __init__(self, *, mongo_uri: str, mongo_client_kwargs: Mapping[str, Any] | None = None):
+        self._get_collection_lock = aiologic.Lock()
+        self._get_client_lock = aiologic.Lock()
+        self._mongo_uri = mongo_uri
+        self._mongo_client_kwargs = dict(mongo_client_kwargs or {})
         self._client_pool: Dict[int, AsyncMongoClient[T_mapping]] = {}
-
         self._collection_pool: Dict[Tuple[int, str, str], AsyncCollection[T_mapping]] = {}
 
     async def __aenter__(self) -> Self:
@@ -263,18 +269,16 @@ class MongoClientPool(Generic[T_mapping]):
         await self.close()
 
     async def close(self) -> None:
-        """Close all clients in the pool (except the base client)."""
+        """Close all clients currently tracked by the pool."""
 
-        with self._lock:
+        async with self._get_client_lock, self._get_collection_lock:
             clients = list(self._client_pool.values())
             self._client_pool.clear()
+            self._collection_pool.clear()
 
         for client in clients:
             try:
-                if client is not self._client_base:
-                    await client.close()
-                else:
-                    logger.debug("Skipping closing base client: %s", client)
+                await client.close()
             except Exception:
                 logger.exception("Error closing MongoDB client: %s", client)
 
@@ -283,31 +287,22 @@ class MongoClientPool(Generic[T_mapping]):
         key = id(loop)
 
         # If there is already a client specifically for this loop, return it.
-        if key in self._client_pool:
-            # Verify that the client still works.
-            await self._client_pool[key].aconnect()
-            return self._client_pool[key]
+        existing = self._client_pool.get(key)
+        if existing is not None:
+            await existing.aconnect()  # This actually does nothing if the client is already connected.
+            return existing
 
-        try:
-            # Try whether the base client will work.
-            await self._client_base.aconnect()
+        async with self._get_client_lock:
+            # Another coroutine may have already created the client.
+            if key in self._client_pool:
+                await self._client_pool[key].aconnect()
+                return self._client_pool[key]
 
-            with self._lock:
-                # If it works, add it to the pool and return it.
-                self._client_pool.setdefault(key, self._client_base)
-            return self._client_base
-        except RuntimeError as exc:
-            if "Cannot use AsyncMongoClient in different event loop" in str(exc):
-                with self._lock:
-                    # Create a new client for this loop.
-                    client = self._client_base._duplicate()  # type: ignore
-                    # Try whether the new client will work.
-                    await client.aconnect()
-                    # Add it to the pool and return it.
-                    self._client_pool.setdefault(key, client)  # type: ignore
-                return client  # type: ignore
-
-            raise
+            # Create a new client for this loop.
+            client = AsyncMongoClient[T_mapping](self._mongo_uri, **self._mongo_client_kwargs)
+            await client.aconnect()
+            self._client_pool[key] = client
+            return client
 
     async def get_collection(self, database_name: str, collection_name: str) -> AsyncCollection[T_mapping]:
         loop = asyncio.get_running_loop()
@@ -315,12 +310,16 @@ class MongoClientPool(Generic[T_mapping]):
         if key in self._collection_pool:
             return self._collection_pool[key]
 
-        # Create a new collection for this loop.
-        client = await self.get_client()
-        collection = client[database_name][collection_name]
-        with self._lock:
+        async with self._get_collection_lock:
+            # Another coroutine may have already created the collection.
+            if key in self._collection_pool:
+                return self._collection_pool[key]
+
+            # Create a new collection for this loop.
+            client = await self.get_client()
+            collection = client[database_name][collection_name]
             self._collection_pool.setdefault(key, collection)
-        return collection
+            return collection
 
 
 class MongoBasedCollection(Collection[T_model]):
@@ -339,7 +338,7 @@ class MongoBasedCollection(Collection[T_model]):
 
     def __init__(
         self,
-        client_pool: MongoClientPool[Mapping[str, Any]] | AsyncMongoClient[Mapping[str, Any]],
+        client_pool: MongoClientPool[Mapping[str, Any]],
         database_name: str,
         collection_name: str,
         partition_id: str,
@@ -349,10 +348,7 @@ class MongoBasedCollection(Collection[T_model]):
         tracker: MetricsBackend | None = None,
     ):
         super().__init__(tracker=tracker)
-        if isinstance(client_pool, AsyncMongoClient):
-            self._client_pool = MongoClientPool(client_pool)
-        else:
-            self._client_pool = client_pool
+        self._client_pool = client_pool
         self._database_name = database_name
         self._collection_name = collection_name
         self._partition_id = partition_id
@@ -547,24 +543,19 @@ class MongoBasedCollection(Collection[T_model]):
 
     @tracked("insert")
     async def insert(self, items: Sequence[T_model]) -> None:
+        """Insert items into the collection.
+
+        The implementation does NOT do checks for duplicate primary keys,
+        neither within the same insert call nor across different insert calls.
+        It relies on the database to enforce uniqueness via indexes.
+        """
         if not items:
             return
 
         collection = await self.ensure_collection()
         docs: List[Mapping[str, Any]] = []
-        pk_conditions: List[Dict[str, Any]] = []
-        seen_primary_keys: set[Tuple[Any, ...]] = set()
         for item in items:
             self._ensure_item_type(item)
-            pk_filter = self._pk_filter(item)
-            pk_values = tuple(pk_filter[pk] for pk in self._primary_keys)
-            if pk_values in seen_primary_keys:
-                raise ValueError(
-                    f"Insert payload contains duplicate primary key(s): {self._render_pk_values(pk_values)}"
-                )
-            seen_primary_keys.add(pk_values)
-            pk_conditions.append({pk: pk_filter[pk] for pk in self._primary_keys})
-
             doc = item.model_dump()
             doc["partition_id"] = self._partition_id
             docs.append(doc)
@@ -572,23 +563,17 @@ class MongoBasedCollection(Collection[T_model]):
         if not docs:
             return
 
-        if len(pk_conditions) == 1:
-            existing_filter: Dict[str, Any] = {"partition_id": self._partition_id, **pk_conditions[0]}
-        else:
-            existing_filter = {"partition_id": self._partition_id, "$or": pk_conditions}
-
-        async with self.tracking_context("insert.find_existing", self._collection_name):
-            existing = await collection.find_one(existing_filter, session=self._session)
-            if existing is not None:
-                existing_values = tuple(existing.get(pk) for pk in self._primary_keys)
-                raise ValueError(f"Item with primary key(s) {self._render_pk_values(existing_values)} already exists")
-
         try:
             async with self.tracking_context("insert.insert_many", self._collection_name):
                 await collection.insert_many(docs, session=self._session)
         except DuplicateKeyError as exc:
             # In case the DB enforces uniqueness via index, normalize to ValueError
-            raise ValueError("Duplicate key error while inserting items") from exc
+            raise DuplicatedPrimaryKeyError("Duplicated primary key(s) while inserting items") from exc
+        except BulkWriteError as exc:
+            write_errors = exc.details.get("writeErrors", [])
+            if write_errors and write_errors[0].get("code") == 11000:
+                raise DuplicatedPrimaryKeyError("Duplicated primary key(s) while inserting items") from exc
+            raise
 
     @tracked("update")
     async def update(self, items: Sequence[T_model], update_fields: Sequence[str] | None = None) -> List[T_model]:
@@ -707,7 +692,7 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
 
     def __init__(
         self,
-        client_pool: MongoClientPool[Mapping[str, Any]] | AsyncMongoClient[Mapping[str, Any]],
+        client_pool: MongoClientPool[Mapping[str, Any]],
         database_name: str,
         collection_name: str,
         partition_id: str,
@@ -723,10 +708,7 @@ class MongoBasedQueue(Queue[T_generic], Generic[T_generic]):
             item_type: The Python type of queue items (primitive or BaseModel subclass).
         """
         super().__init__(tracker=tracker)
-        if isinstance(client_pool, AsyncMongoClient):
-            self._client_pool = MongoClientPool(client_pool)
-        else:
-            self._client_pool = client_pool
+        self._client_pool = client_pool
         self._database_name = database_name
         self._collection_name = collection_name
         self._partition_id = partition_id
@@ -886,7 +868,7 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
 
     def __init__(
         self,
-        client_pool: MongoClientPool[Mapping[str, Any]] | AsyncMongoClient[Mapping[str, Any]],
+        client_pool: MongoClientPool[Mapping[str, Any]],
         database_name: str,
         collection_name: str,
         partition_id: str,
@@ -905,10 +887,7 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
             tracker: The metrics tracker to use.
         """
         super().__init__(tracker=tracker)
-        if isinstance(client_pool, AsyncMongoClient):
-            self._client_pool = MongoClientPool(client_pool)
-        else:
-            self._client_pool = client_pool
+        self._client_pool = client_pool
         self._database_name = database_name
         self._collection_name = collection_name
         self._partition_id = partition_id
@@ -1007,7 +986,65 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
                 )
         except DuplicateKeyError as exc:
             # Very unlikely with replace_one+upsert, but normalize anyway.
-            raise ValueError("Duplicate key error while setting key-value item") from exc
+            raise DuplicatedPrimaryKeyError("Duplicate key error while setting key-value item") from exc
+
+    @tracked("inc")
+    async def inc(self, key: K, amount: V) -> V:
+        assert ensure_numeric(amount, description="amount")
+        collection = await self.ensure_collection()
+        encoded_key = self._key_adapter.dump_python(key, mode="python")
+        encoded_amount = self._value_adapter.dump_python(amount, mode="python")
+        try:
+            async with self.tracking_context("inc.find_one_and_update", self.collection_name):
+                doc = await collection.find_one_and_update(
+                    {
+                        "partition_id": self._partition_id,
+                        "key": encoded_key,
+                    },
+                    {
+                        "$inc": {"value": encoded_amount},
+                    },
+                    upsert=True,
+                    return_document=ReturnDocument.AFTER,
+                    session=self._session,
+                )
+        except OperationFailure as exc:
+            if exc.code == 14 or "Cannot apply $inc" in str(exc):
+                raise TypeError(f"value for key {key!r} is not numeric") from exc
+            raise
+        if doc is None:  # type: ignore
+            raise RuntimeError("Failed to increment value; MongoDB did not return a document")
+        raw_value = doc["value"]
+        return self._value_adapter.validate_python(raw_value)
+
+    @tracked("chmax")
+    async def chmax(self, key: K, value: V) -> V:
+        assert ensure_numeric(value, description="value")
+        collection = await self.ensure_collection()
+        encoded_key = self._key_adapter.dump_python(key, mode="python")
+        encoded_value = self._value_adapter.dump_python(value, mode="python")
+        try:
+            async with self.tracking_context("chmax.find_one_and_update", self.collection_name):
+                doc = await collection.find_one_and_update(
+                    {
+                        "partition_id": self._partition_id,
+                        "key": encoded_key,
+                    },
+                    {
+                        "$max": {"value": encoded_value},
+                    },
+                    upsert=True,
+                    return_document=ReturnDocument.AFTER,
+                    session=self._session,
+                )
+        except OperationFailure as exc:
+            if exc.code == 14 or "Cannot apply $max" in str(exc):
+                raise TypeError(f"value for key {key!r} is not numeric") from exc
+            raise
+        if doc is None:  # type: ignore
+            raise RuntimeError("Failed to update value; MongoDB did not return a document")
+        raw_value = doc["value"]
+        return self._value_adapter.validate_python(raw_value)
 
     @tracked("pop")
     async def pop(self, key: K, default: V | None = None) -> V | None:
@@ -1062,6 +1099,7 @@ class MongoLightningCollections(LightningCollections):
         self._database_name = database_name
         self._partition_id = partition_id
         self._collection_ensured = False
+        self._lock = aiologic.Lock()  # used for generic atomic operations like scan debounce seconds
         self._rollouts = (
             rollouts
             if rollouts is not None
@@ -1227,19 +1265,40 @@ class MongoLightningCollections(LightningCollections):
         await self._span_sequence_ids.ensure_collection()
         self._collection_ensured = True
 
+    @asynccontextmanager
+    async def _lock_manager(self, labels: Optional[Sequence[AtomicLabels]]):
+        if labels is None or "generic" not in labels:
+            yield
+
+        else:
+            # Only lock the generic label.
+            try:
+                async with self.tracking_context("lock", self.collection_name):
+                    await self._lock.async_acquire()
+                yield
+            finally:
+                self._lock.async_release()
+
     @asynccontextmanager
     async def atomic(
-        self, mode: AtomicMode = "rw", snapshot: bool = False, commit: bool = False, *args: Any, **kwargs: Any
+        self,
+        mode: AtomicMode = "rw",
+        snapshot: bool = False,
+        commit: bool = False,
+        labels: Optional[Sequence[AtomicLabels]] = None,
+        *args: Any,
+        **kwargs: Any,
     ):
         """Perform a atomic operation on the collections."""
         if commit:
             raise ValueError("Commit should be used with execute() instead.")
-        async with self.tracking_context("atomic", self.collection_name):
-            # First step: ensure all collections exist before going into the atomic block
-            if not self._collection_ensured:
-                await self._ensure_collections()
-            # Execute directly without commit
-            yield self
+        async with self._lock_manager(labels):
+            async with self.tracking_context("atomic", self.collection_name):
+                # First step: ensure all collections exist before going into the atomic block
+                if not self._collection_ensured:
+                    await self._ensure_collections()
+                # Execute directly without commit
+                yield self
 
     @tracked("execute")
     async def execute(
@@ -1249,6 +1308,7 @@ class MongoLightningCollections(LightningCollections):
         mode: AtomicMode = "rw",
         snapshot: bool = False,
         commit: bool = False,
+        labels: Optional[Sequence[AtomicLabels]] = None,
         **kwargs: Any,
     ) -> T_generic:
         """Execute the given callback within an atomic operation, and with retries on transient errors."""
@@ -1258,7 +1318,8 @@ class MongoLightningCollections(LightningCollections):
 
         # If commit is not turned on, just execute the callback directly.
         if not commit:
-            return await callback(self)
+            async with self._lock_manager(labels):
+                return await callback(self)
 
         # If snapshot is enabled, use snapshot read concern.
         read_concern = ReadConcern("snapshot") if snapshot else ReadConcern("local")
@@ -1268,7 +1329,8 @@ class MongoLightningCollections(LightningCollections):
         async with client.start_session() as session:
             collections = self.with_session(session)
             try:
-                return await self.with_transaction(session, collections, callback, read_concern, write_concern)
+                async with self._lock_manager(labels):
+                    return await self.with_transaction(session, collections, callback, read_concern, write_concern)
             except (ConnectionFailure, OperationFailure) as exc:
                 # Un-retryable errors.
                 raise RuntimeError("Transaction failed with connection or operation error") from exc

agentlightning/store/collection_based.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 import asyncio
 import functools
 import hashlib
+import inspect
 import logging
 import time
 import uuid
@@ -72,7 +73,7 @@ from .base import (
     is_queuing,
 )
 from .collection import FilterOptions, LightningCollections
-from .collection.base import COLLECTION_TRACKING_STORE_METHODS, AtomicLabels
+from .collection.base import AtomicLabels, DuplicatedPrimaryKeyError
 from .utils import LATENCY_BUCKETS, rollout_status_from_attempt, scan_unhealthy_rollouts
 
 T_callable = TypeVar("T_callable", bound=Callable[..., Any])
@@ -124,12 +125,16 @@ def tracked(name: str):
 
         @functools.wraps(func)
         async def wrapper(self: CollectionBasedLightningStore[T_collections], *args: Any, **kwargs: Any) -> Any:
+            # Backtracking where this method comes from
+            public_meth_in_stack, _ = nearest_lightning_store_method_from_stack()
+
             # For backtracking in collection methods.
             # Only track the public methods (+healthcheck)
-            if name in COLLECTION_TRACKING_STORE_METHODS:
-                method_name = name  # pyright: ignore[reportUnusedVariable]
-            else:
-                method_name = None  # pyright: ignore[reportUnusedVariable]
+            if name in COLLECTION_STORE_PUBLIC_METHODS:
+                public_method_name = name  # pyright: ignore[reportUnusedVariable]
+                public_meth_in_stack = name  # We are in a public method already.
+            if name in COLLECTION_STORE_ALL_METHODS:
+                private_method_name = name  # pyright: ignore[reportUnusedVariable]
 
             if self._tracker is None:  # pyright: ignore[reportPrivateUsage]
                 # Skip the tracking because tracking is not configured
@@ -145,10 +150,12 @@ def tracked(name: str):
             finally:
                 elapsed = time.perf_counter() - start_time
                 await self._tracker.inc_counter(  # pyright: ignore[reportPrivateUsage]
-                    "agl.store.total", labels={"method": name, "status": status}
+                    "agl.store.total", labels={"method": name, "store_pubmeth": public_meth_in_stack, "status": status}
                 )
                 await self._tracker.observe_histogram(  # pyright: ignore[reportPrivateUsage]
-                    "agl.store.latency", value=elapsed, labels={"method": name, "status": status}
+                    "agl.store.latency",
+                    value=elapsed,
+                    labels={"method": name, "store_pubmeth": public_meth_in_stack, "status": status},
                 )
 
         return cast(T_callable, wrapper)
@@ -219,27 +226,43 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             all read operations like `query_rollouts` will have better consistency.
             It may use an isolated snapshot that supports repeatable reads.
         tracker: Enable metrics tracking.
+        scan_debounce_seconds: The debounce time for the scan for unhealthy rollouts.
+            Set to 0 to disable debouncing. The debounce is a non-perfect traffic control.
+            It's isolated for each store instance if there are multiple worker replicas.
     """
 
     def __init__(
-        self, collections: T_collections, *, read_snapshot: bool = False, tracker: MetricsBackend | None = None
-    ):
+        self,
+        collections: T_collections,
+        *,
+        read_snapshot: bool = False,
+        tracker: MetricsBackend | None = None,
+        scan_debounce_seconds: float = 10.0,
+    ) -> None:
         # rollouts and spans' storage
         self.collections = collections
         self._read_snapshot = read_snapshot
         self._tracker = tracker
         self._launch_time = time.time()
 
+        # Control scan debounce to avoid overloading the store.
+        self._scan_debounce_seconds = scan_debounce_seconds
+        last_scan_time = self._launch_time
+        if self._scan_debounce_seconds > 0:
+            # Allow the first scan immediately after instantiation
+            last_scan_time -= self._scan_debounce_seconds
+        self._last_scan_entrance_time = last_scan_time
+
         if self._tracker is not None:
             self._tracker.register_histogram(
                 "agl.store.latency",
-                ["method", "status"],
+                ["method", "store_pubmeth", "status"],
                 buckets=LATENCY_BUCKETS,
                 group_level=1,
             )
             self._tracker.register_counter(
                 "agl.store.total",
-                ["method", "status"],
+                ["method", "store_pubmeth", "status"],
                 group_level=1,
             )
             self._tracker.register_counter(
@@ -972,37 +995,36 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         return await self._get_latest_resources()
 
     @tracked("_issue_many_span_sequence_ids")
-    @_with_collections_execute(labels=["span_sequence_ids"])
-    async def _issue_many_span_sequence_ids(self, collections: T_collections, rollout_ids: List[str]) -> List[int]:
+    async def _issue_many_span_sequence_ids(self, rollout_ids: List[str]) -> List[int]:
         """Issue a new span sequence ID for a given rollout."""
-        # Cache the next sequence IDs for the rollouts (for both RW)
-        next_sequence_ids_cache: Dict[str, int] = {}
+        if not rollout_ids:
+            return []
+
+        request_counts: Dict[str, int] = defaultdict(int)
+        for rollout_id in rollout_ids:
+            request_counts[rollout_id] += 1
+
+        latest_values: Dict[str, int] = {}
+        for rollout_id, count in request_counts.items():
+            async with self.collections.atomic(mode="rw", snapshot=False, labels=["span_sequence_ids"]) as collections:
+                latest_values[rollout_id] = await collections.span_sequence_ids.inc(rollout_id, count)
+
+        next_value_tracker: Dict[str, int] = {
+            rollout_id: latest_values[rollout_id] - request_counts[rollout_id] for rollout_id in request_counts
+        }
+
         result: List[int] = []
         for rollout_id in rollout_ids:
-            if rollout_id not in next_sequence_ids_cache:
-                retrieved_id = await collections.span_sequence_ids.get(rollout_id)
-                if retrieved_id is None:
-                    retrieved_id = 0
-                next_sequence_ids_cache[rollout_id] = retrieved_id
-
-            # Increment the sequence ID for the rollout
-            next_sequence_ids_cache[rollout_id] += 1
-            result.append(next_sequence_ids_cache[rollout_id])
-
-        # Propagate the cache to storage
-        for rollout_id, sequence_id in next_sequence_ids_cache.items():
-            await collections.span_sequence_ids.set(rollout_id, sequence_id)
+            next_value_tracker[rollout_id] += 1
+            result.append(next_value_tracker[rollout_id])
 
         return result
 
     @tracked("_sync_span_sequence_id")
-    @_with_collections_execute(labels=["span_sequence_ids"])
-    async def _sync_span_sequence_id(self, collections: T_collections, rollout_id: str, sequence_id: int) -> None:
+    async def _sync_span_sequence_id(self, rollout_id: str, sequence_id: int) -> None:
         """Sync the span sequence ID for a given rollout from the input span sequence ID."""
-        existing_sequence_id = await collections.span_sequence_ids.get(rollout_id)
-        if existing_sequence_id is None:
-            existing_sequence_id = 0
-        await collections.span_sequence_ids.set(rollout_id, max(existing_sequence_id, sequence_id))
+        async with self.collections.atomic(mode="rw", snapshot=False, labels=["span_sequence_ids"]) as collections:
+            await collections.span_sequence_ids.chmax(rollout_id, sequence_id)
 
     @tracked("get_next_span_sequence_id")
     async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
@@ -1084,13 +1106,11 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             try:
                 await collections.spans.insert([span])
                 return True
-            except ValueError as e:
-                if "already exists" in str(e) or "contains duplicate" in str(e):
-                    logger.error(
-                        f"Duplicated span added for rollout={span.rollout_id}, attempt={span.attempt_id}, span={span.span_id}. Skipping."
-                    )
-                    return False
-                raise
+            except DuplicatedPrimaryKeyError:
+                logger.error(
+                    f"Duplicated span added for rollout={span.rollout_id}, attempt={span.attempt_id}, span={span.span_id}. Skipping."
+                )
+                return False
 
         successful_spans: List[Span] = []
         try:
@@ -1098,22 +1118,20 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             async with self.collections.atomic(
                 mode="w", snapshot=self._read_snapshot, commit=False, labels=["spans"]
             ) as collections:
+                # FIXME: Part of the insertion might complete though the full operation fails.
+                # In that case, the "insert spans" return values might not be accurate.
                 await collections.spans.insert(spans)
             successful_spans.extend(spans)
-        except ValueError as e:
-            if "already exists" in str(e) or "contains duplicate" in str(e):
-                # There is a duplicate span, we warn it
-                # We fallback to adding the spans one by one
-                async def _add_many_spans_fallback(collections: T_collections):
-                    for span in spans:
-                        if await _add_span_fallback(collections, span):
-                            successful_spans.append(span)
-
-                await self.collections.execute(
-                    _add_many_spans_fallback, mode="w", snapshot=self._read_snapshot, commit=True, labels=["spans"]
-                )
-            else:
-                raise
+        except DuplicatedPrimaryKeyError:
+            # There is a duplicate span, we warn it
+            # We fallback to adding the spans one by one
+            for span in spans:
+                async with self.collections.atomic(
+                    mode="w", snapshot=self._read_snapshot, labels=["spans"]
+                ) as collections:
+                    # No need to commit here, it will be simple atomic write operations
+                    if await _add_span_fallback(collections, span):
+                        successful_spans.append(span)
 
         return successful_spans
 
@@ -1165,48 +1183,46 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         if not spans:
             return
 
-        async def _update_rollout_attempt(collections: T_collections) -> Optional[Tuple[Rollout, Sequence[str]]]:
-            attempt = await collections.attempts.get(
-                {"rollout_id": {"exact": rollout_id}, "attempt_id": {"exact": attempt_id}}
-            )
-            if attempt is None:
-                return None
-            rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
-            if rollout is None:
-                return None
-
-            # Update attempt heartbeat and ensure persistence
-            attempt.last_heartbeat_time = time.time()
-            if attempt.status in ["preparing", "unresponsive"]:
-                attempt.status = "running"
-            await collections.attempts.update([attempt], update_fields=["last_heartbeat_time", "status"])
-
-            # If the status has already timed out or failed, do not change it (but heartbeat is still recorded)
-
-            # Update rollout status if it's the latest attempt
-            rollout_updated: bool = False
-            updated_fields: List[str] = []
-            latest_attempt = await self._unlocked_get_latest_attempt(collections, rollout.rollout_id)
-            if latest_attempt is not None and attempt.attempt_id == latest_attempt.attempt_id:
-                if rollout.status in ["preparing", "queueing", "requeuing"]:
-                    # If rollout is currently preparing or queuing, set it to running
-                    rollout.status = "running"
-                    await collections.rollouts.update([rollout], update_fields=["status"])
-                    rollout_updated = True
-                    updated_fields = ["status"]
-                # Otherwise, the rollout has succeeded or failed, do nothing
-            return (rollout, updated_fields) if rollout_updated else None
-
-        rollout_update = await self.collections.execute(
-            _update_rollout_attempt,
-            mode="rw",
-            snapshot=self._read_snapshot,
-            commit=True,
-            labels=["rollouts", "attempts"],
-        )
+        rollout_update = await self._on_attempt_heartbeat(rollout_id=rollout_id, attempt_id=attempt_id)
         if rollout_update is not None:
             await self._post_update_rollout([rollout_update])
 
+    @tracked("_on_attempt_heartbeat")
+    @_with_collections_execute(labels=["rollouts", "attempts"])
+    async def _on_attempt_heartbeat(
+        self, collections: T_collections, rollout_id: str, attempt_id: str
+    ) -> Optional[Tuple[Rollout, Sequence[str]]]:
+        attempt = await collections.attempts.get(
+            {"rollout_id": {"exact": rollout_id}, "attempt_id": {"exact": attempt_id}}
+        )
+        if attempt is None:
+            return None
+        rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
+        if rollout is None:
+            return None
+
+        # Update attempt heartbeat and ensure persistence
+        attempt.last_heartbeat_time = time.time()
+        if attempt.status in ["preparing", "unresponsive"]:
+            attempt.status = "running"
+        await collections.attempts.update([attempt], update_fields=["last_heartbeat_time", "status"])
+
+        # If the status has already timed out or failed, do not change it (but heartbeat is still recorded)
+
+        # Update rollout status if it's the latest attempt
+        rollout_updated: bool = False
+        updated_fields: List[str] = []
+        latest_attempt = await self._unlocked_get_latest_attempt(collections, rollout.rollout_id)
+        if latest_attempt is not None and attempt.attempt_id == latest_attempt.attempt_id:
+            if rollout.status in ["preparing", "queueing", "requeuing"]:
+                # If rollout is currently preparing or queuing, set it to running
+                rollout.status = "running"
+                await collections.rollouts.update([rollout], update_fields=["status"])
+                rollout_updated = True
+                updated_fields = ["status"]
+            # Otherwise, the rollout has succeeded or failed, do nothing
+        return (rollout, updated_fields) if rollout_updated else None
+
     @tracked("wait_for_rollouts")
     @healthcheck_before
     async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
@@ -1489,9 +1505,12 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
                         "status": rollout.status,
                         "mode": rollout.mode if rollout.mode is not None else "unknown",
                     }
+                    duration = cast(float, rollout.end_time) - rollout.start_time
                     await self._tracker.inc_counter("agl.rollouts.total", labels=labels)
                     await self._tracker.observe_histogram(
-                        "agl.rollouts.duration", value=cast(float, rollout.end_time) - rollout.start_time, labels=labels
+                        "agl.rollouts.duration",
+                        value=duration,
+                        labels=labels,
                     )
 
         if not skip_enqueue:
@@ -1665,6 +1684,9 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
     @tracked("_scan_for_unhealthy_rollouts")
     async def _scan_for_unhealthy_rollouts(self) -> None:
         """Perform healthcheck against all running rollouts in the store."""
+        if not await self._should_scan_for_unhealthy_rollouts():
+            return
+
         rollouts, attempts_sync_required = await self._find_and_update_unhealthy_rollouts()
 
         if rollouts:
@@ -1674,6 +1696,25 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
         if attempts_sync_required:
             await self._sync_workers_with_attempts(attempts_sync_required)
 
+    @tracked("_should_scan_for_unhealthy_rollouts")
+    async def _should_scan_for_unhealthy_rollouts(self) -> bool:
+        """Check if the scan for unhealthy rollouts should be performed."""
+        if self._scan_debounce_seconds <= 0:
+            return True
+
+        now = time.time()
+        should_scan = now - self._last_scan_entrance_time >= self._scan_debounce_seconds
+        if not should_scan:
+            return False
+
+        # Someone else may be racing for the same scan. Double-check inside the lock.
+        async with self.collections.atomic(mode="rw", snapshot=self._read_snapshot, labels=["generic"]):
+            now = time.time()
+            if now - self._last_scan_entrance_time < self._scan_debounce_seconds:
+                return False
+            self._last_scan_entrance_time = now
+            return True
+
     @tracked("_find_and_update_unhealthy_rollouts")
     @_with_collections_execute(labels=["rollouts", "attempts"])
     async def _find_and_update_unhealthy_rollouts(
@@ -1702,3 +1743,50 @@ class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
             if worker_sync_required:
                 attempts.append(attempt)
         return rollouts, attempts
+
+
+# _scan_for_unhealthy_rollouts is somehow standalone and automatically invoked.
+COLLECTION_STORE_PUBLIC_METHODS = frozenset(
+    [name for name in LightningStore.__dict__ if not name.startswith("_")] + ["_scan_for_unhealthy_rollouts"]
+)
+
+COLLECTION_STORE_ALL_METHODS = frozenset([name for name in CollectionBasedLightningStore.__dict__])
+
+_UNKNOWN_STORE_METHOD = "unknown"
+
+
+def nearest_lightning_store_method_from_stack() -> Tuple[str, str]:
+    """Stack introspection so that we capture the nearest public API method from the
+    call stack whenever metrics are recorded.
+
+    Returns:
+        A tuple of public method name and nearest private method name.
+    """
+    frame = inspect.currentframe()
+    final_public_method_name = final_private_method_name = _UNKNOWN_STORE_METHOD
+    try:
+        if frame is not None:
+            frame = frame.f_back
+            while frame is not None:
+                self_obj = frame.f_locals.get("self")
+                public_method_name = frame.f_locals.get("public_method_name")
+                private_method_name = frame.f_locals.get("private_method_name")
+                if (
+                    final_public_method_name == _UNKNOWN_STORE_METHOD
+                    and public_method_name in COLLECTION_STORE_PUBLIC_METHODS
+                    and isinstance(self_obj, LightningStore)
+                ):
+                    final_public_method_name = public_method_name
+                if (
+                    final_private_method_name == _UNKNOWN_STORE_METHOD
+                    and private_method_name in COLLECTION_STORE_ALL_METHODS
+                    and isinstance(self_obj, LightningStore)
+                ):
+                    final_private_method_name = private_method_name
+                frame = frame.f_back
+    except Exception as exc:
+        logger.debug("Error during stack introspection for LightningStore method: %s", exc)
+    finally:
+        del frame
+
+    return final_public_method_name, final_private_method_name

agentlightning/store/memory.py

@@ -73,12 +73,16 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
     Thread-safe and async-compatible but data is not persistent.
 
     Args:
+        thread_safe: Whether the store is thread-safe.
         eviction_memory_threshold: The threshold for evicting spans in bytes.
             By default, it's 70% of the total VRAM available.
         safe_memory_threshold: The threshold for safe memory usage in bytes.
             By default, it's 80% of the eviction threshold.
         span_size_estimator: A function to estimate the size of a span in bytes.
             By default, it's a simple size estimator that uses sys.getsizeof.
+        tracker: The metrics tracker to use.
+        scan_debounce_seconds: The debounce time for the scan for unhealthy rollouts.
+            Set to 0 to disable debouncing.
     """
 
     def __init__(
@@ -89,10 +93,12 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
         safe_memory_threshold: float | int | None = None,
         span_size_estimator: Callable[[Span], int] | None = None,
         tracker: MetricsBackend | None = None,
+        scan_debounce_seconds: float = 10.0,
     ):
         super().__init__(
             collections=InMemoryLightningCollections(lock_type="thread" if thread_safe else "asyncio", tracker=tracker),
             tracker=tracker,
+            scan_debounce_seconds=scan_debounce_seconds,
         )
 
         self._thread_safe = thread_safe

agentlightning/store/mongo.py

@@ -7,19 +7,7 @@ import hashlib
 import logging
 import time
 import uuid
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    List,
-    Mapping,
-    Optional,
-    Sequence,
-    TypeVar,
-    Union,
-)
-
-from pymongo import AsyncMongoClient
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, TypeVar, Union
 
 from agentlightning.types import Attempt, AttemptedRollout, Rollout
 from agentlightning.utils.metrics import MetricsBackend
@@ -43,26 +31,28 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
     Data is persistent and can be shared between multiple processes.
 
     Args:
-        client: The MongoDB client. Could be a string URI or an instance of AsyncMongoClient.
-        database: The MongoDB database. Could be a string name or an instance of AsyncDatabase.
-            You must provide at least one of client or database.
+        mongo_uri: MongoDB connection string (defaults to local replica set).
+        mongo_client_kwargs: Extra keyword arguments forwarded to `AsyncMongoClient`.
+        database: The MongoDB database name. Defaults to ``agentlightning``.
         partition_id: The partition id. Useful when sharing the database among multiple Agent-lightning trainers.
+        tracker: The metrics tracker to use.
+        scan_debounce_seconds: The debounce time for the scan for unhealthy rollouts.
+            Set to 0 to disable debouncing.
     """
 
     def __init__(
         self,
         *,
-        client: AsyncMongoClient[Mapping[str, Any]] | str,
+        mongo_uri: str = "mongodb://localhost:27017/?replicaSet=rs0",
+        mongo_client_kwargs: Mapping[str, Any] | None = None,
         database_name: str | None = None,
         partition_id: str | None = None,
         tracker: MetricsBackend | None = None,
+        scan_debounce_seconds: float = 10.0,
     ) -> None:
-        self._auto_created_client = False
-        if isinstance(client, str):
-            self._client = AsyncMongoClient[Mapping[str, Any]](client)
-            self._auto_created_client = True
-        else:
-            self._client = client
+        self._mongo_uri = mongo_uri
+        self._mongo_client_kwargs = dict(mongo_client_kwargs or {})
+
         if database_name is None:
             database_name = "agentlightning"
             logger.info("No database name provided, using default 'agentlightning'")
@@ -71,7 +61,10 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
             partition_id = _generate_partition_id()
             logger.info("No partition id provided, generated a new one: %s", partition_id)
 
-        self._client_pool = MongoClientPool(self._client)
+        self._client_pool = MongoClientPool[Mapping[str, Any]](
+            mongo_uri=self._mongo_uri,
+            mongo_client_kwargs=self._mongo_client_kwargs,
+        )
 
         super().__init__(
             collections=MongoLightningCollections(
@@ -81,6 +74,7 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
                 tracker=tracker,
             ),
             tracker=tracker,
+            scan_debounce_seconds=scan_debounce_seconds,
         )
 
     @property
@@ -96,9 +90,6 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
     async def close(self) -> None:
         """Close the store by closing the client pool."""
         await self._client_pool.close()
-        # If I created the client, I should close it too.
-        if self._auto_created_client:
-            await self._client.close()
 
     @tracked("wait_for_rollouts")
     @healthcheck_before

agentlightning/utils/metrics.py

@@ -743,6 +743,9 @@ class PrometheusMetricsBackend(MetricsBackend):
 
     Thread-safety: Registration is protected by a lock. Metric updates assume metrics
     are registered during initialization and then remain stable.
+
+    Due to the nature of Prometheus, this backend is only suitable for recording high-volume metrics.
+    Low-volume metrics might be lost if the event has only appeared once.
     """
 
     def __init__(self) -> None:

docker/compose.grafana.yml

@@ -5,7 +5,7 @@ services:
       - "--config.file=/etc/prometheus/prometheus.yml"
       - "--storage.tsdb.path=/prometheus"
     volumes:
-      - ./prometheus.memory-store.yml:/etc/prometheus/prometheus.yml:ro
+      - ./prometheus/prometheus.base.yml:/etc/prometheus/prometheus.yml:ro
       - ${AGL_MONITORING_DATA_PATH:?Set AGL_MONITORING_DATA_PATH to the metrics directory}/prometheus:/prometheus
     ports:
       - "9090:9090"

docker/compose.prometheus-memory-store.yml

@@ -4,7 +4,27 @@ services:
       file: compose.store.yml
       service: app
 
+    depends_on:
+      - app-exporter  # Wait for the exporter to be ready first
+
     command: agl store --host 0.0.0.0 --port 4747 --tracker console prometheus --backend memory
+    environment:
+      - PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus_multiproc
+    volumes:
+      - prometheus_multiproc:/tmp/prometheus_multiproc
+
+  app-exporter:
+    build:
+      context: ../
+      dockerfile: docker/Dockerfile.dev
+
+    command: agl prometheus --host 0.0.0.0 --port 4748
+    ports:
+      - "4748:4748"
+    environment:
+      - PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus_multiproc
+    volumes:
+      - prometheus_multiproc:/tmp/prometheus_multiproc
 
   node-exporter:
     image: prom/node-exporter:latest
@@ -22,10 +42,11 @@ services:
       - "--storage.tsdb.path=/prometheus"
       - "--storage.tsdb.retention.time=1h"
     volumes:
-      - ./prometheus.memory-store.yml:/etc/prometheus/prometheus.yml:ro
+      - ./prometheus/prometheus.base.yml:/etc/prometheus/prometheus.yml:ro
       - ./data/prometheus:/prometheus
     depends_on:
       - app
+      - app-exporter
       - node-exporter
     ports:
       - "9090:9090"
@@ -51,3 +72,10 @@ services:
       - GF_AUTH_ANONYMOUS_ORG_ROLE=Admin
       - GF_AUTH_DISABLE_LOGIN_FORM=true
       - GF_DASHBOARDS_DEFAULT_HOME_DASHBOARD_PATH=/var/lib/grafana/dashboards/agentlightning.json
+
+volumes:
+  prometheus_multiproc:
+    driver: local
+    driver_opts:
+      type: tmpfs
+      device: tmpfs

docker/compose.prometheus-mongo-store.yml

@@ -21,6 +21,7 @@ services:
 
     depends_on:
       - mongo
+      - app-exporter  # Wait for the exporter to be ready first
 
     command:
       - /bin/bash
@@ -30,6 +31,23 @@ services:
             --tracker console prometheus --backend mongo \
             --mongo-uri mongodb://mongo:27017/?replicaSet=rs0 \
             --n-workers ${AGL_STORE_N_WORKERS:-32}
+    environment:
+      - PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus_multiproc
+    volumes:
+      - prometheus_multiproc:/tmp/prometheus_multiproc
+
+  app-exporter:
+    build:
+      context: ../
+      dockerfile: docker/Dockerfile.dev
+
+    command: agl prometheus --host 0.0.0.0 --port 4748
+    ports:
+      - "4748:4748"
+    environment:
+      - PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus_multiproc
+    volumes:
+      - prometheus_multiproc:/tmp/prometheus_multiproc
 
   mongodb-exporter:
     image: percona/mongodb_exporter:0.47.1
@@ -58,10 +76,11 @@ services:
       - "--storage.tsdb.path=/prometheus"
       - "--storage.tsdb.retention.time=1h"
     volumes:
-      - ./prometheus.mongo-store.yml:/etc/prometheus/prometheus.yml:ro
+      - ./prometheus/prometheus.mongo.yml:/etc/prometheus/prometheus.yml:ro
       - ./data/prometheus:/prometheus
     depends_on:
       - app
+      - app-exporter
       - mongodb-exporter
       - node-exporter
     ports:
@@ -88,3 +107,10 @@ services:
       - GF_AUTH_ANONYMOUS_ORG_ROLE=Admin
       - GF_AUTH_DISABLE_LOGIN_FORM=true
       - GF_DASHBOARDS_DEFAULT_HOME_DASHBOARD_PATH=/var/lib/grafana/dashboards/agentlightning.json
+
+volumes:
+  prometheus_multiproc:
+    driver: local
+    driver_opts:
+      type: tmpfs
+      device: tmpfs

docker/grafana/dashboards/agentlightning.json

@@ -44,7 +44,7 @@
       },
       "targets": [
         {
-          "expr": "sum by (method, path) (rate(agl_http_total[15s]))",
+          "expr": "sum by (method, path) (rate(agl_http_latency_count[15s]))",
           "legendFormat": "{{method}} {{path}}"
         }
       ],
@@ -112,7 +112,7 @@
       },
       "targets": [
         {
-          "expr": "sum by (method, path, status) (rate(agl_http_total{status!~\"2..\"}[15s]))",
+          "expr": "sum by (method, path, status) (rate(agl_http_latency_count{status!~\"2..\"}[15s]))",
           "legendFormat": "{{method}} {{path}} {{status}}"
         }
       ],
@@ -155,7 +155,7 @@
       },
       "targets": [
         {
-          "expr": "sum by (method) (rate(agl_store_total[15s]))",
+          "expr": "sum by (method) (rate(agl_store_latency_count[15s]))",
           "legendFormat": "{{method}}"
         }
       ],
@@ -223,7 +223,7 @@
       },
       "targets": [
         {
-          "expr": "sum by (method, status) (rate(agl_store_total{status!=\"OK\"}[15s]))",
+          "expr": "sum by (method, status) (rate(agl_store_latency_count{status!=\"OK\"}[15s]))",
           "legendFormat": "{{method}} {{status}}"
         }
       ],
@@ -266,8 +266,8 @@
       },
       "targets": [
         {
-          "expr": "sum by (store_method, operation) (rate(agl_collections_total[15s]))",
-          "legendFormat": "{{store_method}} {{operation}}"
+          "expr": "sum by (store_pubmeth, operation) (rate(agl_collections_latency_count[15s]))",
+          "legendFormat": "{{store_pubmeth}} {{operation}}"
         }
       ],
       "title": "Collection Ops / Sec (by Store Method)",
@@ -296,8 +296,8 @@
       },
       "targets": [
         {
-          "expr": "sum by (store_method, collection, status) (rate(agl_collections_total{status!=\"OK\"}[15s]))",
-          "legendFormat": "{{store_method}} {{collection}} {{status}}"
+          "expr": "sum by (store_pubmeth, collection, status) (rate(agl_collections_latency_count{status!=\"OK\"}[15s]))",
+          "legendFormat": "{{store_pubmeth}} {{collection}} {{status}}"
         }
       ],
       "title": "Collection Errors / Sec",
@@ -364,8 +364,8 @@
       },
       "targets": [
         {
-          "expr": "histogram_quantile(0.95, sum by (le, store_method) (rate(agl_collections_latency_bucket[15s])))",
-          "legendFormat": "{{store_method}}"
+          "expr": "histogram_quantile(0.95, sum by (le, store_pubmeth) (rate(agl_collections_latency_bucket[15s])))",
+          "legendFormat": "{{store_pubmeth}}"
         }
       ],
       "title": "Collection Latency P95 (by Store Method)",
@@ -407,7 +407,7 @@
       },
       "targets": [
         {
-          "expr": "sum by (status) (rate(agl_rollouts_total[15s]))",
+          "expr": "sum by (status) (rate(agl_rollouts_duration_count[15s]))",
           "legendFormat": "{{status}}"
         }
       ],
@@ -437,7 +437,7 @@
       },
       "targets": [
         {
-          "expr": "sum by (mode) (rate(agl_rollouts_total[15s]))",
+          "expr": "sum by (mode) (rate(agl_rollouts_duration_count[15s]))",
           "legendFormat": "{{mode}}"
         }
       ],

docker/prometheus/prometheus.base.yml

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

docker/prometheus/prometheus.mongo.yml

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

docs/reference/cli.md

@@ -14,17 +14,18 @@
 ## agl
 
 ```text
-usage: agl [-h] {vllm,store,agentops}
+usage: agl [-h] {vllm,store,prometheus,agentops}
 
 Agent Lightning CLI entry point.
 
 Available subcommands:
-  vllm      Run the vLLM CLI with Agent Lightning instrumentation.
-  store     Run a LightningStore server.
-  agentops  Start the AgentOps server manager.
+  vllm        Run the vLLM CLI with Agent Lightning instrumentation.
+  store       Run a LightningStore server.
+  prometheus  Serve Prometheus metrics from the multiprocess registry.
+  agentops    Start the AgentOps server manager.
 
 positional arguments:
-  {vllm,store,agentops}
+  {vllm,store,prometheus,agentops}
                         Subcommand to run.
 
 options:
@@ -73,6 +74,26 @@ options:
   --port PORT  Port to run the server on
 ```
 
+## agl prometheus
+
+Expose the Prometheus multiprocess registry on a dedicated FastAPI server. This is useful when the main LightningStore service is under heavy load; exporters can scrape this auxiliary endpoint instead.
+
+```text
+usage: agl prometheus [-h] [--host HOST] [--port PORT] [--metrics-path METRICS_PATH] [--log-level {DEBUG,INFO,WARNING,ERROR}] [--access-log]
+
+Serve Prometheus metrics outside the LightningStore server.
+
+options:
+  -h, --help            show this help message and exit
+  --host HOST           Host to bind the metrics server to.
+  --port PORT           Port to expose the Prometheus metrics on.
+  --metrics-path METRICS_PATH
+                        HTTP path used to expose metrics. Must start with '/' and not be the root path.
+  --log-level {DEBUG,INFO,WARNING,ERROR}
+                        Configure the logging level for the metrics server.
+  --access-log          Enable uvicorn access logs. Disabled by default to reduce noise.
+```
+
 ## agl agentops
 
 Start a mock AgentOps server to bypass the online service of AgentOps.

tests/benchmark/benchmark_store.py

@@ -157,11 +157,13 @@ class AlgorithmBatch(agl.Algorithm):
 
             pending = {rollout_id: task_name for rollout_id, task_name in batch_rollouts}
             completed_ids: Set[str] = set()
+            completed_ids_last_updated: int = 0
             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:
@@ -171,6 +173,18 @@ class AlgorithmBatch(agl.Algorithm):
                     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
+
+                # Check and warn for stale rollouts
+                if complete_ids_updated:
+                    completed_ids_last_updated = 0
+                else:
+                    completed_ids_last_updated += 1
+                    if completed_ids_last_updated >= 10:
+                        unfinished_ids = set(rollout_id for rollout_id, _ in batch_rollouts) - completed_ids
+                        print(f"Stale rollouts: {unfinished_ids}")
+                        completed_ids_last_updated = 0
+
                 await asyncio.sleep(5.0)
 
     async def algorithm_batch_with_completion_threshold(self, total_tasks: int, batch_size: int, remaining_tasks: int):

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/store/conftest.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import os
 import time
 from itertools import count
-from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Sequence
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Mapping, Sequence
 from unittest.mock import Mock
 from uuid import uuid4
 
@@ -15,6 +15,7 @@ from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel, Field
 from pytest import FixtureRequest
 
+from agentlightning.store import collection_based
 from agentlightning.store.base import LightningStore
 from agentlightning.store.collection import DequeBasedQueue, DictBasedKeyValue, KeyValue, ListBasedCollection, Queue
 from agentlightning.store.collection.base import Collection
@@ -26,6 +27,10 @@ if TYPE_CHECKING:
 
 __all__ = [
     "inmemory_store",
+    "inmemory_debounced_store",
+    "mongo_debounced_store",
+    "debounced_store",
+    "fake_time",
     "mock_readable_span",
     "sample_items",
     "sample_collection",
@@ -35,16 +40,25 @@ __all__ = [
     "dict_key_value",
     "dict_key_value_data",
     "temporary_mongo_database",
+    "mongo_uri",
+    "mongo_client_kwargs",
 ]
 
 
 mongo_uri = os.getenv("AGL_TEST_MONGO_URI", "mongodb://localhost:27017/?replicaSet=rs0")
+mongo_client_kwargs: Dict[str, Any] = {"serverSelectionTimeoutMS": 5000}
 
 
 @pytest.fixture
 def inmemory_store() -> InMemoryLightningStore:
     """Create a fresh InMemoryLightningStore instance."""
-    return InMemoryLightningStore()
+    return InMemoryLightningStore(scan_debounce_seconds=0)
+
+
+@pytest.fixture
+def inmemory_debounced_store(fake_time: _FakeTime) -> InMemoryLightningStore:
+    """Create an InMemoryLightningStore configured with scan debouncing."""
+    return InMemoryLightningStore(scan_debounce_seconds=5.0)
 
 
 @pytest_asyncio.fixture
@@ -52,7 +66,29 @@ async def mongo_store(temporary_mongo_database: AsyncDatabase[Any]):
     """Fixture for MongoDB store implementation."""
     from agentlightning.store.mongo import MongoLightningStore
 
-    db = MongoLightningStore(client=temporary_mongo_database.client, database_name=temporary_mongo_database.name)
+    db = MongoLightningStore(
+        mongo_uri=mongo_uri,
+        mongo_client_kwargs=mongo_client_kwargs,
+        database_name=temporary_mongo_database.name,
+        scan_debounce_seconds=0,
+    )
+    try:
+        yield db
+    finally:
+        await db.close()
+
+
+@pytest_asyncio.fixture
+async def mongo_debounced_store(fake_time: _FakeTime, temporary_mongo_database: AsyncDatabase[Any]):
+    """Fixture for MongoDB store implementation with scan debouncing."""
+    from agentlightning.store.mongo import MongoLightningStore
+
+    db = MongoLightningStore(
+        mongo_uri=mongo_uri,
+        mongo_client_kwargs=mongo_client_kwargs,
+        database_name=temporary_mongo_database.name,
+        scan_debounce_seconds=5.0,
+    )
     try:
         yield db
     finally:
@@ -70,6 +106,17 @@ def store_fixture(request: FixtureRequest) -> AsyncGenerator[LightningStore, Non
     return request.getfixturevalue(request.param)
 
 
+@pytest.fixture(
+    params=[
+        "inmemory_debounced_store",
+        pytest.param("mongo_debounced_store", marks=pytest.mark.mongo),
+    ]
+)
+def debounced_store(request: FixtureRequest) -> LightningStore:
+    """Parameterized fixture for debounced store implementations."""
+    return request.getfixturevalue(request.param)
+
+
 @pytest.fixture
 def mock_readable_span() -> ReadableSpan:
     """Create a mock ReadableSpan for testing."""
@@ -123,11 +170,35 @@ class QueueItem(BaseModel):
     idx: int
 
 
+class _FakeTime:
+    """Simple controllable clock for scan debouncing tests."""
+
+    def __init__(self, start: float = 0.0) -> None:
+        self._value = start
+
+    def time(self) -> float:
+        return self._value
+
+    def set(self, value: float) -> None:
+        self._value = value
+
+    def advance(self, delta: float) -> None:
+        self._value += delta
+
+
+@pytest.fixture
+def fake_time(monkeypatch: pytest.MonkeyPatch) -> _FakeTime:
+    """Patch collection_based.time.time with a controllable clock."""
+    controller = _FakeTime()
+    monkeypatch.setattr(collection_based.time, "time", controller.time)
+    return controller
+
+
 @pytest_asyncio.fixture
 async def mongo_client():
     from pymongo import AsyncMongoClient
 
-    client = AsyncMongoClient[Any](mongo_uri, serverSelectionTimeoutMS=5000)
+    client = AsyncMongoClient[Any](mongo_uri, **mongo_client_kwargs)
     try:
         await client.admin.command("ping")
     except Exception as exc:  # depends on external service
@@ -270,7 +341,9 @@ def sample_collection_memory(sample_items: Sequence[SampleItem]) -> ListBasedCol
 async def sample_collection_mongo(temporary_mongo_database: AsyncDatabase[Any], sample_items: Sequence[SampleItem]):
     from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool
 
-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
         collection = MongoBasedCollection(
             client_pool,
             temporary_mongo_database.name,
@@ -307,7 +380,9 @@ def deque_queue_memory() -> DequeBasedQueue[QueueItem]:
 async def deque_queue_mongo(temporary_mongo_database: AsyncDatabase[Any]):
     from agentlightning.store.collection.mongo import MongoBasedQueue, MongoClientPool
 
-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
         queue = MongoBasedQueue[QueueItem](
             client_pool,
             temporary_mongo_database.name,
@@ -347,7 +422,9 @@ def dict_key_value_memory(dict_key_value_data: Dict[str, int]) -> DictBasedKeyVa
 async def dict_key_value_mongo(temporary_mongo_database: AsyncDatabase[Any], dict_key_value_data: Dict[str, int]):
     from agentlightning.store.collection.mongo import MongoBasedKeyValue, MongoClientPool
 
-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
         key_value = MongoBasedKeyValue[str, int](
             client_pool,
             temporary_mongo_database.name,

tests/store/test_collection_memory.py

@@ -18,20 +18,22 @@ from typing import (
     Sequence,
     Tuple,
     Union,
+    cast,
 )
 from uuid import uuid4
 
+import pydantic
 import pytest
 from pydantic import BaseModel
 
 import agentlightning.store.collection.memory as memory_module
-from agentlightning.store.collection import DequeBasedQueue, DictBasedKeyValue, ListBasedCollection
+from agentlightning.store.collection import DequeBasedQueue, DictBasedKeyValue, KeyValue, ListBasedCollection
 from agentlightning.store.collection.base import Collection
 from agentlightning.store.collection.memory import _item_matches_filters  # pyright: ignore[reportPrivateUsage]
 from agentlightning.store.collection.memory import _LoopAwareAsyncLock  # pyright: ignore[reportPrivateUsage]
 from agentlightning.store.collection.memory import _ThreadSafeAsyncLock  # pyright: ignore[reportPrivateUsage]
 from agentlightning.types import Rollout
-from tests.store.conftest import QueueItem, SampleItem
+from tests.store.conftest import QueueItem, SampleItem, mongo_client_kwargs, mongo_uri
 
 if TYPE_CHECKING:
     from pymongo.asynchronous.database import AsyncDatabase
@@ -117,11 +119,13 @@ async def test_list_collection_insert_rejects_duplicate_payload(sample_collectio
     dup_a = SampleItem(partition="omega", index=1, name="dup-a", status="new")
     dup_b = SampleItem(partition="omega", index=1, name="dup-b", status="new")
 
-    with pytest.raises(ValueError, match="duplicate primary key"):
+    with pytest.raises(ValueError, match=r"(duplicated|Duplicated) primary key"):
         await sample_collection.insert([dup_a, dup_b])
 
-    assert await sample_collection.size() == starting_size
-    assert await sample_collection.get({"partition": {"exact": "omega"}}) is None
+    if isinstance(sample_collection, ListBasedCollection):
+        # Only ListBasedCollection supports this rejecting duplicate items within the same insert batch.
+        assert await sample_collection.size() == starting_size
+        assert await sample_collection.get({"partition": {"exact": "omega"}}) is None
 
 
 @pytest.mark.asyncio()
@@ -865,6 +869,150 @@ async def test_dict_key_value_pop_returns_default(dict_key_value: DictBasedKeyVa
     assert await dict_key_value.size() == 1
 
 
+@pytest.mark.asyncio()
+async def test_key_value_inc_updates_existing(dict_key_value: KeyValue[str, int]) -> None:
+    new_value = await dict_key_value.inc("alpha", 2)
+    assert new_value == 3
+    assert await dict_key_value.get("alpha") == 3
+
+
+@pytest.mark.asyncio()
+async def test_key_value_inc_initializes_missing(dict_key_value: KeyValue[str, int]) -> None:
+    new_value = await dict_key_value.inc("gamma", 4)
+    assert new_value == 4
+    assert await dict_key_value.has("gamma")
+
+
+@pytest.mark.asyncio()
+async def test_key_value_inc_rejects_non_numeric_amount(dict_key_value: KeyValue[str, int]) -> None:
+    with pytest.raises(TypeError):
+        await dict_key_value.inc("alpha", cast(Any, "invalid"))
+
+
+@pytest.mark.asyncio()
+async def test_key_value_chmax_updates_existing(dict_key_value: KeyValue[str, int]) -> None:
+    new_value = await dict_key_value.chmax("alpha", 10)
+    assert new_value == 10
+    assert await dict_key_value.get("alpha") == 10
+
+
+@pytest.mark.asyncio()
+async def test_key_value_chmax_ignores_smaller(dict_key_value: KeyValue[str, int]) -> None:
+    initial = await dict_key_value.get("alpha")
+    result = await dict_key_value.chmax("alpha", 0)
+    assert result == initial
+    assert await dict_key_value.get("alpha") == initial
+
+
+@pytest.mark.asyncio()
+async def test_key_value_chmax_initializes_missing(dict_key_value: KeyValue[str, int]) -> None:
+    result = await dict_key_value.chmax("gamma", 7)
+    assert result == 7
+    assert await dict_key_value.get("gamma") == 7
+
+
+@pytest.mark.asyncio()
+async def test_key_value_chmax_rejects_non_numeric_value(dict_key_value: KeyValue[str, int]) -> None:
+    with pytest.raises(TypeError):
+        await dict_key_value.chmax("alpha", cast(Any, "wrong"))
+
+
+@pytest.mark.asyncio()
+async def test_dict_key_value_inc_rejects_non_numeric_value(dict_key_value_memory: DictBasedKeyValue[str, Any]) -> None:
+    await dict_key_value_memory.set("alpha", cast(Any, "na"))
+    with pytest.raises(TypeError):
+        await dict_key_value_memory.inc("alpha", 1)
+
+
+@pytest.mark.asyncio()
+async def test_dict_key_value_chmax_rejects_non_numeric_value(
+    dict_key_value_memory: DictBasedKeyValue[str, Any],
+) -> None:
+    await dict_key_value_memory.set("alpha", cast(Any, "na"))
+    with pytest.raises(TypeError):
+        await dict_key_value_memory.chmax("alpha", 1)
+
+
+@pytest.mark.mongo
+@pytest.mark.asyncio()
+async def test_mongo_key_value_inc_rejects_non_numeric_value(temporary_mongo_database: AsyncDatabase[Any]) -> None:
+    from agentlightning.store.collection.mongo import MongoBasedKeyValue, MongoClientPool
+
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
+        key_value = MongoBasedKeyValue[str, int](
+            client_pool,
+            temporary_mongo_database.name,
+            f"kv-inc-{uuid4().hex}",
+            "partition-inc",
+            str,
+            int,
+        )
+        collection = await key_value.ensure_collection()
+        await collection.insert_one(
+            {
+                "partition_id": "partition-inc",
+                "key": "alpha",
+                "value": "oops",
+            }
+        )
+
+        with pytest.raises(TypeError):
+            await key_value.inc("alpha", 1)
+
+
+@pytest.mark.mongo
+@pytest.mark.asyncio()
+async def test_mongo_key_value_chmax_behaves_like_max(temporary_mongo_database: AsyncDatabase[Any]) -> None:
+    from agentlightning.store.collection.mongo import MongoBasedKeyValue, MongoClientPool
+
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
+        key_value = MongoBasedKeyValue[str, int](
+            client_pool,
+            temporary_mongo_database.name,
+            f"kv-chmax-{uuid4().hex}",
+            "partition-chmax",
+            str,
+            int,
+        )
+        assert await key_value.chmax("alpha", 5) == 5
+        assert await key_value.chmax("alpha", 3) == 5
+        assert await key_value.chmax("alpha", 9) == 9
+        assert await key_value.get("alpha") == 9
+
+
+@pytest.mark.mongo
+@pytest.mark.asyncio()
+async def test_mongo_key_value_chmax_rejects_non_numeric_value(temporary_mongo_database: AsyncDatabase[Any]) -> None:
+    from agentlightning.store.collection.mongo import MongoBasedKeyValue, MongoClientPool
+
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
+        key_value = MongoBasedKeyValue[str, int](
+            client_pool,
+            temporary_mongo_database.name,
+            f"kv-chmax-bad-{uuid4().hex}",
+            "partition-chmax-bad",
+            str,
+            int,
+        )
+        collection = await key_value.ensure_collection()
+        await collection.insert_one(
+            {
+                "partition_id": "partition-chmax-bad",
+                "key": "alpha",
+                "value": "oops",
+            }
+        )
+
+        with pytest.raises((TypeError, pydantic.ValidationError)):
+            await key_value.chmax("alpha", 1)
+
+
 def test_thread_safe_async_lock_blocks_threads() -> None:
     lock = _ThreadSafeAsyncLock()
     allow_second = threading.Event()
@@ -1071,7 +1219,9 @@ async def test_mongo_based_sanity_check(temporary_mongo_database: AsyncDatabase[
         MongoClientPool,
     )
 
-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
         collection = MongoBasedCollection[Any](
             client_pool, temporary_mongo_database.name, "test", "test-123", ["rollout_id"], Rollout
         )
@@ -1110,31 +1260,6 @@ async def test_mongo_based_sanity_check(temporary_mongo_database: AsyncDatabase[
         assert not await span_kv.has("span-123")
 
 
-@pytest.mark.mongo
-@pytest.mark.asyncio()
-async def test_mongo_based_collection_rejects_duplicate_payload(temporary_mongo_database: AsyncDatabase[Any]) -> None:
-    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool
-
-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
-        collection = MongoBasedCollection[Any](
-            client_pool,
-            temporary_mongo_database.name,
-            f"duplicate-check-{uuid4().hex}",
-            "partition-dup",
-            ["rollout_id"],
-            Rollout,
-        )
-        await collection.ensure_collection()
-        start_time = time.time()
-        first = Rollout(rollout_id="dup-rollout", input="payload", start_time=start_time, status="running")
-        duplicate = Rollout(rollout_id="dup-rollout", input="payload", start_time=start_time, status="running")
-
-        with pytest.raises(ValueError, match="duplicate primary key"):
-            await collection.insert([first, duplicate])
-
-        assert await collection.size() == 0
-
-
 @pytest.mark.mongo
 @pytest.mark.asyncio()
 async def test_mongo_ensure_collection_creates_partition_scoped_index(
@@ -1143,7 +1268,9 @@ async def test_mongo_ensure_collection_creates_partition_scoped_index(
     from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool
 
     collection_name = f"ensure-{uuid4().hex}"
-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
         collection = MongoBasedCollection[Any](
             client_pool,
             temporary_mongo_database.name,
@@ -1173,7 +1300,9 @@ async def test_mongo_ensure_collection_survives_concurrent_calls(temporary_mongo
     collection_name = f"ensure-{uuid4().hex}"
 
     async def ensure_once() -> None:
-        async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+        async with MongoClientPool[Mapping[str, Any]](
+            mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+        ) as client_pool:
             collection = MongoBasedCollection(
                 client_pool,
                 temporary_mongo_database.name,
@@ -1204,7 +1333,9 @@ async def test_mongo_ensure_collection_repeats_without_altering_indexes(
     from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool
 
     collection_name = f"ensure-{uuid4().hex}"
-    async with MongoClientPool(temporary_mongo_database.client) as client_pool:
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
         collection = MongoBasedCollection(
             client_pool, temporary_mongo_database.name, collection_name, "partition-repeat", ["index"], SampleItem
         )
@@ -1225,7 +1356,9 @@ async def _with_mongo_collections(
 ) -> Any:
     from agentlightning.store.collection.mongo import MongoClientPool, MongoLightningCollections
 
-    async with MongoClientPool(db.client) as client_pool:
+    async with MongoClientPool[Mapping[str, Any]](
+        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
+    ) as client_pool:
         collections = MongoLightningCollections(
             client_pool=client_pool,
             database_name=db.name,

tests/store/test_implementation.py

@@ -19,12 +19,13 @@ import asyncio
 import logging
 import sys
 import time
-from typing import List, Optional, Sequence, cast
+from typing import Any, List, Optional, Protocol, Sequence, cast
 from unittest.mock import Mock
 
 import pytest
 from pydantic import BaseModel
 
+from agentlightning.store import CollectionBasedLightningStore
 from agentlightning.store.base import UNSET, LightningStore
 from agentlightning.store.memory import InMemoryLightningStore, estimate_model_size
 from agentlightning.types import (
@@ -45,7 +46,11 @@ from agentlightning.types import (
     TraceStatus,
 )
 
-# Typing tests
+
+class FakeTimeController(Protocol):
+    def set(self, value: float) -> None: ...
+
+    def advance(self, delta: float) -> None: ...
 
 
 def test_paginated_result_behaves_like_sequence() -> None:
@@ -1308,7 +1313,8 @@ async def test_add_many_spans_handles_mixed_rollouts_and_attempts(store_fixture:
     duplicate_first = _build_span(1, first.rollout_id, first.attempt.attempt_id)
 
     stored = await store_fixture.add_many_spans([span_first, span_retry, span_second, duplicate_first])
-    assert {span.span_id for span in stored} == {span_first.span_id, span_retry.span_id, span_second.span_id}
+    if isinstance(store_fixture, InMemoryLightningStore):
+        assert {span.span_id for span in stored} == {span_first.span_id, span_retry.span_id, span_second.span_id}
 
     spans_first = await store_fixture.query_spans(first.rollout_id)
     assert {span.span_id for span in spans_first} >= {span_first.span_id, span_retry.span_id}
@@ -3035,7 +3041,7 @@ async def test_healthcheck_unresponsive_behavior(store_fixture: LightningStore,
     """Test that healthcheck detects and handles unresponsive conditions."""
     # Create rollout with short unresponsive timeout but no retry for unresponsive
     config = RolloutConfig(
-        unresponsive_seconds=0.1,  # Very short unresponsive timeout
+        unresponsive_seconds=0.2,  # Very short unresponsive timeout
         max_attempts=3,
         retry_condition=["timeout"],  # Note: "unresponsive" not in retry_condition
     )
@@ -3054,7 +3060,7 @@ async def test_healthcheck_unresponsive_behavior(store_fixture: LightningStore,
     assert running_attempts[0].last_heartbeat_time is not None
 
     # Wait for unresponsive timeout
-    await asyncio.sleep(0.15)  # Wait longer than unresponsive_seconds
+    await asyncio.sleep(0.25)  # Wait longer than unresponsive_seconds
 
     # Verify attempt was marked as unresponsive
     attempts_after = await store_fixture.query_attempts(rollout.rollout_id)
@@ -3383,3 +3389,56 @@ async def test_query_resources_returns_all_fields(store_fixture: LightningStore)
         assert res.update_time > 0
         assert res.version >= 1
         assert res.resources is not None
+
+
+@pytest.mark.asyncio
+async def test_scan_debounce_allows_initial_scan(
+    fake_time: FakeTimeController, debounced_store: CollectionBasedLightningStore[Any]
+) -> None:
+    """The first watchdog scan should run immediately even with debouncing enabled."""
+    fake_time.set(100.0)
+
+    assert await debounced_store._should_scan_for_unhealthy_rollouts()  # pyright: ignore[reportPrivateUsage]
+
+
+@pytest.mark.asyncio
+async def test_scan_debounce_blocks_until_interval(
+    fake_time: FakeTimeController, debounced_store: CollectionBasedLightningStore[Any]
+) -> None:
+    """Subsequent scans wait until the debounce window elapses."""
+    fake_time.set(200.0)
+
+    assert await debounced_store._should_scan_for_unhealthy_rollouts()  # pyright: ignore[reportPrivateUsage]
+    assert not await debounced_store._should_scan_for_unhealthy_rollouts()  # pyright: ignore[reportPrivateUsage]
+
+    fake_time.advance(debounced_store._scan_debounce_seconds - 1)  # pyright: ignore[reportPrivateUsage]
+    assert not await debounced_store._should_scan_for_unhealthy_rollouts()  # pyright: ignore[reportPrivateUsage]
+
+    fake_time.advance(1.0)
+    assert await debounced_store._should_scan_for_unhealthy_rollouts()  # pyright: ignore[reportPrivateUsage]
+
+
+@pytest.mark.asyncio
+async def test_scan_debounce_allows_single_concurrent_scan(
+    fake_time: FakeTimeController, debounced_store: CollectionBasedLightningStore[Any]
+) -> None:
+    """When multiple coroutines race, only one should trigger the scan."""
+    fake_time.set(300.0)
+
+    results = await asyncio.gather(
+        debounced_store._should_scan_for_unhealthy_rollouts(),  # pyright: ignore[reportPrivateUsage]
+        debounced_store._should_scan_for_unhealthy_rollouts(),  # pyright: ignore[reportPrivateUsage]
+        debounced_store._should_scan_for_unhealthy_rollouts(),  # pyright: ignore[reportPrivateUsage]
+    )
+
+    assert results.count(True) == 1
+    assert results.count(False) == 2
+
+
+@pytest.mark.asyncio
+async def test_scan_debounce_disabled_when_zero(store_fixture: CollectionBasedLightningStore[Any]) -> None:
+    """Setting debounce to zero should run the scan every time."""
+    assert store_fixture._scan_debounce_seconds == 0  # pyright: ignore[reportPrivateUsage]
+
+    assert await store_fixture._should_scan_for_unhealthy_rollouts()  # pyright: ignore[reportPrivateUsage]
+    assert await store_fixture._should_scan_for_unhealthy_rollouts()  # pyright: ignore[reportPrivateUsage]