chore: import upstream snapshot with attribution
Deploy Documentation / deploy (push) Has been cancelled
CPU Test / Test (Utilities, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (LLM proxy, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Others, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Store, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Utilities, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Weave, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (AgentOps, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (LLM proxy, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Others, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Weave, latest, Python 3.13) (push) Has been cancelled
Dashboard / Chromatic (push) Has been cancelled
CPU Test / Lint - fast (push) Has been cancelled
CPU Test / Lint - next (push) Has been cancelled
CPU Test / Lint - slow (push) Has been cancelled
CPU Test / Lint - JavaScript (push) Has been cancelled
CPU Test / Build documentation (push) Has been cancelled
CPU Test / Test (AgentOps, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (LLM proxy, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (Others, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (Store, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (Weave, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (AgentOps, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Store, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Utilities, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Weave, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (AgentOps, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (LLM proxy, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (Others, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (Store, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (Utilities, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (JavaScript) (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 12:44:17 +08:00
commit 85742ab165
588 changed files with 320176 additions and 0 deletions
+539
View File
@@ -0,0 +1,539 @@
# The Bird's Eye View of Agent-lightning
!!! warning "High Volume of Information Ahead"
This article provides an in-depth exploration of the Agent-lightning architecture.
It is not intended as a beginners guide or usage tutorial.
This article summarizes how Agent-lightning (as of v0.2) wires the [Algorithm][agentlightning.Algorithm], [Runner][agentlightning.Runner], and [LightningStore][agentlightning.LightningStore] loop together and shows where auxiliary components (the [Tracer][agentlightning.Tracer], [Adapter][agentlightning.Adapter], and [LLM Proxy][agentlightning.LLMProxy]) plug into the loop. Each section provides a diagram for a different perspective of the system.
## Algorithm ↔ Runner ↔ Store data flow
At its heart, Agent-lightning is built on three main components that work in a coordinated loop:
* **[Algorithm][agentlightning.Algorithm]:** The "brain" of the system. It decides what tasks to run, learns from the results, and updates resources (like AI models or prompts).
* **[Runner][agentlightning.Runner]:** The "worker" of the system. It executes tasks assigned by the algorithm, runs the agent, and records the results.
* **[LightningStore][agentlightning.LightningStore]:** The central "database" and message queue. It acts as the single source of truth, storing tasks, results, and resources, and enabling communication between the Algorithm and Runner.
The typical data flow in a training loop is as follows: The **[Algorithm][agentlightning.Algorithm]** enqueues tasks (called **[Rollouts][agentlightning.Rollout]**) into the **[LightningStore][agentlightning.LightningStore]**. A **[Runner][agentlightning.Runner]** then dequeues a task, executes it, and streams the results (called **[Spans][agentlightning.Span]**) back to the store. Once the task is complete, the algorithm can query the new data from the store to learn and update its resources.
The diagram below shows this fundamental interaction in a simple, non-parallel setup.
```mermaid
sequenceDiagram
autonumber
participant Algo as Algorithm
participant Store as LightningStore
participant Runner
participant Agent
loop Over the dataset
Algo-->>Store: add_resources + enqueue_rollout
Store-->>Runner: dequeue_rollout → AttemptedRollout
Store-->>Runner: get_latest_resources
Runner-->>Store: update_attempt("running", worker_id)
Runner->>Agent: rollout + resources
Agent->>Runner: reward / spans
Runner-->>Store: add_span or add_otel_span
Runner-->>Store: update_attempt("finished", status)
Store-->>Algo: query_rollouts + spans
Algo-->>Algo: Update resources (optional)
end
```
Solid lines represent direct calls, while dashed lines are asynchronous or long-running operations.
### Key Terminology
We define the following terms, which may be helpful for understanding the diagram above.
* **[Resources][agentlightning.Resource]:** A collection of assets to be tuned or trained. Agents perform rollouts against resources and collect span data. Algorithms use those data to update the resources. In RL training, the resources are a tunable model. In prompt tuning, the resources are prompt templates.
* **[Rollout][agentlightning.Rollout]:** A unit of work that an agent performs against a resource. A rollout (noun) can be incomplete, in which case it is also known as a **task**, **sample**, or **job** (these terms are used interchangeably). The agent executes its own defined workflow against the rollout — the process is also called "to rollout" (verb). After execution, the rollout (noun) is considered *complete*.
* **[Attempt][agentlightning.Attempt]:** A single execution of a rollout. One rollout can have multiple attempts in case of failures or timeouts.
* **[Span][agentlightning.Span]:** During the rollout, the agent can generate multiple spans (also known as "traces" or "events"). The recorded spans are collected in the store, which is crucial for understanding agent behavior and optimizing agents.
* **[Reward][agentlightning.emit_reward]:** A special span that is defined as a number judging the quality of the rollout during some period of the rollout.
* **[Dataset][agentlightning.Dataset]:** A collection of incomplete rollouts (i.e., tasks) for the agent to process. The dual datasets (train, val) serve as the initial input for the algorithm to enqueue the first batch of rollouts.
### Store
As discussed previously, the [LightningStore][agentlightning.LightningStore] is the central hub for all data in Agent-lightning. The store exposes a set of APIs for algorithms and runners to interact with the data; the most important ones are:
```python
from agentlightning.types import AttemptedRollout, ResourcesUpdate, Span, TaskInput
class LightningStore:
async def enqueue_rollout(self, input: TaskInput, ...) -> Rollout: ...
async def dequeue_rollout(self) -> AttemptedRollout | None: ...
async def add_span(self, span: Span) -> Span: ...
async def get_latest_resources(self) -> Optional[ResourcesUpdate]: ...
async def wait_for_rollouts(self, rollout_ids: List[str], ...): ...
async def query_spans(self, rollout_id: str, ...): ...
async def update_attempt(self, rollout_id: str, attempt_id: str, status: str, ...): ...
...
```
These interfaces operate on [`AttemptedRollout`][agentlightning.AttemptedRollout], [`ResourcesUpdate`][agentlightning.ResourcesUpdate], [`Span`][agentlightning.Span], and [`TaskInput`][agentlightning.TaskInput] instances from `agentlightning.types`.
As the APIs show, the store essentially provides a queue for rollouts and storage for resources, spans, and attempts. Developers should implement the store carefully to ensure data integrity and consistency, especially when multiple runners work in parallel across multiple attempts.
The store is designed to be extensible. Users can implement their own store by inheriting from [`LightningStore`][agentlightning.LightningStore] and overriding methods. Agent-lightning provides a few reference implementations, such as [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore] (default) and `SqliteLightningStore` (under construction). When parallelized, the store may need special wrappers to ensure thread/process safety or delegate computation to a store in another process or machine.
## Supporting Components in the Loop
While the core loop is simple, Agent-lightning provides several components to make development easier and more powerful.
### Tracer
The [`Tracer`][agentlightning.Tracer] is a component within the [`Runner`][agentlightning.Runner] that records detailed spans (events) during an agent's execution and sends them to the [`LightningStore`][agentlightning.LightningStore]. Instead of requiring the agent to manually log every span, the tracer automatically instruments key methods (e.g., LLM calls) and captures their inputs, outputs, and metadata. This provides a detailed log of the agent's behavior with minimal effort.
```mermaid
sequenceDiagram
autonumber
participant Store
participant Runner
participant Tracer
participant Agent
Note over Runner,Tracer: Runner manages tracer as member
Tracer->>Agent: Apply instrumentation
loop Until no more rollouts
Store-->>Runner: dequeue_rollout → AttemptedRollout
Store-->>Runner: get_latest_resources
Runner->>Agent: training_rollout / validation_rollout
loop For each finished span
Agent-->>Tracer: openai.chat.completion invoked<br>agent.execute invoked<br>...
Agent->>Tracer: emit intermediate reward
Tracer-->>Store: add_otel_span(rollout_id, attempt_id, span)
end
Agent->>Runner: final reward + extra spans (if any)
Runner-->>Store: add_span(rollout_id, attempt_id, span)
Runner-->>Store: update_attempt(status)
end
Tracer->>Agent: Unapply instrumentation
```
The above diagram shows the overall data flow between store, tracer and agent. In realistic, it's a bit more complicated than that. Spans are not emitted actively by the agent; they are intercepted by the tracer by hooking and instrumenting key methods used in the agents. The tracer uses a callback (called exporter) to monitor events and log to the store. Before a rollout starts, the runner enters a [`trace_context`][agentlightning.Tracer.trace_context] before invoking the agent, wiring store identifiers into the tracer. Each span completion streams back to the store through `LightningSpanProcessor`, so the agents instrumentation lands in [`add_otel_span`][agentlightning.LightningStore.add_otel_span]. If the agents rollout method returns a numeric reward, the runner emits one more OpenTelemetry span before finalizing the attempt.
### Hooks
[`Hook`][agentlightning.Hook] implementations are user-defined callback functions that allow you to augment a [`Runner`][agentlightning.Runner]'s behavior at specific points in its lifecycle. You can use hooks to add custom logging, set up resources before a rollout begins, or tear them down after it ends. Hooks can be triggered at four key moments: `on_rollout_start`, `on_trace_start`, `on_trace_end`, and `on_rollout_end`.
Users should pay special attention to the difference between `on_trace_end` and `on_rollout_end`. The former is called right before the tracer exits the trace context, while the latter is called after the runner processes the final leftover rewards and spans, and finalizes the attempt in the store.
```mermaid
sequenceDiagram
autonumber
participant Store
participant Hooks
participant Runner
participant Tracer
participant Agent
Note over Runner,Hooks: Runner manages hooks as member
loop Until no more rollouts
Store-->>Runner: dequeue_rollout → AttemptedRollout
Store-->>Runner: get_latest_resources
Runner->>Hooks: on_rollout_start(agent, runner, rollout)
Runner->>Agent: training_rollout / validation_rollout
Tracer->>Agent: enter_trace_context
activate Tracer
Runner->>Hooks: on_trace_start(agent, runner, tracer, rollout)
Note over Runner,Agent: Agent rollout omitted
Runner->>Hooks: on_trace_end(agent, runner, tracer, rollout)
Tracer->>Agent: exit_trace_context
deactivate Tracer
Agent->>Runner: final reward + extra spans (if any)
Runner-->>Store: add_span(rollout_id, attempt_id, span)
Runner->>Hooks: on_rollout_end(agent, runner, rollout, status)
end
```
### Adapter
The [`Adapter`][agentlightning.Adapter] is a component used by the [`Algorithm`][agentlightning.Algorithm] to transform raw data from the [`LightningStore`][agentlightning.LightningStore] into a format suitable for learning. Runners stream raw spans into the store during execution. Later, the algorithm queries these spans and uses an adapter to convert them into structured data, like training examples for a reinforcement learning model.
For instance, the [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet] processes OpenTelemetry spans to create `(prompt, response, reward)` triplets, which are the fundamental data structure for many RL fine-tuning algorithms.
```mermaid
flowchart LR
Runner -- (1) add_otel_span --> Store
Store -- (2) query_spans --> Algorithm
Algorithm -- (3) spans --> Adapter
Adapter -- (4) transformed data --> Algorithm
```
### LLM Proxy
The [`LLMProxy`][agentlightning.LLMProxy] is an optional bridge component that sits between an agent and the algorithms' resources. It acts as a centralized endpoint for all LLM calls. Usually the proxy URL is added to the store as a special resource, so that the [`Runner`][agentlightning.Runner] can fetch it along with other resources when dequeuing a rollout. During rollouts, the runner invokes the proxy's HTTP endpoint instead of calling a model backend directly.
This design offers several benefits:
1. **Instrumentation:** It automatically captures detailed traces of LLM interactions (prompts, responses, metadata) and sends them to the store, complementing the tracer, especially when the agent's code is hard to instrument directly.
2. **Backend Abstraction:** It provides a unified interface for various LLM backends (OpenAI, Anthropic, local models) and can add features like retry logic, rate limiting, and caching.
3. **Resource Management:** The algorithm can dynamically update which LLM the agent uses (e.g., swapping to a newly fine-tuned model) by simply swapping the backend model the proxy is using, without interrupting the agent's code.
The benefits above seem to be all discussed within the context of model fine-tuning. As a matter of fact, the proxy can be useful for prompt tuning as well. The algorithm can register one of the following two types of endpoints into the proxy:
1. **Endpoint served by the algorithm:** If the algorithm is internally updating the LLM weights (e.g., RL), it can launch an LLM inference engine (i.e., a model server) and register the endpoint URL with the proxy. The proxy then forwards all LLM calls to that endpoint.
2. **Third-party LLM endpoint:** If the algorithm is not updating the LLM weights (e.g., prompt tuning), it can register a third-party LLM endpoint into the proxy.
We show a diagram below that illustrates how the proxy fits into the overall data flow.
```mermaid
sequenceDiagram
autonumber
participant Algo as Algorithm
participant LLMProxy as LLM Proxy
participant Store
participant Runner
participant Agent
Note over Algo,LLMProxy: Algorithm manages LLMProxy as member
loop Over the Dataset
Algo->>Algo: Launch LLM Inference Engine<br>(optional)
Algo->>LLMProxy: Register Inference Engine<br>(optional)
Algo-->>Store: enqueue_rollout
LLMProxy->>Store: Proxy URL added as Resource
Store-->>Runner: dequeue_rollout → AttemptedRollout
Store-->>Runner: get_latest_resources
Runner->>Agent: rollout + resources<br>(LLM Proxy URL as resource)
loop Defined by Agent
Agent-->>LLMProxy: LLM calls
activate LLMProxy
LLMProxy-->>Store: add_span or add_otel_span
LLMProxy-->>Agent: LLM responses
deactivate LLMProxy
Agent-->>Runner: rewards
Runner-->>Store: add_span or add_otel_span
end
Runner-->>Store: update_attempt("finished", status)
Store-->>Algo: query_rollouts + spans
Algo-->>Algo: Update LLM Weights<br>(optional)
end
```
In this diagram, the store receives spans from both the proxy and the runner. We will see a problem later with parallelism where the proxy and runner are in different machines, and spans need to obtain a special counter from the store to ensure the ordering of spans.
### Trainer
The [Trainer][agentlightning.Trainer] is the high-level orchestrator that initializes and connects all major components -- [Algorithm][agentlightning.Algorithm], [Runner][agentlightning.Runner], [LightningStore][agentlightning.LightningStore], [Tracer][agentlightning.Tracer], [Adapter][agentlightning.Adapter], [LLM Proxy][agentlightning.LLMProxy], and [Hook][agentlightning.Hook]. The components can have a lifecycle as long as the trainer. The trainer manages their lifecycles and handles dependency injection, ensuring that every part of the system operates within a consistent and shared environment.
Below, we demonstrate how the components relate to each other and their roles. We first clarify the roles and relationships shown in the diagram:
1. **Owns:** components that the trainer constructs and manages directly (e.g., runner, tracer).
2. **Injects:** components passed into others as dependencies.
3. **References:** weak links for coordination without ownership.
4. **Uses:** components that are temporarily interacted with.
For example, the [LightningStore][agentlightning.LightningStore] is injected into the [Algorithm][agentlightning.Algorithm] and [Runner][agentlightning.Runner]. The [Tracer][agentlightning.Tracer] and [LitAgent][agentlightning.LitAgent] are injected into the runner. The [Adapter][agentlightning.Adapter] and [LLM Proxy][agentlightning.LLMProxy] are injected into the algorithm. The store is further injected into the tracer, adapter and LLM proxy by the runner and algorithm respectively.
```mermaid
flowchart TD
%% === Left side: Algorithm domain ===
subgraph L["Algorithm Side"]
Algorithm["Algorithm<br>(no default)"]
Adapter["Adapter<br>(TracerTraceToTriplet*)"]
LLMProxy["LLM Proxy<br>(no default)"]
Algorithm -.injects.-> Adapter
Algorithm -.injects.-> LLMProxy
end
linkStyle 0,1 stroke:#896978,stroke-width:2px;
%% === Middle: Core trainer and store ===
subgraph M["Core"]
Trainer["Trainer"]
Store["LightningStore<br>(InMemory* default)"]
Trainer --has--> Algorithm
Trainer --has--> Store
Trainer --has--> Adapter
Trainer --has--> LLMProxy
end
linkStyle 2,3,4,5 stroke:#839791,stroke-width:2px;
%% === Right side: Runner side ===
subgraph R["Runner Side"]
Runner["Runner<br>(LitAgentRunner* default)"]
Tracer["Tracer<br>(AgentOpsTracer*)"]
Hooks["Hooks (empty default)"]
Agent["Agent<br>(LitAgent*)"]
Runner -.injects.-> Tracer
Runner -.injects.-> Store
Runner -.injects.-> Agent
Runner -.injects.-> Hooks
Tracer -.injects.-> Store
Hooks -.uses.-> Runner
Hooks -.uses.-> Agent
Hooks -.uses.-> Tracer
end
linkStyle 6,7,8,9,10 stroke:#896978,stroke-width:2px;
linkStyle 11,12,13 stroke:#7a89c2,stroke-width:2px;
%% === Cross-section connections ===
Trainer --has--> Runner
Trainer --has--> Tracer
Trainer --has--> Hooks
Trainer --uses--> Agent
Algorithm -.injects.-> Store
LLMProxy -.injects.-> Store
Agent -.references.-> Trainer
Runner -.references.-> Trainer
Algorithm -.references.-> Trainer
linkStyle 14,15,16 stroke:#839791,stroke-width:2px;
linkStyle 17,20,21,22 stroke:#7a89c2,stroke-width:2px;
linkStyle 18,19 stroke:#896978,stroke-width:2px;
style L fill:none;
style M fill:none;
style R fill:none;
```
## Putting It All Together: A Reinforcement Learning Example (VERL)
[](){ #birds-eye-view-verl-example }
VERL shows how an algorithm consumes the shared infrastructure. For historical reasons, code lives in `agentlightning.algorithm.verl` and `agentlightning.verl`. The latter is legacy and reuses terms like `Trainer` in confusing ways. The former is a thin wrapper that conforms to the new algorithm interface. Future versions will merge the two.
Reinforcement learning aims to learn a policy that takes actions in states to maximize expected reward. For agents, the policy is usually a language model. Inputs are prompts (state). Outputs are generated text (action). A numeric score judges quality (reward). The `(state, action, reward)` **triplet** is the basic learning unit.
In Agent-lightning, the environment is implicit in the agents workflow, which orchestrates one or more LLM calls and often self-judges using rules or additional model calls. During a rollout, the agent emits spans that contain everything needed for RL training, including LLM call traces and numeric judge/reward signals. The "algorithm", on the other hand, have more responsibilities.
1. Providing a language model deployment that is currently learning and improving for the agent to interact with;
2. Preparing the tasks that the agents will perform;
3. Querying the spans generated, extracting triplets, and converting them into a format that the underlying RL library can consume;
4. Updating the language model based on the learning signals.
In the VERL integration, the algorithm launches a chat completion endpoint using `vLLM` and wraps training with `FSDP` for distributed optimization. It enqueues tasks from the dataset. After rollouts finish, it queries spans and converts them to triplets with `TracerTraceToTriplet`. VERLs native training loop then consumes these triplets to update model weights. The workflow can be summarized in the following diagram.
```mermaid
sequenceDiagram
autonumber
participant vLLM as vLLM Chat<br>Completion Endpoint
participant FSDP as FSDP / Megatron<br>Weights Optimizer
participant Algo as Algorithm<br>Main Controller<br>(Main Process)
participant Adapter as TracerTraceToTriplet
participant LLMProxy as LLM Proxy
participant Store as LightningStore
participant Runner as Runner + Agent
Note over Algo,LLMProxy: LLMProxy and Adapter are injected by Trainer as member
Note over vLLM,Algo: Algorithm creates and owns vLLM and FSDP
loop Over the Dataset in Batches
Algo->>vLLM: Create Chat Completion Endpoint
activate vLLM
vLLM->>LLMProxy: Registered as Backend Endpoint
LLMProxy->>Store: Proxy URL added as Resource
par Over data samples in the batch
Algo-->>Store: enqueue_rollout
Store-->>Runner: Dequeue Rollout +<br>Resources (i.e., URL)
loop One Rollout Attempt
Runner-->>LLMProxy: LLM calls
LLMProxy-->>vLLM: Forwarded LLM calls
vLLM-->>LLMProxy: LLM responses
LLMProxy-->>Store: add_span / add_otel_span
LLMProxy-->>Runner: Forwarded LLM responses
Runner-->>Store: add_span / add_otel_span <br> (by tracer, including rewards)
end
Runner-->>Store: update_attempt("finished", status)
end
Algo-->>Store: Poll for completed rollouts + spans
Algo->>vLLM: Chat Completion Endpoint Sleeps
deactivate vLLM
Algo->>Adapter: adapt(spans)
Adapter->>FSDP: Triplets (state, action, reward)
activate FSDP
FSDP-->>Algo: Updated LLM weights
deactivate FSDP
end
```
**Notes:**
1. There are interactions between different components injected into or owned by algorithms in the diagram, such as the output of the adapter feeding into the FSDP optimizer. This is for simplicity of illustration and slightly different from the actual implementation, where it's the algorithm main controller that orchestrates the data flow between components.
2. **On mapping to VERL.** VERL uses a classic RLHF setup where each action is a single token, the state is the full conversation history up to that token, and reward is given at the end. This is very different from our setup where each action is actually a chunk of text, although they are both called RL! Therefore, after the adapter produces triplets, the algorithm converts each `(state, action, reward)` into a VERL trajectory (`DataProto`) with keys like `input_ids`, `position_ids`, `attention_mask`, and `token_level_scores`. That conversion happens after triplet generation and is not shown in the diagram.
## Execution Strategies and Parallelism
Readers might have observed from the diagram above that there is absolutely no communication between (1) runner and agents and (2) algorithm. The only overlap of them is the [Trainer][agentlightning.Trainer] and [LightningStore][agentlightning.LightningStore]. This observation is very clear with the diagram within the trainer section. This design allows us to flexibly scale the runner and algorithm independently, which is crucial for large-scale training.
Agent-lightning packages two executable bundles: a runner bundle ([Runner][agentlightning.Runner], [Tracer][agentlightning.Tracer], [Hook][agentlightning.Hook], [LitAgent][agentlightning.LitAgent]) and an algorithm bundle ([Algorithm][agentlightning.Algorithm], [Adapter][agentlightning.Adapter], [LLM Proxy][agentlightning.LLMProxy]). Both share the [LightningStore][agentlightning.LightningStore]. The trainer initializes and connects the bundles.
```mermaid
graph TD
subgraph Runner_Side["Runner Bundle"]
direction LR
R[Runner] --- T[Tracer] --- H[Hooks] --- A1[Agent]
end
subgraph Algorithm_Side["Algorithm Bundle"]
direction LR
ALG[Algorithm] --- AD[Adapter] --- LLM[LLM Proxy]
end
S[(Store)]
TR[Trainer]
Runner_Side <--> S
Algorithm_Side <--> S
TR --> Runner_Side
TR --> Algorithm_Side
linkStyle 0,1,2,3,4 opacity:0;
```
An [execution strategy][agentlightning.ExecutionStrategy], defined and owned by the trainer, governs how algorithm and runner bundles are placed, connected, scaled, and aborted. It serves four primary purposes.
Execution strategies first determine **bundle placement** — whether the two bundles run in the same thread, process, machine, or across separate machines. They also define **store management**, wrapping the store and specifying how data is shared between bundles.
In terms of **scalability**, the strategy can replicate the runner bundle across multiple threads, processes, or machines to expand throughput on the runner side. The algorithm side remains single-process due to the complexity of parallelization. Mature frameworks such as *DeepSpeed* and *Megatron* already support distributed model training, so scaling of the algorithm bundle is delegated to those implementations.
**Abort handling** is another core responsibility. Aborts may be triggered by normal exits, failures in either bundle, or user interrupts. The trainer must include cancellation interfaces for the bundles so that bundles can be cleanly aborted. When the algorithm bundle exits normally, the strategy signals the runner bundle to terminate. If the runner exits first, no signal is sent to the algorithm, as it may still be processing completed rollouts. In cases of failure or user interruption, the strategy signals both bundles to abort; if a bundle fails to respond, the strategy should attempt a forceful termination.
Agent-lightning currently provides two execution strategies: **shared-memory** and **client-server**, described in the following sections.
### Shared-memory Strategy
[`SharedMemoryExecutionStrategy`][agentlightning.SharedMemoryExecutionStrategy] runs algorithm and runner bundles as threads in one process. The strategy wraps the store with [`LightningStoreThreaded`][agentlightning.LightningStoreThreaded], which guards calls with a lock for safe concurrency.
This is good for lightweight debugging because components share one Python heap and avoid serialization. It is not suitable for heavy RL training or compute-intensive agents.
```mermaid
flowchart TB
subgraph MainProcess
direction TB
subgraph AlgorithmThread [Thread 0]
Algorithm[Algorithm bundle]
end
subgraph RunnerThread1 [Thread 1]
Runner1[Runner bundle #1]
end
subgraph RunnerThread2 [Thread 2]
Runner2[Runner bundle #2]
end
subgraph RunnerThread3 [Thread 3]
RunnerN[Runner bundle #N]
end
LightningStoreFacade[LightningStoreThreaded]
BaseStore[Underlying LightningStore]
end
Algorithm -- async calls --> LightningStoreFacade
Runner1 -- async calls --> LightningStoreFacade
Runner2 -- async calls --> LightningStoreFacade
RunnerN -- async calls --> LightningStoreFacade
LightningStoreFacade -->|thread-safe delegates| BaseStore
```
You can configure which role runs on the main thread. If the main thread runs the algorithm, it is able to spawn multiple runner threads. If it runs a runner, `n_runners` must be 1 and the runner lives on the main thread.
### Client-server Strategy
[](){ #birds-eye-view-client-server-strategy }
[`ClientServerExecutionStrategy`][agentlightning.ClientServerExecutionStrategy] splits concerns across processes. The algorithm bundle starts a [`LightningStoreServer`][agentlightning.LightningStoreServer] (HTTP API) that wraps the underlying store. Runners connect via [`LightningStoreClient`][agentlightning.LightningStoreClient] to call the same interface over REST. The server embeds a client to support algorithm-launched subprocesses (e.g., an LLM proxy worker) that need to talk back to the algorithms process through the same API.
Currently this design introduces an extra wrapper in the Server side (as shown in the diagram), which helps debugging and improves fault tolerance. We might revisit this design in the future and enforce the client to be the only way to communicate with the store.
```mermaid
flowchart TD
subgraph Algorithm Process Group
subgraph StoreServer[LightningStoreServer]
StoreHttpClient[HTTP Client]
StoreHttpServer[HTTP Server]
StoreWrapper[LightningStore Wrapper]
StoreHttpClient -- HTTP --> StoreHttpServer
end
subgraph Algorithm Bundle
Algorithm[Algorithm Main Process]
subgraph Another subprocess
LLMProxy[LLM Proxy]
end
end
LLMProxy -- async calls --> StoreHttpClient
Algorithm -- async calls --> StoreWrapper
end
subgraph RunnerSide ["Runner Side"]
subgraph Runner Process 1
Runner1[Runner bundle #1]
Runner1 -- async calls --> LightningStoreClient1
LightningStoreClient1[LightningStoreClient]
end
subgraph Runner Process 2
Runner2[Runner bundle #2]
Runner2 -- async calls --> LightningStoreClient2
LightningStoreClient2[LightningStoreClient]
end
subgraph Runner Process N
RunnerN[Runner bundle #N]
RunnerN -- async calls --> LightningStoreClientN
LightningStoreClientN[LightningStoreClient]
end
end
LocalStore[Underlying LightningStore]
StoreHttpServer -->|delegates| StoreWrapper
StoreWrapper -->|delegates| LocalStore
LightningStoreClient1 -- HTTP --> StoreHttpServer
LightningStoreClient2 -- HTTP --> StoreHttpServer
LightningStoreClientN -- HTTP --> StoreHttpServer
style RunnerSide fill:none;
```
## Online/Continuous Learning
Continuous learning keeps the algorithm loop running while runners report tasks and spans opportunistically. Key differences from batch mode:
1. The algorithm does not enqueue rollouts from a fixed dataset. Runners report tasks/rollouts and spans spontaneously.
2. The algorithm can wait for rollouts with a expected set of rollout IDs, but more often polls for new rollouts and spans or waits for a count to arrive.
3. The [`Runner`][agentlightning.Runner] processes one rollout at a time via [`step(task)`][agentlightning.Runner.step] instead of exhausting a task queue. It notifies the store when starting a rollout so the store records it.
4. A user or higher-level loop controls which resources the next step uses and when to retry.
[Spans][agentlightning.Span], [Adapter][agentlightning.Adapter] implementations, and the [LLM Proxy][agentlightning.LLMProxy] work the same way.
```mermaid
sequenceDiagram
autonumber
actor User
participant Runner
participant Agent
participant Store as LightningStore
participant Algorithm
Note over Algorithm: Algorithm is long-running and loops continuously
loop Continuous Learning Loop
activate User
opt Decide what to do next
User-->>Store: get_resources_by_id
Store-->>User: Resources
User-->>User: Prepare input for next step
end
User->>Runner: step(input, resources)
activate Runner
Runner-->>Store: Notify: start_rollout(input)
Runner->>Agent: rollout(input, resources)
Agent-->>Runner: add_span / reward spans
Runner-->>Store: add_span or add_otel_span
Runner-->>Store: update_attempt(status="finished")
deactivate Runner
deactivate User
Algorithm->>Store: poll for new rollouts and spans
opt If there is enough new data
Store-->>Algorithm: new spans
Algorithm->>Algorithm: adapt spans → learning signal
Algorithm->>Store: update_resources
end
end
```
+135
View File
@@ -0,0 +1,135 @@
# Serving LLMs under Agent-lightning
Agent-lightning focuses on data, learning signals, and control flow — **not** on running model inference. This deep dive explains how to **serve** a model alongside Agent-lightning so runners can call it reliably, how the **LLM Proxy** fits into the loop, and why **token IDs** matter if you care about correctness in training and evaluation.
## General background on LLM serving
[](){ #general-llm-serving-background }
Serving a model is essential if you want to train it, especially when you use the models own generations as training data. Well briefly review the general background to ensure all readers are aligned.
Modern LLM servers solve a difficult scheduling problem: keeping GPUs fully utilized while handling prompts of different lengths, streaming tokens as they arrive, and fitting large KV caches into limited memory. Techniques like [**continuous batching**](https://www.anyscale.com/blog/continuous-batching-llm-inference) and [**paged attention**](https://arxiv.org/abs/2309.06180) address these challenges. Continuous batching interleaves decoding across requests to reuse weights efficiently; with careful memory planning, it achieves major throughput gains without increasing latency. PagedAttention reduces KV-cache fragmentation so batching remains effective as sequences grow. See [vLLMs PagedAttention paper](https://arxiv.org/abs/2309.06180) and [industry analyses](https://www.baseten.co/blog/continuous-vs-dynamic-batching-for-ai-inference/) for details. Balancing inference correctness and efficiency is difficult — a [recent blog](https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/) from Thinking Machines Labs highlights how inference nondeterminism ultimately affects training.
Beyond scheduling, servers expose an HTTP API, often **OpenAI-compatible** (`/v1/chat/completions` and `/v1/responses`), which is itself a complex stack. In addition to text prompts and chat messages, the API defines many parameters and response fields such as [tool calls](https://platform.openai.com/docs/guides/function-calling), [structured output](https://platform.openai.com/docs/guides/structured-outputs), and [multimodal support](https://platform.openai.com/docs/guides/images-vision). Much effort has been put into implementing all these parameters for many frameworks. Popular engines like **vLLM** and [**SGLang**](https://github.com/sgl-project/sglang) ship with OpenAI-compatible frontends so you can reuse existing client code. [Ollama](https://ollama.com/blog/openai-compatibility) and [llama.cpp](https://llama-cpp-python.readthedocs.io/en/latest/server/) provide similar capabilities. However, because models differ internally, each framework interprets and implements the API slightly differently. Even with identical requests, the tokens passed to the model can vary substantially across frameworks.
## What Agent-lightning expects from a served LLM
Most of the issues above either have workarounds or remain open research problems. Keep them in mind, but the key question is: what does Agent-lightning expect from a served LLM? The answer includes at least two things:
* An OpenAI-compatible **Chat Completions** or **Responses** endpoint the agent can call during rollouts.
* Optional training and debugging signals: **logprobs**, **usage**, and ideally **token IDs**. (OpenAIs public API exposes usage and logprobs, but **not** token IDs — more on [why IDs matter][token-ids-matter] later.)
## Launching a serving framework
For many algorithms, youll start an engine (e.g., **vLLM** or **SGLang**) before rollouts, then shut it down afterward to free GPU memory. Most frameworks provide a one-line “serve” command to launch the OpenAI-compatible server. You can use those to bring up `/v1/chat/completions` with your checkpoint, ensuring streaming and any required tool-calling features are enabled. A working example is shown in [Unsloth SFT](../how-to/unsloth-sft.md).
Weight updates — which occur after each training step — are trickier. Some frameworks like [vLLM](https://vllm.ai/) support hot-updating model weights, but its usually simpler and more reliable to restart the engine to load new weights. For medium-sized tasks (hundreds of rollouts taking 10+ minutes), the restart overhead (under 30 seconds) is typically negligible.
If youre using Agent-lightnings [**VERL**][agentlightning.algorithm.verl.VERL] integration, the algorithm can **manage the server automatically**. The [VERL framework](https://github.com/volcengine/verl) intelligently allocates compute resources and wraps vLLM/SGLang behind an `AsyncLLMServer` abstraction. You can directly use this as the LLM endpoint for agents. Since VERL can spawn multiple vLLM replicas, using [`LLMProxy`][agentlightning.LLMProxy] to manage them adds an additional safety layer.
A full sequence diagram of how [VERL][agentlightning.algorithm.verl.VERL] interacts with the LLM server and proxy is available [here][birds-eye-view-verl-example].
## LLM Proxy
The **LLM Proxy** is a utility class in Agent-lightning, built on [LiteLLM](https://docs.litellm.ai/), that sits between runners and your backend engine(s) or server(s). In Agent-lightning it acts as a single URL registered as a [`Resource`][agentlightning.Resource] in the store, offering three key benefits:
1. **Unified endpoint & hot-swaps.** You can redirect traffic between OpenAI, Anthropic, local vLLM/SGLang, or canary checkpoints without modifying agent code — simply repoint the proxy.
2. **First-class tracing.** The proxy emits **OpenTelemetry** spans for every call and sends them to the [`LightningStore`][agentlightning.LightningStore]. It includes rollout and attempt identifiers in request headers so spans are correctly attributed. Sequence numbers are allocated monotonically via the store to [prevent clock-skew issues][distributed-tracing] and allow reliable reconstruction of execution trees.
3. **Token IDs.** The proxy can return prompt and response token IDs along with the model output. More details are available in the [next section][token-ids-matter].
Operationally, running the proxy alongside the algorithm works best: the algorithm registers the backend (e.g., the vLLM URL) via [`LLMProxy.update_model_list`][agentlightning.LLMProxy.update_model_list], publishes the proxy URL as a resource via [`LightningStore.add_resources`][agentlightning.LightningStore.add_resources], and runners simply use that URL during rollouts. This mirrors many production clientserver setups.
## Token IDs and why they matter
[](){ #token-ids-matter }
This section explains how Agent-lightning handles and uses token IDs — a subtle but important detail for training stability and accuracy.
Most agents interact with LLMs via **Chat Completion APIs**, exchanging chat messages. There are two main approaches to collecting training data from such agents.
!!! note
Tokenization here refers to the process of converting **Chat Messages** into **Token IDs**. Detokenization is the reverse process of converting **Token IDs** back to **Chat Messages**. Normally, the tokenizer is published along with the pretrained model, which includes a vocabulary, special tokens, and a chat template to dealing with chat messages.
**1. Retokenizing chat messages.**
In this approach, you store chat messages as text and let training algorithms **retokenize** them later, as done in many SFT workflows (e.g., [HuggingFace SFT](https://huggingface.co/docs/trl/sft_trainer)).
In practice, weve found this method unstable and less accurate. The chart below compares training results. The retokenization approach is run twice. All settings are the same except for the retokenization approach.
<div style="height:400px">
<canvas data-chart='{
"type": "line",
"data": {
"labels": [0.0, 32.0, 64.0, 96.0, 128.0, 160.0, 192.0, 224.0, 256.0, 288.0, 320.0, 352.0, 384.0, 416.0, 448.0, 480.0],
"datasets": [
{
"label": "With Token IDs from Retokenization",
"data": [0.49, 0.512, 0.54, 0.532, 0.54, 0.466, 0.328, 0.358, 0.348, 0.35, 0.346, 0.372, 0.346, 0.33, 0.346, 0.332],
"spanGaps": true
},
{
"label": "Retokenization (Second Run)",
"data": [0.494, 0.526, 0.536, 0.554, 0.544, 0.556, 0.568, 0.552, 0.45, 0.466, 0.474, 0.47, 0.464, 0.476, 0.488, 0.432],
"spanGaps": true
},
{
"label": "With Token IDs from Engine",
"data": [0.494, 0.522, 0.514, 0.538, 0.53, 0.564, 0.564, 0.586, 0.594, 0.604, 0.618, 0.584, 0.606, 0.558, 0.612, 0.588],
"spanGaps": true
}
]
},
"options": {
"interaction": {
"mode": "nearest",
"intersect": false
},
"plugins": {
"legend": {
"display": true,
"position": "top"
},
"title": {
"display": true,
"text": "Agent Training Results Comparison"
}
},
"scales": {
"x": {
"title": {
"display": true,
"text": "Step"
}
},
"y": {
"title": {
"display": true,
"text": "Reward"
}
}
}
}
}'></canvas>
</div>
This instability has three causes. Firstly, chat template used in different frameworks could be slightly different. For example, one single LLaMA model can work with multiple chat templates (multiple in [vLLM](https://github.com/vllm-project/vllm/tree/1d165d6d859d3c50720f0c07209db2363c4fd33b/examples) and one in [HuggingFace](https://huggingface.co/meta-llama)). It's possible that the chat template used in detokenization is different from the one used in tokenization (this is actually an implementation bug).
Secondly, a word might be generated as two tokens (e.g., `H + AVING`) but later retokenized as `HAV + ING`. The text looks identical, but the token IDs differ from what the model originally produced.
Thirdly, a generated tool call text like `<tool_call>{ "name": ... }</tool_call>` is parsed by tool call parser into an object that is required by chat completion API. Later, the object is rendered back to `<tool_call>{ "name": ... }</tool_call>` and retokenized again, tool call parsing and re-rendering might cause changes in whitespace and formatting. In some situations, JSON errors may even be auto-corrected by the tool call parser — masking the models true generation errors and preventing them from being trained away.
----
**2. Saving token IDs directly.**
The alternative is to save the token IDs generated by the model, as done in RL setups like [Tinker](https://thinkingmachines.ai/tinker/). This requires a training pipeline that treats tokens as first-class entities, meaning agents must communicate with the inference engine at the token level.
However, most agents — especially those built with frameworks like LangChain — rely on OpenAI-compatible APIs and cant tokenize or detokenize themselves. As mentioned [earlier][general-llm-serving-background], implementing this layer manually is complex and error-prone. Some frameworks implement custom solutions (e.g., [VERL Agent Loop](https://github.com/volcengine/verl/blob/4da0d3d3188072772cb2ec817b3d6cf4a463821f/recipe/langgraph_agent/chat_model.py), [Tinker Renderer](https://github.com/thinking-machines-lab/tinker-cookbook/blob/34a6588d7055040c259985d98e71c0140b389ba7/tinker_cookbook/renderers.py)), while others leave it to users (e.g., [SkyRL Search-R1](https://novasky-ai.notion.site/skyrl-searchr1)).
----
A better solution is to use an **OpenAI-compatible API that returns token IDs directly.** This lets agents continue using familiar APIs while capturing token IDs via [tracing](../tutorials/traces.md) for training. The limitation, of course, is that the serving framework must actually support this capability.
When Agent-lightning was first released, we implemented an [instrumented vLLM server](https://github.com/microsoft/agent-lightning/blob/v0.1/agentlightning/instrumentation/vllm.py) that monkey-patched vLLMs OpenAI server to return token IDs. Since then, the Agent-lightning and vLLM teams have collaborated to add this feature directly to [vLLM core](https://github.com/vllm-project/vllm/pull/22587). Starting with **vLLM v0.10.2**, the OpenAI-compatible API includes a [`return_token_ids` parameter](https://docs.vllm.ai/en/v0.10.2/serving/openai_compatible_server.html#api-reference), allowing token IDs to be requested alongside chat messages. SGLang has tracked [similar feature requests](https://github.com/sgl-project/sglang/issues/2634), though its OpenAI-compatible layer doesnt yet support them.
In short, when using vLLM v0.10.2 or newer, [`LLMProxy`][agentlightning.LLMProxy] automatically adds `return_token_ids` to each request so the engine includes token IDs in its response. For older vLLM versions, you still need the instrumented version (via `agl vllm` CLI command).
Finally, if you only save token IDs in spans, it will have its own limitations — if you train one model using spans from another model with a different tokenizer, incompatibilities can arise. In practice, though, spans in Agent-lightning always store both chat messages and token IDs (actually the full request and response objects), allowing you to fall back to retokenization when necessary.
+410
View File
@@ -0,0 +1,410 @@
# 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).
## What's in the Store?
![Store Architecture](../assets/store-api-visualized.svg){ .center }
At a high level:
* **Task Queue** [`enqueue_rollout`][agentlightning.LightningStore.enqueue_rollout] adds work; workers poll with [`dequeue_rollout`][agentlightning.LightningStore.dequeue_rollout]. When a rollout is dequeued, it automatically creates a new attempt associated with itself.
* **Rollouts** A rollout is one unit of work. It has input, metadata, links to resources, and a lifecycle (`queuing → preparing → running → ...`). Valid [RolloutStatus][agentlightning.RolloutStatus] are **`queuing`**, `preparing`, `running`, `succeeded`, `failed`, **`requeuing`**, **`cancelled`**. For algorithms and runners, the rollout can be seen as a whole, without worrying about the internal attempts.
* **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.
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
The status model is intentionally small and operationally clear.
```mermaid
stateDiagram-v2
direction LR
[*] --> preparing: <b>Runner calls</b> dequeue_rollout()<br>or start_rollout()<br>or start_attempt()
preparing --> running: <b>Runner calls</b><br>add_[otel_]span()<br>for the first time
state c_runner <<choice>>
state c_watch <<choice>>
preparing --> c_runner: <b>Runner calls</b><br>update_attempt(...)</b>
running --> c_runner: <b>Runner calls</b><br>update_attempt(...)
running --> c_watch: <b>Watchdog checks</b>
preparing --> c_watch: <b>Watchdog checks</b>
state "Client-set outcome" as Client {
direction TB
succeeded
failed
}
state "Watchdog / policy" as Watch {
direction TB
timeout
unresponsive
}
c_runner --> succeeded: update_attempt(status=succeeded)
c_runner --> failed: update_attempt(status=failed)
c_watch --> timeout: now - start_time > timeout_seconds
c_watch --> unresponsive: now - last_heartbeat > unresponsive_seconds
unresponsive --> running: <b>Runner calls</b><br>add_[otel_]span()
```
Each attempt begins in **preparing**, created either when a rollout is dequeued or explicitly started. It transitions to **running** the first time a span is recorded. From there, a few clear rules govern how it can change:
* When the runner explicitly marks completion, the attempt becomes **succeeded** or **failed** (when the runner catches exception thrown out by the agent).
* When the watchdog detects that the total elapsed time since start exceeds the configured limit, it marks the attempt as **timeout**.
* If heartbeats stop arriving for too long, the watchdog marks it **unresponsive**.
* A new span from the runner can immediately revive an **unresponsive** attempt back to **running**.
!!! info "What's a Watchdog?"
The watchdog enforces timing and liveness rules defined by each rollouts [`RolloutConfig`][agentlightning.RolloutConfig]. Its not a separate thread or service, but a function periodically invoked (e.g., before store mutations) to keep attempts healthy and consistent.
This simple model allows the system to distinguish between normal termination, abnormal stalling, and recoverable interruption without additional state flags.
## Worker Telemetry
Workers track runner-level activity timestamps (`last_heartbeat_time`, `last_dequeue_time`, `last_busy_time`, `last_idle_time`) plus their current rollout assignment. Those fields are now derived automatically:
- [`dequeue_rollout(worker_id=...)`][agentlightning.LightningStore.dequeue_rollout] records which worker polled the queue and refreshes `last_dequeue_time`.
- [`update_attempt(..., worker_id=...)`][agentlightning.LightningStore.update_attempt] drives the worker status machine. Assigning an attempt marks the worker **busy** and stamps `last_busy_time`; finishing with `status in {"succeeded","failed"}` switches to **idle**, while watchdog transitions such as `timeout`/`unresponsive` make the worker **unknown** and clear `current_rollout_id` / `current_attempt_id`.
- [`update_worker(...)`][agentlightning.LightningStore.update_worker] is reserved for heartbeats. It snapshots optional `heartbeat_stats` and always updates `last_heartbeat_time`.
Because every transition flows through these APIs, worker status is derived automatically from rollout execution and heartbeats. Note, however, that calling `update_worker` with a new `worker_id` will create a new worker record with status "unknown" if one does not exist. Thus, while manual status changes are not allowed, new worker records can be created externally via heartbeats.
## Rollout Transition Map
Rollout status is an **aggregated view** of its latest attempts status, with additional transitions for queueing and explicit cancellation.
A rollouts retry behavior is controlled by [`Rollout.config`][agentlightning.types.Rollout] (a [`RolloutConfig`][agentlightning.types.RolloutConfig]). The key fields are:
* [`timeout_seconds`][agentlightning.RolloutConfig.timeout_seconds] maximum wall-clock time for an attempt before it is marked `timeout`.
* [`unresponsive_seconds`][agentlightning.RolloutConfig.unresponsive_seconds] maximum silence between heartbeats before an attempt is marked `unresponsive`.
* [`max_attempts`][agentlightning.RolloutConfig.max_attempts] total number of attempts allowed for the rollout (including the first).
* [`retry_condition`][agentlightning.RolloutConfig.retry_condition] which attempt terminal statuses should trigger a retry (e.g., `["failed", "timeout", "unresponsive"]`).
**How it plays out:** The runner works on attempt `k`. If the attempt ends in a status that is listed in `retry_condition`, and `k < max_attempts`, the rollout moves to **requeuing** and the store creates attempt `k+1`. Otherwise, the rollout becomes **failed** (or **succeeded** if the runner marked it so). `timeout_seconds` and `unresponsive_seconds` are enforced by the watchdog and feed into the same decision flow.
A minimal example of how to use `RolloutConfig`:
```python
from agentlightning import RolloutConfig
# Retry on explicit failures or timeouts, up to 3 attempts in total.
cfg = RolloutConfig(
timeout_seconds=600,
unresponsive_seconds=120,
max_attempts=3,
retry_condition=["failed", "timeout"]
)
# When creating/enqueuing a rollout, attach this config.
# The store will propagate attempt outcomes according to cfg.
rollout = await store.enqueue_rollout(input, config=cfg)
```
| Latest attempt status | Rollout transition | Notes / guards |
| ------------------------------------- | ----------------------------------------- | ------------------------------------------------------------------------------------------------- |
| N/A | `queuing` | Created by `enqueue_rollout()`. |
| `preparing` | `queuing/requeuing``preparing` | Typically `dequeue_rollout()` or `start_rollout()`/`start_attempt()` creates a new attempt. |
| `running` | `preparing/queuing/requeuing``running` | First `add_[otel_]span()` flips the attempt to `running`; rollout follows via `rollout_status_from_attempt`. |
| `succeeded` | `*``succeeded` | Terminal. Rollout `end_time` set. |
| `failed` / `timeout` / `unresponsive` | `*``requeuing` | **Only if** `status ∈ retry_condition ∧ sequence_id < max_attempts`. |
| `failed` / `timeout` / `unresponsive` | `*``failed` | Otherwise (no retries left or retries disabled). |
| `*` | `*``cancelled` | Explicitly set by `update_rollout(status=cancelled)`. |
!!! note "Why aggregation?"
In code, we use `rollout_status_from_attempt()` which actively updates the rollout based on the latest attempt. Reading the table above is usually easier than reverse-engineering the propagation logic in the code: think of the rollouts transitions as *callbacks* on attempt state changes, plus queue/cancel paths.
## Spans
Every traceable operation in a rollout is stored as a [Span][agentlightning.Span]. Spans not only capture fine-grained instrumentation but also act as periodic heartbeats that demonstrate liveness. The first span marks activation; each subsequent one both extends the trace and refreshes the attempts [`last_heartbeat_time`][agentlightning.Attempt.last_heartbeat_time]. If no span arrives within the configured [`unresponsive_seconds`][agentlightning.RolloutConfig.unresponsive_seconds], the watchdog downgrades the attempt to **unresponsive** until activity resumes.
Spans are indexed by `(rollout_id, attempt_id, sequence_id)` where the sequence is monotonic. Tracing analysis tools like [Adapter][agentlightning.Adapter] usually rely on "time order" to reconstruct the trace. However, in a distributed system, the recorded start time and end time of a span are not necessarily in the right order when they aggregated into a central store. Therefore, we enforce every span creation to retrieve a monotonically increasing [`sequence_id`][agentlightning.Span.sequence_id] from the store before adding the span.
!!! note
In practice, one `sequence_id` can be used to create multiple spans. In that case, the orders between the multiple spans are determined by the order of `start_time` and `end_time` of the spans.
### OpenTelemetry conversion
Runners often produce [OpenTelemetry `ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/#attributes) objects directly. The store normalizes them into [`Span`][agentlightning.Span] as follows:
1. The runner first requests [`get_next_span_sequence_id`][agentlightning.LightningStore.get_next_span_sequence_id] via `sequence_id = await store.get_next_span_sequence_id(rollout_id, attempt_id)`. This guarantees ordering within the attempt regardless of clock skew.
2. `trace_id`, `span_id`, `parent_id`, `name`, `status`, timestamps, attributes, events, links, and resource are copied from the OTEL span. Timestamps are auto-normalized to seconds (nanoseconds are converted).
3. OTEL `SpanContext` and parent context are preserved so downstream tools can correlate traces across systems.
4. Any additional serializable fields present on the `ReadableSpan` are retained in the stored span (after safe JSON serialization), which keeps the representation forward-compatible.
Programmatically this is encapsulated by [`Span.from_opentelemetry(readable_span, rollout_id, attempt_id, sequence_id)`][agentlightning.Span.from_opentelemetry]; [`store.add_otel_span(...)`][agentlightning.LightningStore.add_otel_span] simply wraps the fetch-then-add flow. The end result is a store span that is stable to sort, merge, and query, while still preserving the richness of the original OTEL payload.
!!! tip
[`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
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.
Check whether the store supports OTLP traces via the [`capabilities["otlp_traces"]`][agentlightning.LightningStore.capabilities] property.
## Implementation Overview
The `agentlightning.store` module is organized into two distinct layers plus optional wrappers:
```mermaid
classDiagram
direction TB
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
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.
Database-based stores like [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore] are inherently thread-safe through database atomicity guarantees.
## Process Safety and Client-server Store
**[`LightningStoreServer`][agentlightning.LightningStoreServer]** wraps another underlying store and runs a FastAPI app to expose the store API over HTTP. [`LightningStoreClient`][agentlightning.LightningStoreClient] is a small [`LightningStore`][agentlightning.LightningStore] implementation that talks to the HTTP API.
!!! warning
The server HTTP API is not considered a stable API at this moment. Users are encouraged to use the [`LightningStoreClient`][agentlightning.LightningStoreClient] to communicate with the server as a stable interface.
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 wont 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 `/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) dont hang by reusing a session from another loop.
Minimal lifecycle:
```python
import agentlightning as agl
# Server (owner process)
in_memory_store = agl.InMemoryLightningStore()
server = agl.LightningStoreServer(store=in_memory_store, host="0.0.0.0", port=4747)
await server.start() # starts uvicorn in a daemon thread and waits for /health
# or keep your own event loop and stop via await server.stop()
# await server.run_forever()
# Client (same or different process)
client = agl.LightningStoreClient("http://localhost:4747")
print(await client.query_rollouts(status_in=["queuing"]))
await client.close()
await server.stop()
```
Another approach is to use a dedicated command line to start a long running server process, possibly sharable across multiple processes. In the main process, you can always use [`LightningStoreClient`][agentlightning.LightningStoreClient] to communicate with the server.
```bash
agl store --port 4747
```
!!! note
[`LightningStoreClient.wait_for_rollouts`][agentlightning.LightningStoreClient.wait_for_rollouts] intentionally enforces a tiny timeout (≤ 0.1s) to avoid blocking event loops. Poll with short timeouts or compose with `asyncio.wait_for` at a higher layer.