Agent Lightning

https://microsoft.github.io/agent-lightning/0.2.1
Admin
Agent Lightning is a trainer for AI agents that enables optimization with minimal code changes,...
Tokens:66,590
Snippets:305
Trust Score:8.3
Update:11 months ago
Show doc for...
Tokens:
Raw
### Install Agent-Lightning with VERL Support (Automated)

Source: https://microsoft.github.io/agent-lightning/0.2.1/tutorials/installation

Installs Agent-Lightning with dependencies for the VERL framework. Automated installation may lead to version conflicts; manual setup is recommended for stability.

```bash
pip install agentlightning[verl]

```

--------------------------------

### Install Agent-Lightning Nightly Build (Test PyPI)

Source: https://microsoft.github.io/agent-lightning/0.2.1/tutorials/installation

Installs the latest experimental features from Agent-Lightning's nightly builds available on Test PyPI. Use with caution as these may be unstable.

```bash
pip install --upgrade --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ agentlightning

```

--------------------------------

### Install Agent-Lightning with VERL Support (Manual Setup)

Source: https://microsoft.github.io/agent-lightning/0.2.1/tutorials/installation

Manually installs Agent-Lightning and its VERL dependencies step-by-step for enhanced stability, ensuring compatibility with specific PyTorch and CUDA versions.

```bash
pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cu128
pip install flash-attn --no-build-isolation
pip install vllm==0.10.2
pip install verl==0.5.0

```

--------------------------------

### Minimal Developer Installation using uv

Source: https://microsoft.github.io/agent-lightning/0.2.1/tutorials/installation

Clones the Agent-Lightning repository and sets up a minimal development environment using 'uv sync --group dev'. This approach is recommended for developers and contributors.

```bash
git clone https://github.com/microsoft/agent-lightning
cd agent-lightning
uv sync --group dev

```

--------------------------------

### Install Agent-Lightning Stable Release (PyPI)

Source: https://microsoft.github.io/agent-lightning/0.2.1/tutorials/installation

Installs or upgrades Agent-Lightning to the latest stable version from PyPI. This is the recommended method for general use.

```bash
pip install --upgrade agentlightning

```

--------------------------------

### Start Runner Process

Source: https://microsoft.github.io/agent-lightning/0.2.1/how-to/write-first-algorithm

Exports the OpenAI API key and starts the runner process which connects to the store to wait for rollouts. Requires an OpenAI API key.

```bash
export OPENAI_API_KEY=sk-... # Your OpenAI API key
python apo_custom_algorithm.py runner
```

--------------------------------

### Install Agent-Lightning with APO Support

Source: https://microsoft.github.io/agent-lightning/0.2.1/tutorials/installation

Installs Agent-Lightning with additional dependencies required for the APO algorithm module. Ensure OpenAI Python SDK version 2.0 or newer is installed separately.

```bash
pip install agentlightning[apo]

```

--------------------------------

### Start LightningStore Server

Source: https://microsoft.github.io/agent-lightning/0.2.1/how-to/write-first-algorithm

Starts the LightningStore server that facilitates communication between the algorithm and runner components. The store listens on port 4747 by default.

```bash
agl store
```

--------------------------------

### vLLM Server Startup with subprocess

Source: https://microsoft.github.io/agent-lightning/0.2.1/how-to/unsloth-sft

This code snippet demonstrates starting a vLLM inference server using subprocess. It waits for the server to be ready and then initializes an OpenAI client to interact with it. This setup incorporates auto tool choice and a Hermes tool-call parser.

```Python
from openai import OpenAI

subprocess.Popen(
    [\n        "vllm", "serve", model_path, "--port", str(port),\n        "--enable-auto-tool-choice", "--tool-call-parser", "hermes"\n    ]
)

# Wait for the server to be ready
url = f"http://localhost:{port}/health"
start = time.time()
client = httpx.Client()

while True:\n    if client.get(url).status_code == 200:\n        break

server_address = f"http://localhost:{port}/v1"

# Try using the vLLM server
openai = OpenAI(base_url=server_address)
...

```

--------------------------------

### Launch Algorithm Process in Bash for Agent-Lightning

Source: https://microsoft.github.io/agent-lightning/0.2.1/tutorials/parallelize

This code snippet sets environment variables to run the algorithm role on a designated server machine. It binds to all interfaces for the LightningStoreServer and starts the training script. Requires Python environment with Agent-Lightning. Inputs are host, port, and role; outputs are none implicit. Limited to single algorithm process per setup.

```bash
export AGL_SERVER_HOST=0.0.0.0
export AGL_SERVER_PORT=4747
export AGL_CURRENT_ROLE=algorithm

python train_calc_agent.py

```

--------------------------------

### Configure and Run Agent Trainer with Training Data

Source: https://microsoft.github.io/agent-lightning/0.2.1/how-to/train-first-agent

Shows how to configure the Trainer component which orchestrates the entire training workflow. The Trainer manages parallel runners, applies algorithms to improve prompts, and executes the training loop. Configuration includes algorithm setup, parallel execution settings, initial prompt resources, and data adapters. The example demonstrates loading datasets and starting the training process with a room selector agent.

```python
# 1. Configure the Trainer with the algorithm and initial prompt
trainer = agl.Trainer(
    algorithm=algo,
    n_runners=8, # Run 8 agents in parallel to try out the prompts
    initial_resources={
        # The initial prompt template to be tuned
        "prompt_template": prompt_template_baseline()
    },
    # This is used to convert the span data into a message format consumable by APO algorithm
    adapter=agl.TraceToMessages(),
)

# 2. Load datasets: They can be list of task objects consumable by `room_selector`.
dataset_train, dataset_val = ...

# 3. Start the training process!
trainer.fit(
    agent=room_selector,
    train_dataset=dataset_train,
    val_dataset=dataset_val
)
```

--------------------------------

### Start LightningStoreServer and Client in Python

Source: https://microsoft.github.io/agent-lightning/0.2.1/deep-dive/store

Demonstrates initializing and starting a LightningStoreServer and connecting with a LightningStoreClient in the same or different processes. The server runs on a specified host and port using FastAPI, while the client communicates with the server via HTTP. Includes basic query operation and cleanup.

```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=["queuing"]))

await client.close()
await server.stop()

```

--------------------------------

### SFT Dataset Example

Source: https://microsoft.github.io/agent-lightning/0.2.1/how-to/unsloth-sft

This example shows the format of a dataset expected by the TRL SFT Trainer, consisting of input/output message pairs. The LLM learns to generate the 'assistant' response based on this dataset.

```JSON
{\n  "messages": [{\n    "role": "user",\n    "content": "What color is the sky?\"\n  },\n  {\n    "role": "assistant",\n    "content": "It is blue.\"\n  }\n  ]\n\n}
```

--------------------------------

### Download Unsloth Model from Hub (Bash)

Source: https://microsoft.github.io/agent-lightning/0.2.1/how-to/unsloth-sft

Downloads base model for Unsloth training using Hugging Face hub CLI. Ensures local checkpoints are available for both initial and subsequent training iterations. Uses unsloth/Qwen3-4B-Instruct-2507 as example model.

```bash
hf download unsloth/Qwen3-4B-Instruct-2507 --local-dir models/version_0
```

--------------------------------

### Create VERL instance using Agent-lightning (Python)

Source: https://microsoft.github.io/agent-lightning/0.2.1/algorithm-zoo/verl

Demonstrates how to import the agentlightning library and instantiate a VERL object using the shortcut `agl.VERL(...)`. Requires the `agentlightning` package with VERL extras installed. Returns a VERL instance ready for configuration.

```Python
import agentlightning as agl

agl.VERL(...)
```

--------------------------------

### Install APO package via pip

Source: https://microsoft.github.io/agent-lightning/0.2.1/algorithm-zoo/apo

Shows the pip command to install the APO component of the agentlightning package with its dependencies.

```Bash
pip install agentlightning[apo]
```

--------------------------------

### Run Agent with Tracer and LightningStoreClient

Source: https://microsoft.github.io/agent-lightning/0.2.1/how-to/write-first-algorithm

Sets up a runner that connects to the LightningStore and continuously polls for tasks. Uses a Tracer to instrument the agent execution. Requires a running LightningStore server and an agent function like 'simple_agent'.

```python
# Connecting to Store
store = agl.LightningStoreClient("http://localhost:4747")  # or some other address
runner = LitAgentRunner[str](tracer=AgentOpsTracer())
with runner.run_context(agent=simple_agent, store=store):  # <-- where the wrapping and instrumentation happens
    await runner.iter()  # polling for new tasks forever
```

--------------------------------

### Start LightningStoreServer via CLI

Source: https://microsoft.github.io/agent-lightning/0.2.1/deep-dive/store

Command to start a long-running LightningStoreServer process that can be shared across multiple processes. The server listens on the specified port for client connections.

```bash
agl store --port 4747

```

--------------------------------

### Install Agent-lightning VERL dependencies (Shell)

Source: https://microsoft.github.io/agent-lightning/0.2.1/algorithm-zoo/verl

Shows the pip command to install the `agentlightning` package with VERL extras. Ensure a compatible PyTorch version is present to avoid compatibility issues.

```Shell
pip install agentlightning[verl]
```

--------------------------------

### Launch SQL Agent Training in Bash

Source: https://microsoft.github.io/agent-lightning/0.2.1/how-to/train-sql-agent

Executes training scripts from examples/spider for Qwen or LLaMA models using LitSQLAgent and trainer.fit. Requires GPU with 40GB+ memory and VERL extras installed. Use --active-agent flag to train specific variants; export HF_TOKEN for LLaMA.

```bash
python train_sql_agent.py qwen   # Default Qwen-2.5-Coder-1.5B run
python train_sql_agent.py llama  # LLaMA-3.2-1B with llama3_json tool parser

```

--------------------------------

### Configuring Execution Strategies for Trainer in Python

Source: https://microsoft.github.io/agent-lightning/0.2.1/tutorials/parallelize

These snippets demonstrate ways to set up execution strategies for the Trainer, including passing instances, dictionary configurations, and environment variables. The purpose is to define runner execution and server settings. Dependencies include the agl library and an algorithm instance; inputs are strategy parameters, outputs a configured Trainer object. Limitations: Environment vars override some dict settings; requires proper role and port setup for multi-machine use.

```python
strategy = agl.SharedMemoryExecutionStrategy(main_thread="algorithm", n_runners=4)
trainer = agl.Trainer(algorithm=algorithm, strategy=strategy)
```

```python
trainer = agl.Trainer(
    algorithm=algorithm,
    n_runners=8,
    strategy={"type": "cs", "server_port": 9999},
)
```

```python
import os

os.environ["AGL_SERVER_PORT"] = "10000"
os.environ["AGL_CURRENT_ROLE"] = "algorithm"
os.environ["AGL_MANAGED_STORE"] = "0"

trainer = agl.Trainer(algorithm=algorithm, n_runners=8, strategy="cs")
```

```python
trainer = agl.Trainer(
    algorithm=algorithm,
    n_runners=1,
    strategy={
        "type": "cs",
        "server_host": "0.0.0.0",
        "server_port": 9999,
        "main_process": "runner",
    },
)
```

--------------------------------

### trace_context

Source: https://microsoft.github.io/agent-lightning/0.2.1/reference/runner

Starts a new tracing context. This method should be used as a context manager to ensure proper start and capture of traces within a code block. It allows for customization with name, store, rollout ID, and attempt ID.

```APIDOC
## POST /trace_context

### Description
Starts a new tracing context. This should be used as a context manager.

### Method
POST

### Endpoint
/trace_context

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **name** (str, Optional) - Optional name for the tracing context.
- **store** (LightningStore, Optional) - Optional store to add spans.
- **rollout_id** (str, Optional) - Optional rollout ID for the spans.
- **attempt_id** (str, Optional) - Optional attempt ID for the spans.

### Request Example
{
  "name": "my_traced_task",
  "store": "my_store",
  "rollout_id": "123",
  "attempt_id": "456"
}

### Response
#### Success Response (200)
- **AsyncGenerator[LightningSpanProcessor, None]** - The LightningSpanProcessor instance to collect spans.

#### Response Example
{
  "example": "N/A"
}
```

--------------------------------

### Initializing VERL Algorithm with Configuration - Python

Source: https://microsoft.github.io/agent-lightning/0.2.1/algorithm-zoo/verl

The VERL class creates an algorithm instance that uses VERL's PPO runner for training, accepting a config dictionary with overrides for algorithm, data, actor, reference model, and trainer settings merged via Hydra. It requires dependencies like datasets and a trainer. The code demonstrates setting up for a Qwen model with specific batch sizes, optimization parameters, and logging; limitations include advanced customization needing source code modifications until native hooks are available.

```python
from agentlightning.algorithm.verl import VERL

algorithm = VERL(
    config={
        "algorithm": {
            "adv_estimator": "grpo",
            "use_kl_in_reward": False,
        },
        "data": {
            "train_batch_size": 32,
            "max_prompt_length": 4096,
            "max_response_length": 2048,
        },
        "actor_rollout_ref": {
            "rollout": {
                "tensor_model_parallel_size": 1,
                "n": 4,
                "log_prob_micro_batch_size_per_gpu": 4,
                "multi_turn": {"format": "hermes"},
                "name": "vllm",
                "gpu_memory_utilization": 0.6,
            },
            "actor": {
                "ppo_mini_batch_size": 32,
                "ppo_micro_batch_size_per_gpu": 4,
                "optim": {"lr": 1e-6},
                "use_kl_loss": False,
                "kl_loss_coef": 0.0,
                "entropy_coeff": 0,
                "clip_ratio_low": 0.2,
                "clip_ratio_high": 0.3,
                "fsdp_config": {
                    "param_offload": True,
                    "optimizer_offload": True,
                },
            },
            "ref": {
                "log_prob_micro_batch_size_per_gpu": 8,
                "fsdp_config": {"param_offload": True},
            },
            "model": {
                "path": "Qwen/Qwen2.5-1.5B-Instruct",
                "use_remove_padding": True,
                "enable_gradient_checkpointing": True,
            },
        },
        "trainer": {
            "n_gpus_per_node": 1,
            "val_before_train": True,
            "critic_warmup": 0,
            "logger": ["console", "wandb"],
            "project_name": "AgentLightning",
            "experiment_name": "calc_x",
            "nnodes": 1,
            "save_freq": 64,
            "test_freq": 32,
            "total_epochs": 2,
        },
    }
)
trainer.fit(algorithm, train_dataset=my_train_dataset)
```

--------------------------------

### Run Pre-Written Algorithm Script

Source: https://microsoft.github.io/agent-lightning/0.2.1/how-to/write-first-algorithm

Runs the pre-written algorithm script using a single command. Assumes the script handles internal configuration and execution logic.

```bash
python apo_custom_algorithm.py algo
```

--------------------------------

### algo Decorator Usage Examples

Source: https://microsoft.github.io/agent-lightning/0.2.1/reference/algorithm

Examples demonstrating how to use the algo decorator with both synchronous and asynchronous algorithm functions. The decorated functions can expect optional parameters including store, train_dataset, val_dataset, llm_proxy, adapter, and initial_resources. Parameters not expected by the function will not be injected.

```python
from agentlightning.algorithm.decorator import algo

@algo
def batching_algorithm(*, store, train_dataset, val_dataset):
    for sample in train_dataset:
        store.enqueue_rollout(input=sample, mode="train")

@algo
async def async_algorithm(*, store, train_dataset=None, val_dataset=None):
    await store.enqueue_rollout(input={"prompt": "hello"}, mode="train")
```

--------------------------------

### Execute Rollout with Resources

Source: https://microsoft.github.io/agent-lightning/0.2.1/tutorials/debug

Demonstrates how to execute a single rollout within a runner context using the step method. Shows resource management by creating a PromptTemplate and passing it to the step method along with the task input for execution.

```python
with runner.run_context(agent=apo_rollout, store=store):
    resource = agl.PromptTemplate(template="You are a helpful assistant. {any_question}", engine="f-string")
    rollout = await runner.step(
        "Explain why the sky appears blue using principles of light scattering in 100 words.",
        resources={"main_prompt": resource},
    )
```

--------------------------------

### trace_run

Source: https://microsoft.github.io/agent-lightning/0.2.1/reference/runner

A convenience wrapper to trace the execution of a synchronous function. This method automatically starts and stops tracing for the designated function.

```APIDOC
## POST /trace_run

### Description
A convenience wrapper to trace the execution of a synchronous function.

### Method
POST

### Endpoint
/trace_run

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **func** (Callable[..., Any]) - The synchronous function to execute and trace.
- **args** (Any, Optional) - Positional arguments to pass to the function.
- **kwargs** (Any, Optional) - Keyword arguments to pass to the function.

### Request Example
{
  "func": "my_synchronous_function",
  "args": ["arg1", "arg2"],
  "kwargs": {"key1": "value1", "key2": "value2"}
}

### Response
#### Success Response (200)
- **Any** - The return value of the function.

#### Response Example
{
  "example": "The function's return value"
}
```

--------------------------------

### Algorithm Utility Functions

Source: https://microsoft.github.io/agent-lightning/0.2.1/reference/algorithm

Utility functions for setting initial resources, LLM proxy, store, and trainer for algorithms.

```APIDOC
## PUT /agentlightning/algorithm/set_initial_resources

### Description
Set the initial resources for this algorithm.

### Method
PUT

### Endpoint
/agentlightning/algorithm/set_initial_resources

### Parameters
#### Request Body
- **`resources`** (NamedResources) - Required - The initial resources to set.

## PUT /agentlightning/algorithm/set_llm_proxy

### Description
Set the LLM proxy for this algorithm to reuse when available.

### Method
PUT

### Endpoint
/agentlightning/algorithm/set_llm_proxy

### Parameters
#### Request Body
- **`llm_proxy`** (LLMProxy | None) - Required - The LLMProxy instance configured by the trainer, if any.

## PUT /agentlightning/algorithm/set_store

### Description
Set the store for this algorithm to communicate with the runners. Store is set directly instead of using weakref because its copy is meant to be maintained throughout the algorithm's lifecycle.

### Method
PUT

### Endpoint
/agentlightning/algorithm/set_store

### Parameters
#### Request Body
- **`store`** (LightningStore) - Required - The store instance to communicate with runners.

## PUT /agentlightning/algorithm/set_trainer

### Description
Set the trainer for this algorithm.

### Method
PUT

### Endpoint
/agentlightning/algorithm/set_trainer

### Parameters
#### Request Body
- **`trainer`** (Trainer) - Required - The `Trainer` instance that will handle training and validation.
```

--------------------------------

### trace_run_async

Source: https://microsoft.github.io/agent-lightning/0.2.1/reference/runner

A convenience wrapper to trace the execution of an asynchronous function. This method automatically starts and stops tracing for the designated asynchronous function.

```APIDOC
## POST /trace_run_async

### Description
A convenience wrapper to trace the execution of a single asynchronous function.

### Method
POST

### Endpoint
/trace_run_async

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **func** (Callable[..., Awaitable[Any]]) - The asynchronous function to execute and trace.
- **args** (Any, Optional) - Positional arguments to pass to the function.
- **kwargs** (Any, Optional) - Keyword arguments to pass to the function.

### Request Example
{
  "func": "my_asynchronous_function",
  "args": ["arg1", "arg2"],
  "kwargs": {"key1": "value1", "key2": "value2"}
}

### Response
#### Success Response (200)
- **Any** - The return value of the function.

#### Response Example
{
  "example": "The function's return value"
}
```

--------------------------------

### GET /dequeue_rollout

Source: https://microsoft.github.io/agent-lightning/0.2.1/how-to/write-first-algorithm

This endpoint is used to dequeue a rollout. The Store Server provides this endpoint to allow Runners to retrieve tasks from the queue.

```APIDOC
## GET /dequeue_rollout

### Description
Dequeues a rollout.  The Store Server provides this endpoint to allow Runners to retrieve tasks from the queue.

### Method
GET

### Endpoint
/dequeue_rollout

### Parameters
#### Path Parameters
- None

#### Query Parameters
- None

#### Request Body
- None

### Request Example
{
  "example": "N/A - No request body is defined in the documentation" 
}

### Response
#### Success Response (200)
- None

#### Response Example
{
  "example": "N/A - No response body explicitly mentioned. Observed status code is 200." 
}
```

--------------------------------

### Dry-run Trainer with Baseline Algorithm - Python

Source: https://microsoft.github.io/agent-lightning/0.2.1/tutorials/debug

Demonstrates how to initialize a `Trainer` and run it in development mode with a sample dataset and prompt resources. This method allows exercising the training infrastructure without full RL or SFT stacks, using the default `agentlightning.Baseline` algorithm.

```python
import agentlightning as agl

dataset: agl.Dataset[str] = [
    "Explain why the sky appears blue using principles of light scattering in 100 words.",
    "What's the capital of France?",
]
resource = agl.PromptTemplate(template="You are a helpful assistant. {any_question}", engine="f-string")

trainer = agl.Trainer(
    n_runners=1,
    initial_resources={"main_prompt": resource},
)
trainer.dev(apo_rollout, dataset)
```

--------------------------------

### GET /query_spans/{ro-c67eaa9016b6}

Source: https://microsoft.github.io/agent-lightning/0.2.1/how-to/write-first-algorithm

This endpoint retrieves spans associated with a specific rollout ID. It's used for querying span data related to a particular rollout.

```APIDOC
## GET /query_spans/{ro-c67eaa9016b6}

### Description
Retrieves spans associated with a specific rollout ID. It's used for querying span data related to a particular rollout.

### Method
GET

### Endpoint
/query_spans/{ro-c67eaa9016b6}

### Parameters
#### Path Parameters
- **ro-c67eaa9016b6** (string) - Required - The rollout ID to query spans for.

#### Query Parameters
- None

#### Request Body
- None

### Request Example
{
  "example": "N/A - No request body is defined in the documentation" 
}

### Response
#### Success Response (200)
- None

#### Response Example
{
  "example": "N/A - No response body explicitly mentioned. Observed status code is 200." 
}
```
Tokens:
Raw