resolve copilot comments

.github/workflows/examples-calc-x.yml

@@ -171,12 +171,12 @@ jobs:
       - name: Sync dependencies (latest)
         run: |
           uv sync --frozen --no-default-groups --extra verl \
-            --group dev --group experiment --group agents --extra weave --group torch-gpu-stable
+            --group dev --group experiment --group agents --extra weave --extra mongo --group torch-gpu-stable
         if: matrix.setup-script == 'latest'
       - name: Sync dependencies (stable & legacy)
         run: |
           uv sync --frozen --no-default-groups --extra verl \
-            --group dev --group experiment --group agents --extra weave --group torch-gpu-${{ matrix.setup-script }}
+            --group dev --group experiment --group agents --extra weave --extra mongo --group torch-gpu-${{ matrix.setup-script }}
         if: matrix.setup-script != 'latest'
       - name: Freeze dependencies
         run: |
@@ -270,6 +270,33 @@ jobs:
           WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
           WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
 
+      - name: Setup Docker environments
+        run: ./scripts/mongodb_docker_run.sh
+        shell: bash
+
+      - name: Training with MongoDB
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/calc_x
+          ../../scripts/restart_ray.sh
+          sleep 5
+          PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --mongo-uri mongodb://localhost:27017/?replicaSet=rs0
+          sleep 10
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: calc_x_train_mongo
+
+      - name: Validate training with MongoDB
+        run: |
+          set -ex
+          uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_mongo.outputs.project_name }} ${{ steps.calc_x_train_mongo.outputs.run_name }}
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+
       - name: Training with LoRA
         run: |
           set -ex

.github/workflows/tests-full.yml

@@ -132,49 +132,7 @@ jobs:
         run: cd dashboard && npm run build
 
       - name: Setup Docker environments
-        run: |
-          set -euo pipefail
-
-          cd docker
-
-          # Setup data directories
-          ./setup.sh
-
-          # Start Dockers
-          docker compose -f compose.mongo.yml up -d
-
-          SERVICE_NAME=mongo
-          TIMEOUT=60   # seconds
-          SLEEP=2
-
-          cid="$(docker compose -f compose.mongo.yml ps -q "$SERVICE_NAME")"
-          if [ -z "$cid" ]; then
-            echo "Service $SERVICE_NAME is not running"
-            exit 1
-          fi
-
-          echo "Waiting for $SERVICE_NAME to become healthy..."
-          end=$((SECONDS + TIMEOUT))
-
-          while [ "$SECONDS" -lt "$end" ]; do
-            status="$(docker inspect -f '{{.State.Health.Status}}' "$cid")"
-            echo "Current status: $status"
-
-            if [ "$status" = "healthy" ]; then
-              echo "$SERVICE_NAME is healthy ✅"
-              exit 0
-            elif [ "$status" = "unhealthy" ]; then
-              echo "$SERVICE_NAME is unhealthy ❌"
-              docker logs "$cid" || true
-              exit 1
-            fi
-
-            sleep "$SLEEP"
-          done
-
-          echo "Timed out waiting for $SERVICE_NAME to become healthy after ${TIMEOUT}s"
-          docker logs "$cid" || true
-          exit 1
+        run: ./scripts/mongodb_docker_run.sh
         shell: bash
 
       - name: Launch LiteLLM Proxy

agentlightning/algorithm/verl/interface.py

@@ -32,6 +32,26 @@ class VERL(Algorithm):
         trainer_cls: Optional override for the trainer class. Experimental.
         daemon_cls: Optional override for the daemon class. Experimental.
 
+    !!! note "Trajectory aggregation (experimental)"
+
+        Trajectory-level aggregation merges an entire multi-turn rollout into a single,
+        masked training sample so GPU time is spent once per trajectory rather than N times
+        per turn. Enable it via:
+
+        ```python
+        config["agentlightning"]["trace_aggregator"] = {
+            "level": "trajectory",
+            "trajectory_max_prompt_length": ...,
+            "trajectory_max_response_length": ...,
+        }
+        ```
+
+        Keep conversations structured (message lists rather than manual string
+        concatenation) so prefix matching can stitch traces, and toggle `debug=True` plus
+        `unmatch_log_dir` when you need to inspect retokenization or chat-template
+        mismatches. See [this blog post](https://agent-lightning.github.io/posts/trajectory_level_aggregation/)
+        for more details.
+
     Examples:
         ```python
         from agentlightning.algorithm.verl import VERL

agentlightning/store/collection/mongo.py

@@ -1077,7 +1077,7 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
 class MongoLightningCollections(LightningCollections):
     """Mongo implementation of LightningCollections using MongoDB collections.
 
-    Serves as the storage base for [`MongoLightningStore`][agentlightning.store.MongoLightningStore].
+    Serves as the storage base for [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore].
     """
 
     def __init__(

agentlightning/store/mongo.py

@@ -33,7 +33,7 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
     Args:
         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``.
+        database_name: 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.

docs/deep-dive/store.md

@@ -1,8 +1,8 @@
 # Understanding Store
 
-The **[`LightningStore`][agentlightning.LightningStore]** is the central coordination point for Agent-lightning. It holds the task queue, rollouts, attempts, spans, and versioned resources, and exposes a small API both Runners and Algorithms use to communicate. This document explains what’s in the store, how statuses transition, how spans are recorded, and the concurrency model (threads & processes).
+The **[`LightningStore`][agentlightning.LightningStore]** is the central coordination point for Agent-lightning. It holds the task queue, rollouts, attempts, spans, and versioned resources, and exposes a small API both Runners and Algorithms use to communicate. This document explains what's in the store, how statuses transition, how spans are recorded, and the concurrency model (threads & processes).
 
-## What’s in the Store?
+## What's in the Store?
 
 ![Store Architecture](../assets/store-api-visualized.svg){ .center }
 
@@ -13,12 +13,11 @@ At a high level:
 * **Attempts** – Each rollout can have multiple executions (retries). Attempts track [`status`][agentlightning.Attempt.status], [`start_time`][agentlightning.Attempt.start_time], [`end_time`][agentlightning.Attempt.end_time], [`last_heartbeat_time`][agentlightning.Attempt.last_heartbeat_time] and link to spans. Valid [AttemptStatus][agentlightning.AttemptStatus] are `preparing`, `running`, `succeeded`, `failed`, `requeuing`, `cancelled`.
 * **Spans** – Structured trace events produced by the Tracer during an attempt. Spans are ordered by a **monotonic sequence id** per `(rollout_id, attempt_id)`.
 * **Resources** – Versioned, named bundles (e.g., prompt templates) referenced by rollouts.
+* **Workers** – Metadata about runner instances: heartbeat timestamps, current assignment, and status.
 
-Rollout and Task share the same surface in practice: [`Rollout.input`][agentlightning.types.Rollout] is the task input. The queue stores rollouts that are not yet running; [Runners][agentlightning.Runner] dequeue them and update the same rollout’s status as work progresses.
+Rollout and Task share the same surface in practice: [`Rollout.input`][agentlightning.types.Rollout] is the task input. The queue stores rollouts that are not yet running; [Runners][agentlightning.Runner] dequeue them and update the same rollout's status as work progresses.
 
-All [`LightningStore`][agentlightning.LightningStore] implementations must inherit from [`LightningStore`][agentlightning.LightningStore] and override the methods to implement the storage logic.
-
-Before we look at status transitions, it helps to keep in mind that rollouts are the “outside view,” while attempts are the “inside view.” Attempts are what actually run; rollouts summarize the latest attempt plus a small set of control actions like queueing and cancellation.
+Before we look at status transitions, it helps to keep in mind that rollouts are the "outside view," while attempts are the "inside view." Attempts are what actually run; rollouts summarize the latest attempt plus a small set of control actions like queueing and cancellation.
 
 ## Attempt Status Transitions
 
@@ -152,31 +151,220 @@ Programmatically this is encapsulated by [`Span.from_opentelemetry(readable_span
 
     [`add_span`][agentlightning.LightningStore.add_span] or [`add_otel_span`][agentlightning.LightningStore.add_otel_span] both appends a span *and* acts as a heartbeat that can revive `unresponsive` → `running`.
 
-## OTLP Compatibility
+### OTLP Compatibility
 
-Some of the LightningStore implementations support exporting traces via the [OTLP/HTTP specification](https://opentelemetry.io/docs/specs/otlp/). For example, [`LightningStoreServer`][agentlightning.LightningStoreServer] exposes `/v1/traces` endpoint, it implements the binary Protobuf variant defined by the spec, including the required `Content-Type: application/x-protobuf`, optional `Content-Encoding: gzip`, and status responses encoded as `google.rpc.Status`. Agent-lightning helps parsing `ExportTraceServiceRequest` messages, validate identifiers, normalize resource metadata, and allocate sequence
-numbers so store implementations only need to persist [`Span`][agentlightning.Span] objects in order.
+Some of the LightningStore implementations support exporting traces via the [OTLP/HTTP specification](https://opentelemetry.io/docs/specs/otlp/). For example, [`LightningStoreServer`][agentlightning.LightningStoreServer] exposes `/v1/traces` endpoint, it implements the binary Protobuf variant defined by the spec, including the required `Content-Type: application/x-protobuf`, optional `Content-Encoding: gzip`, and status responses encoded as `google.rpc.Status`. Agent-lightning helps parsing `ExportTraceServiceRequest` messages, validate identifiers, normalize resource metadata, and allocate sequence numbers so store implementations only need to persist [`Span`][agentlightning.Span] objects in order.
 
 Because the interface speaks standard OTLP, any OpenTelemetry-compatible SDK or collector can emit spans directly to a LightningStore OTLP endpoint without custom shims. The server responds according to the OTLP contract (status code, encoding, and error payloads), which keeps Agent-lightning interoperable with existing observability tooling. This compatibility serves as a strong complement to the OpenTelemetry conversion discussed above.
 
-## Store Implementations
+Check whether the store supports OTLP traces via the [`capabilities["otlp_traces"]`][agentlightning.LightningStore.capabilities] property.
 
-Currently, the only out-of-the-box implementation is [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore]:
+## Implementation Overview
 
-- Fast startup, zero external dependencies, and ideal for local development, CI, and unit tests.
-- Fully asyncio-safe for writes; most reader operations can iterate without locks, except those that need to perform multiple queries.
-- Includes a best-effort span eviction policy once memory crosses a configured watermark; querying evicted spans raises a clear error so callers can fall back.
+The `agentlightning.store` module is organized into two distinct layers plus optional wrappers:
 
-For production you will likely want persistence. We’re actively building a SQLite-backed store that keeps the same API surface while adding durability, crash recovery, and better historical span queries. If you need something sooner, implement your own store by subclassing [`LightningStore`][agentlightning.LightningStore] and providing concrete storage for the small set of abstract methods (`enqueue_rollout`, `dequeue_rollout`, `update_attempt`, `add_span`, etc.). This document plus the tests in `tests/store/` illustrate the expected behavior.
+```mermaid
+classDiagram
+    direction TB
 
-Different store implementations may have different capabilities. For example, [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore] does not support exporting traces via OTLP. Try to distinguish the capabilities of a store implementation by checking the [`capabilities`][agentlightning.LightningStore.capabilities] property.
+    class LightningStore {
+        <<abstract>>
+        +enqueue_rollout()
+        +dequeue_rollout()
+        +update_attempt()
+        +add_span()
+        +query_rollouts()
+        ...
+    }
+
+    class LightningCollections {
+        <<abstract>>
+        +rollouts: Collection
+        +attempts: Collection
+        +spans: Collection
+        +resources: Collection
+        +workers: Collection
+        +rollout_queue: Queue
+        +span_sequence_ids: KeyValue
+        +atomic()
+    }
+
+    class CollectionBasedLightningStore~T~ {
+        +collections: T
+        -healthcheck_before()
+        -tracked()
+    }
+
+    class InMemoryLightningStore
+    class MongoLightningStore
+    class InMemoryLightningCollections
+    class MongoLightningCollections
+
+    class LightningStoreServer {
+        +store: LightningStore
+        +start()
+        +stop()
+    }
+    class LightningStoreClient {
+        +server_address: str
+    }
+    class LightningStoreThreaded {
+        +store: LightningStore
+    }
+
+    LightningStore <|-- CollectionBasedLightningStore
+    LightningStore <|-- LightningStoreServer
+    LightningStore <|-- LightningStoreClient
+    LightningStore <|-- LightningStoreThreaded
+
+    CollectionBasedLightningStore <|-- InMemoryLightningStore
+    CollectionBasedLightningStore <|-- MongoLightningStore
+
+    LightningCollections <|-- InMemoryLightningCollections
+    LightningCollections <|-- MongoLightningCollections
+
+    InMemoryLightningStore ..> InMemoryLightningCollections : uses
+    MongoLightningStore ..> MongoLightningCollections : uses
+
+    LightningStoreServer o-- LightningStore : wraps
+    LightningStoreThreaded o-- LightningStore : wraps
+```
+
+1. **Collections Layer** – Low-level storage primitives ([`LightningCollections`][agentlightning.store.collection.LightningCollections]) providing CRUD operations via [`Collection`][agentlightning.store.collection.Collection], [`Queue`][agentlightning.store.collection.Queue], and [`KeyValue`][agentlightning.store.collection.KeyValue] interfaces. Each backend (in-memory, MongoDB) implements these primitives.
+
+2. **Store Layer** – All [`LightningStore`][agentlightning.LightningStore] implementations must inherit from [`LightningStore`][agentlightning.LightningStore] and override the methods to implement the storage logic. [`CollectionBasedLightningStore`][agentlightning.CollectionBasedLightningStore] builds on collections to implement the full [`LightningStore`][agentlightning.LightningStore] API, including business logic like status transitions, watchdog health checks, and retry policies.
+
+3. **Wrappers** – Cross-cutting concerns live in thin wrappers:
+    - [`LightningStoreThreaded`][agentlightning.LightningStoreThreaded] adds mutex-based thread safety.
+    - [`LightningStoreServer`][agentlightning.LightningStoreServer] / [`LightningStoreClient`][agentlightning.LightningStoreClient] enable multi-process access over HTTP.
+
+## Collections
+
+The collections layer provides storage primitives that [`CollectionBasedLightningStore`][agentlightning.CollectionBasedLightningStore] builds upon. This separation keeps business logic (status transitions, watchdog, retries) in the store layer while allowing different backends to focus purely on persistence.
+
+The off-the-shelf implementations are [`InMemoryLightningCollections`][agentlightning.store.collection.InMemoryLightningCollections] and [`MongoLightningCollections`][agentlightning.store.collection.mongo.MongoLightningCollections], which are the underlying collections for [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore] and [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore], respectively.
+
+### Collection Primitives
+
+[`LightningCollections`][agentlightning.store.collection.LightningCollections] bundles three primitive types:
+
+| Primitive | Purpose | Methods |
+|-----------|---------|---------|
+| [`Collection[T]`][agentlightning.store.collection.Collection] | Indexed storage with primary keys | [`query()`][agentlightning.store.collection.Collection.query], [`get()`][agentlightning.store.collection.Collection.get], [`insert()`][agentlightning.store.collection.Collection.insert], [`update()`][agentlightning.store.collection.Collection.update], [`upsert()`][agentlightning.store.collection.Collection.upsert], [`delete()`][agentlightning.store.collection.Collection.delete] |
+| [`Queue[T]`][agentlightning.store.collection.Queue] | FIFO queue for task scheduling | [`enqueue()`][agentlightning.store.collection.Queue.enqueue], [`dequeue()`][agentlightning.store.collection.Queue.dequeue], [`peek()`][agentlightning.store.collection.Queue.peek], [`size()`][agentlightning.store.collection.Queue.size] |
+| [`KeyValue[K, V]`][agentlightning.store.collection.KeyValue] | Simple key-value store | [`get()`][agentlightning.store.collection.KeyValue.get], [`set()`][agentlightning.store.collection.KeyValue.set], [`inc()`][agentlightning.store.collection.KeyValue.inc], [`chmax()`][agentlightning.store.collection.KeyValue.chmax], [`pop()`][agentlightning.store.collection.KeyValue.pop] |
+
+Every [`LightningCollections`][agentlightning.store.collection.LightningCollections] instance exposes these named collections:
+
+- `rollouts` – [`Collection[Rollout]`][agentlightning.store.collection.Collection] keyed by `rollout_id`
+- `attempts` – [`Collection[Attempt]`][agentlightning.store.collection.Collection] keyed by `(rollout_id, attempt_id)`
+- `spans` – [`Collection[Span]`][agentlightning.store.collection.Collection] keyed by `(rollout_id, attempt_id, span_id)`
+- `resources` – [`Collection[ResourcesUpdate]`][agentlightning.store.collection.Collection] keyed by `resources_id`
+- `workers` – [`Collection[Worker]`][agentlightning.store.collection.Collection] keyed by `worker_id`
+- `rollout_queue` – [`Queue[str]`][agentlightning.store.collection.Queue] holding rollout IDs awaiting execution
+- `span_sequence_ids` – [`KeyValue[str, int]`][agentlightning.store.collection.KeyValue] tracking monotonic sequence counters
+
+### Atomic Operations
+
+Collections support atomic operations through the [`atomic()`][agentlightning.store.collection.LightningCollections.atomic] context manager:
+
+```python
+async with collections.atomic(mode="rw", labels=["rollouts", "attempts"]) as ctx:
+    rollout = await ctx.rollouts.get(filter={"rollout_id": {"exact": rollout_id}})
+    # modify and update within the same transaction
+    await ctx.rollouts.update([updated_rollout])
+```
+
+The arguments passed to [`atomic()`][agentlightning.store.collection.LightningCollections.atomic] are quite arbitrary and flexible. Different implementations may have different interpretations of the arguments. For example, to [`InMemoryLightningCollections`][agentlightning.store.collection.InMemoryLightningCollections], the `mode` parameter controls locking behavior (`"r"` for read-only, `"rw"` for read-write), while `labels` specifies which collections to lock. Acquiring locks in sorted order prevents deadlocks when multiple operations run concurrently.
+
+### Implementing a Custom Backend
+
+To add a new storage backend, implement [`LightningCollections`][agentlightning.store.collection.LightningCollections]:
+
+```python
+from agentlightning.store.collection import LightningCollections, Collection, Queue, KeyValue
+
+class MyLightningCollections(LightningCollections):
+    @property
+    def rollouts(self) -> Collection[Rollout]:
+        return self._rollouts  # your implementation
+
+    @property
+    def rollout_queue(self) -> Queue[str]:
+        return self._queue  # your implementation
+
+    # ... implement remaining properties
+
+    async def atomic(self, *, mode, snapshot=False, labels=None, **kwargs):
+        # provide transaction / locking semantics
+        ...
+```
+
+Then instantiate your store:
+
+```python
+from agentlightning.store.collection_based import CollectionBasedLightningStore
+
+store = CollectionBasedLightningStore(collections=MyLightningCollections())
+```
+
+The store layer handles all business logic; your collections just need to provide correct CRUD semantics.
+
+## Collection-based Store Implementations
+
+Agent-lightning ships with two collection-based store implementations:
+
+### InMemoryLightningStore
+
+[`InMemoryLightningStore`][agentlightning.InMemoryLightningStore] uses [`InMemoryLightningCollections`][agentlightning.store.collection.InMemoryLightningCollections] backed by Python data structures. It supports **fast startup** with zero external dependencies—ideal for local development, CI, and unit tests. It also provides two lock modes, configurable between `"asyncio"` (single-thread, multiple coroutines) and `"thread"` (multi-threaded via [aiologic](https://github.com/x42005e1f/aiologic)).
+[`InMemoryLightningCollections`][agentlightning.store.collection.InMemoryLightningCollections] use nested dictionaries for O(1) primary-key lookup and `deque` for the task queue.
+
+### MongoLightningStore
+
+[`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore] uses [`MongoLightningCollections`][agentlightning.store.collection.mongo.MongoLightningCollections] backed by MongoDB. It supports **persistent storage** suitable for production deployments and **multi-process safe** via database-level atomicity. It also supports **partition support** via `partition_id` for running multiple trainers against the same database.
+
+```python
+from agentlightning.store.mongo import MongoLightningStore
+
+store = MongoLightningStore(
+    mongo_uri="mongodb://localhost:27017/?replicaSet=rs0",
+    database_name="agentlightning",
+    partition_id="trainer-1",  # optional: isolate data per trainer
+)
+```
+
+!!! note
+
+    [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore] requires the `mongo` optional dependency. Install with `pip install agentlightning[mongo]`.
+
+### Capabilities
+
+[](){ #store-capabilities }
+
+Different stores have different capabilities. Check the [`capabilities`][agentlightning.LightningStore.capabilities] property to understand what a store supports:
+
+| Capability | Description | InMemory | Mongo | Server | Client |
+|------------|-------------|----------|-------|--------|--------|
+| `thread_safe` | Safe for concurrent access from multiple threads | configurable | ✓ | ✓ | ✓ |
+| `async_safe` | Safe for concurrent access from multiple coroutines | ✓ | ✓ | ✓ | ✓ |
+| `zero_copy` | Can be shared across processes without serialization | ✗ | ✓ | ✓ | ✓ |
+| `otlp_traces` | Exposes an OTLP-compatible `/v1/traces` endpoint | ✗ | ✗ | ✓ | ✓ |
 
 ## Thread Safety
 
-**[`LightningStoreThreaded`][agentlightning.LightningStoreThreaded]** is a subclass of [`LightningStore`][agentlightning.LightningStore] that wraps another underlying store to make a store instance safe for multi-threaded callers. It wraps every state-mutating call in a mutex. Specifically:
+Thread safety can be achieved at different layers:
+
+**At the collections layer**: [`InMemoryLightningCollections`][agentlightning.store.collection.InMemoryLightningCollections] accepts a `lock_type` parameter:
+
+- `"asyncio"` – Uses per-event-loop `asyncio.Lock` for single-threaded, multi-coroutine scenarios.
+- `"thread"` – Uses `aiologic.Lock` for true multi-threaded access.
+
+**At the store layer**: [`LightningStoreThreaded`][agentlightning.LightningStoreThreaded] wraps any [`LightningStore`][agentlightning.LightningStore] to add mutex-based thread safety:
 
 * Methods like [`start_rollout`][agentlightning.LightningStore.start_rollout], [`enqueue_rollout`][agentlightning.LightningStore.enqueue_rollout], [`update_attempt`][agentlightning.LightningStore.update_attempt], [`add_span`][agentlightning.LightningStore.add_span], etc. are guarded by a lock.
-* Non-mutating, potentially blocking calls remain pass-through by design (e.g., [`wait_for_rollouts`][agentlightning.LightningStore.wait_for_rollouts]), as they don’t modify shared state and should not hold the lock for long periods.
+* Non-mutating, potentially blocking calls remain pass-through by design (e.g., [`wait_for_rollouts`][agentlightning.LightningStore.wait_for_rollouts]), as they don't modify shared state and should not hold the lock for long periods.
+
+Database-based stores like [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore] are inherently thread-safe through database atomicity guarantees.
 
 ## Process Safety and Client-server Store
 
@@ -188,7 +376,7 @@ Different store implementations may have different capabilities. For example, [`
 
 The server tracks the creator PID. In the owner process it delegates directly to the in-memory store; in other processes it lazily constructs a [`LightningStoreClient`][agentlightning.LightningStoreClient] to talk to the HTTP API. This prevents accidental cross-process mutation of the wrong memory image. When the server is pickled (e.g., via `multiprocessing`), only the minimal fields are serialized, but **NOT** the FastAPI/uvicorn objects. Subprocesses won’t accidentally carry live server state. Forked subprocess should also use [`LightningStoreClient`][agentlightning.LightningStoreClient] to communicate with the server in the main process.
 
-On the client side, the client retries network/5xx failures using a small backoff, and probes `/health` between attempts. Application exceptions inside the server are wrapped as HTTP 400 with a traceback—these are **not retried**. The client also maintains a **per-event-loop** `aiohttp.ClientSession` map so that tracer callbacks (often on separate loops/threads) don’t hang by reusing a session from another loop.
+On the client side, the client retries network/5xx failures using a small backoff, and probes `/v1/agl/health` between attempts. Application exceptions inside the server are wrapped as HTTP 400 with a traceback—these are **not retried**. The client also maintains a **per-event-loop** `aiohttp.ClientSession` map so that tracer callbacks (often on separate loops/threads) don’t hang by reusing a session from another loop.
 
 Minimal lifecycle:
 

docs/reference/cli.md

@@ -1,7 +1,5 @@
 # Command Line Interface
 
-<!-- TODO: This document should be auto-generated. -->
-
 !!! warning
 
     This document is a work in progress and might not be updated with the latest changes.
@@ -65,15 +63,40 @@ Agent-lightning's LightningStore CLI. Use it to start an independent LightningSt
 Currently the store data are stored in memory and will be lost when the server is stopped.
 
 ```text
-usage: agl store [-h] [--port PORT]
+usage: agl store [-h] [--host HOST] [--port PORT] [--cors-origin CORS_ORIGINS] [--log-level {DEBUG,INFO,WARNING,ERROR}] [--tracker {prometheus,console} [{prometheus,console} ...]] [--n-workers N_WORKERS] [--backend {memory,mongo}]
+                 [--mongo-uri MONGO_URI]
 
 Run a LightningStore server
 
 options:
-  -h, --help   show this help message and exit
-  --port PORT  Port to run the server on
+  -h, --help            show this help message and exit
+  --host HOST           Host to bind the server to
+  --port PORT           Port to run the server on
+  --cors-origin CORS_ORIGINS
+                        Allowed CORS origin. Repeat for multiple origins. Use '*' to allow all origins.
+  --log-level {DEBUG,INFO,WARNING,ERROR}
+                        Configure the logging level for the store.
+  --tracker {prometheus,console} [{prometheus,console} ...]
+                        Enable metrics tracking. Repeat for multiple trackers.
+  --n-workers N_WORKERS
+                        Number of workers to run in the server. When it's greater than 1, the server will be run using `mp` launch mode. Only applicable for zero-copy stores such as MongoDB backend.
+  --backend {memory,mongo}
+                        Backend to use for the store.
+  --mongo-uri MONGO_URI
+                        MongoDB URI to use for the store. Applicable only if --backend is 'mongo'.
 ```
 
+!!! tip
+
+    After launching the store via CLI, you can tell the [`Trainer`][agentlightning.Trainer] to use the store by passing the store address to the trainer.
+
+    ```python
+    store_client = agl.LightningStoreClient("http://localhost:4747")
+    trainer = agl.Trainer(store=store_client, ...)
+    ```
+
+    See [using external store][debug-with-external-store] for more details.
+
 ## 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.
@@ -93,18 +116,3 @@ options:
                         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.
-
-```text
-usage: agl agentops [-h] [--daemon] [--port PORT]
-
-Start AgentOps server
-
-options:
-  -h, --help   show this help message and exit
-  --daemon     Run server as a daemon
-  --port PORT  Port to run the server on
-```

docs/reference/store.md

@@ -8,6 +8,8 @@
 
 ::: agentlightning.InMemoryLightningStore
 
+::: agentlightning.store.mongo.MongoLightningStore
+
 ::: agentlightning.CollectionBasedLightningStore
 
 ## Client-Server and Thread-safe Wrappers
@@ -39,3 +41,13 @@
 ::: agentlightning.store.collection.DictBasedKeyValue
 
 ::: agentlightning.store.collection.InMemoryLightningCollections
+
+::: agentlightning.store.collection.mongo.MongoBasedCollection
+
+::: agentlightning.store.collection.mongo.MongoBasedQueue
+
+::: agentlightning.store.collection.mongo.MongoBasedKeyValue
+
+::: agentlightning.store.collection.mongo.MongoClientPool
+
+::: agentlightning.store.collection.mongo.MongoLightningCollections

docs/tutorials/emitter.md

@@ -1,4 +1,4 @@
-# Using the Emitter
+# Using Emitters
 
 [](){ #using-emitter }
 

docs/tutorials/parallelize.md

@@ -1,4 +1,4 @@
-# Scaling out Algorithms and Rollouts
+# Scaling out Agent-lightning
 
 Agent-lightning splits training into an **algorithm bundle** and a **runner bundle** that exchange work through the [`LightningStore`][agentlightning.LightningStore]. This tutorial shows how to increase rollout throughput, place bundles across processes or machines, and keep the algorithm side scalable with external frameworks.
 
@@ -226,3 +226,109 @@ Agent-lightning strives to make algorithms’ own parallelization work well unde
 !!! note
 
     The [birds' eye view][birds-eye-view-client-server-strategy] illustrates how adapters, proxies, and stores interact when the algorithm spawns additional workers. Use that diagram as a checklist when introducing new distributed components.
+
+## Parallelizing [`LightningStore`][agentlightning.LightningStore]
+
+By default, Agent-lightning persists rollouts and spans in an in-memory store. [`Trainer.fit`][agentlightning.Trainer.fit] spins it up automatically, or you can launch it yourself via the [`agl store` command](../reference/cli.md). [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore] keeps all state inside the current process, which makes local iteration fast but introduces two production constraints:
+
+1. Spans are evicted once the process crosses its memory cap, so long runs risk data loss unless the host has abundant RAM.
+2. Although the store is well optimized via asynchronous programming, the store lives in a single process and remains bound by the GIL, preventing it from saturating multi-core machines.
+
+!!! note "General note for all server-client stores"
+
+    If your algorithm and runners communicate through HTTP protocol (which should be the default for 99% of the cases), you need to ensure the file limit is sufficiently large to avoid the "Too many open files" error. You can set the file limit by running the following command:
+
+    ```bash
+    ulimit -n 100000
+    ```
+
+For resilient runs, switch to a persistent backend such as [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore], which writes data to MongoDB instead of local RAM. Agent-lightning relies on [pymongo](https://pymongo.readthedocs.io/en/stable/) to interact with MongoDB, which can be installed via:
+
+```bash
+pip install agentlightning[mongo]
+```
+
+To use the MongoDB store, you need to pass the MongoDB URI to the store constructor. The URI should be in the format of `mongodb://<host>:<port>/<database>?replicaSet=<replicaSet>`.
+
+```python
+from agentlightning.store.mongo import MongoLightningStore
+
+trainer = agl.Trainer(
+    algorithm=algorithm,
+    store=MongoLightningStore(mongo_uri="mongodb://localhost:27017/?replicaSet=rs0"),
+)
+```
+
+!!! tip "Setting up MongoDB"
+
+    MongoDB is a popular document-oriented database. Before running Agent-lightning with [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore], make sure that you've already had a MongoDB instance running. Setting up can be conveniently done via Docker Compose via [compose.mongo.yml]({{ src("docker/compose.mongo.yml") }}). Unless targeting serious production use, we recommend creating the data folders and setting them to `777` permission to avoid permission issues.
+
+    ```bash
+    mkdir -p data/mongo-host
+    chmod 777 data/mongo-host
+    docker compose -f compose.mongo.yml up -d
+    ```
+
+    Alternatively, you can also install MongoDB manually following the [official documentation](https://www.mongodb.com/docs/manual/installation/). If you installed MongoDB manually, an important note is that you need to ensure that the MongoDB instance has enabled replica set feature, since Agent-lightning uses the transactional operations internally. The simplest approach is to use the following script (executed in the MongoDB shell) to initialize the replica set:
+
+    ```javascript
+    rs.initiate({
+      _id: "rs0",
+      members: [{ _id: 0, host: "localhost:27017" }],
+    });
+    ```
+
+To scale out further, launch the store server via [`agl store --backend mongo`](../reference/cli.md) (see [Debugging with External Store][debug-with-external-store]). The CLI accepts `--n-workers`, which starts the server under `gunicorn` with multiple worker processes so concurrent runners can push and pull at higher throughput. This option applies only to persistent backends; an in-memory store, on the other hand, cannot be sharded across workers because its state lives inside one process.
+
+!!! note
+
+    The `--n-workers` here is the number of worker processes for the store server, NOT related to the number of rollout runners.
+
+## Increasing Throughput of LLM Proxy
+
+Agent-lightning includes an optional [`LLMProxy`][agentlightning.LLMProxy] that wraps [LiteLLM](https://docs.litellm.ai/) to provide a unified OpenAI-compatible endpoint for your agents. When rollout throughput increases, the proxy can become a bottleneck. You can scale it out using the same pattern as the store server.
+
+To increase proxy throughput, pass `num_workers` when constructing the proxy:
+
+```python
+import agentlightning as agl
+
+proxy = agl.LLMProxy(
+    port=4000,
+    launch_mode="mp",  # multiprocessing mode
+    num_workers=4,  # four gunicorn workers handle concurrent requests
+)
+```
+
+You can also configure the proxy through [`Trainer`][agentlightning.Trainer]:
+
+```python
+trainer = agl.Trainer(
+    algorithm=algorithm,
+    n_runners=8,  # The runners here is the rollout runners, not related to LLM proxy replicas
+    llm_proxy={"port": 4000, "num_workers": 4},  # launch mode is actually mp by default
+)
+```
+
+When `num_workers > 1`, the launcher starts gunicorn with the specified number of worker processes. Each worker runs its own event loop, allowing the proxy to handle many concurrent LLM requests without being blocked by Python's GIL.
+
+!!! tip
+
+    When using `mp` launch mode, [`LLMProxy`][agentlightning.LLMProxy] will start the server in a separate process. To make sure the proxy is still accessing the same store as the main process, you need to set the store to be [zero-copy compatible][store-capabilities], which means, either the store is a native zero-copy store like [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore] or the store is wrapped via [`LightningStoreServer`][agentlightning.LightningStoreServer] or [`LightningStoreClient`][agentlightning.LightningStoreClient].
+
+!!! note "Shared Server Infrastructure"
+
+    Both [`LightningStoreServer`][agentlightning.LightningStoreServer] and [`LLMProxy`][agentlightning.LLMProxy] rely on a common utility called [`PythonServerLauncherArgs`][agentlightning.utils.server_launcher.PythonServerLauncherArgs]. This dataclass captures the settings needed to launch a FastAPI application:
+
+    ```python
+    from agentlightning.utils import PythonServerLauncherArgs
+
+    args = PythonServerLauncherArgs(
+        port=8000,
+        host="0.0.0.0",
+        n_workers=4,          # spawn 4 gunicorn workers
+        launch_mode="thread", # or "mp" for multiprocessing, "asyncio" for in-loop
+    )
+    ```
+
+    Under the hood, [`PythonServerLauncher`][agentlightning.utils.server_launcher.PythonServerLauncher] reads these arguments and chooses between uvicorn (single worker) and gunicorn (multiple workers) automatically.

examples/calc_x/train_calc_agent.py

@@ -122,6 +122,7 @@ def train(
     lora_adapter_path: Optional[str],
     trajectory_level: bool = False,
     weave: bool,
+    mongo_uri: Optional[str],
 ):
     """The training entrypoint function for Calc-X agent with VERL algorithm.
 
@@ -139,6 +140,7 @@ def train(
         lora_adapter_path: Optional path to a pre-trained LoRA adapter to load.
         trajectory_level: Whether to enable trajectory level in trace aggregator.
         weave: Whether to enable Weave tracing.
+        mongo_uri: MongoDB URI to use for the store.
     """
     # Load datasets (respect CLI file paths)
     train_dataset = cast(agl.Dataset[MathProblem], HuggingFaceDataset.from_parquet(train_file).to_list())  # type: ignore
@@ -216,6 +218,10 @@ def train(
 
     if external_store_address:
         store: Optional[agl.LightningStore] = agl.LightningStoreClient(external_store_address)
+    elif mongo_uri:
+        from agentlightning.store.mongo import MongoLightningStore
+
+        store = MongoLightningStore(mongo_uri=mongo_uri)
     else:
         store = None
 
@@ -278,6 +284,12 @@ def main():
         action="store_true",
         help="Enable trajectory level in trace aggregator.",
     )
+    parser.add_argument(
+        "--mongo-uri",
+        type=str,
+        default=None,
+        help="MongoDB URI to use for the store.",
+    )
 
     args = parser.parse_args()
 
@@ -308,6 +320,7 @@ def main():
         lora_adapter_path=args.lora_adapter_path,
         trajectory_level=args.trajectory_level,
         weave=args.weave,
+        mongo_uri=args.mongo_uri,
     )
 
 

scripts/mongo/README.md

@@ -1,9 +0,0 @@
-# MongoDB Development Setup
-
-This script is used to setup MongoDB for development.
-
-## Usage
-
-```bash
-docker compose up -d
-```

scripts/mongo/docker-compose.yml

@@ -1,10 +0,0 @@
-services:
-  mongo:
-    image: mongo:latest
-    container_name: mongo-dev
-    ports:
-      - "27017:27017"
-    command: ["mongod", "--bind_ip_all", "--replSet", "rs0"]
-    volumes:
-      - ./data:/data/db
-      - ./init-rs.js:/docker-entrypoint-initdb.d/init-rs.js:ro

scripts/mongo/init-rs.js

@@ -1,6 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-rs.initiate({
-  _id: "rs0",
-  members: [{ _id: 0, host: "localhost:27017" }],
-});

scripts/mongodb_docker_run.sh

@@ -0,0 +1,44 @@
+#!/bin/bash
+
+set -euo pipefail
+
+cd docker
+
+# Setup data directories
+./setup.sh
+
+# Start Dockers
+docker compose -f compose.mongo.yml up -d
+
+SERVICE_NAME=mongo
+TIMEOUT=60   # seconds
+SLEEP=2
+
+cid="$(docker compose -f compose.mongo.yml ps -q "$SERVICE_NAME")"
+if [ -z "$cid" ]; then
+    echo "Service $SERVICE_NAME is not running"
+    exit 1
+fi
+
+echo "Waiting for $SERVICE_NAME to become healthy..."
+end=$((SECONDS + TIMEOUT))
+
+while [ "$SECONDS" -lt "$end" ]; do
+    status="$(docker inspect -f '{{.State.Health.Status}}' "$cid")"
+    echo "Current status: $status"
+
+    if [ "$status" = "healthy" ]; then
+        echo "$SERVICE_NAME is healthy ✅"
+        exit 0
+    elif [ "$status" = "unhealthy" ]; then
+        echo "$SERVICE_NAME is unhealthy ❌"
+        docker logs "$cid" || true
+        exit 1
+    fi
+
+    sleep "$SLEEP"
+done
+
+echo "Timed out waiting for $SERVICE_NAME to become healthy after ${TIMEOUT}s"
+docker logs "$cid" || true
+exit 1