github-featured/microsoft--agent-lightning
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

Embed algorithm into trainer (#99)

Browse Source
This commit is contained in:
Yuge Zhang
2025-09-19 18:42:58 +08:00
committed by GitHub
parent 0a92600a4c
commit ddc8997b8c
19 changed files with 989 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agentlightning/__init__.py
+1 -1
View File
@@ -4,7 +4,7 @@ __version__ = "0.1.2"
from .client import AgentLightningClient, DevTaskLoader
from .config import lightning_cli
from .litagent import LitAgent
from .litagent import *
from .logging import configure_logger
from .reward import reward
from .server import AgentLightningServer
agentlightning/algorithm/__init__.py
+5
View File
@@ -0,0 +1,5 @@
# Copyright (c) Microsoft. All rights reserved.
from .base import BaseAlgorithm
__all__ = ["BaseAlgorithm"]
agentlightning/algorithm/base.py
+99
View File
@@ -0,0 +1,99 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import weakref
from typing import Any, Optional, TYPE_CHECKING
from agentlightning.client import AgentLightningClient
from agentlightning.types import Dataset
if TYPE_CHECKING:
from agentlightning.trainer import Trainer
class BaseAlgorithm:
"""Algorithm is the strategy, or tuner to train the agent."""
_trainer_ref: weakref.ReferenceType[Trainer] | None = None
def set_trainer(self, trainer: Trainer) -> None:
"""
Set the trainer for this algorithm.
Args:
trainer: The Trainer instance that will handle training and validation.
"""
self._trainer_ref = weakref.ref(trainer)
@property
def trainer(self) -> Trainer:
"""
Get the trainer for this algorithm.
Returns:
The Trainer instance associated with this agent.
"""
if self._trainer_ref is None:
raise ValueError("Trainer has not been set for this agent.")
trainer = self._trainer_ref()
if trainer is None:
raise ValueError("Trainer reference is no longer valid (object has been garbage collected).")
return trainer
def __call__(self, *args: Any, **kwargs: Any) -> Any:
return self.run(*args, **kwargs)
def run(
self,
train_dataset: Optional[Dataset[Any]] = None,
validation_dataset: Optional[Dataset[Any]] = None,
dev_dataset: Optional[Dataset[Any]] = None,
) -> None:
"""Subclasses should implement this method to implement the algorithm.
Args:
train_dataset: The dataset to train on. Not all algorithms require a training dataset.
val_dataset: The dataset to validate on. Not all algorithms require a validation dataset.
Returns:
Algorithm should refrain from returning anything. It should just run the algorithm.
"""
raise NotImplementedError("Subclasses must implement run().")
def get_client(self) -> AgentLightningClient:
"""Get the client to communicate with the algorithm.
If the algorithm does not require a server-client communication, it can also create a mock client
that never communicates with itself.
Returns:
The AgentLightningClient instance associated with this algorithm.
"""
raise NotImplementedError("Subclasses must implement get_client().")
def fit(
self,
agent: Any,
train_data: Optional[Dataset[Any]] = None,
test_data: Optional[Dataset[Any]] = None,
dev_data: Optional[Dataset[Any]] = None,
trainer: Optional[Trainer] = None,
) -> None:
"""Fit the algorithm with the provided agent and datasets.
Args:
agent: The agent to train.
train_data: The training dataset.
test_data: The test dataset.
dev_data: The development dataset.
trainer: The trainer instance.
"""
if trainer is not None:
self.set_trainer(trainer)
self.run(
train_dataset=train_data,
validation_dataset=test_data,
dev_dataset=dev_data,
)
agentlightning/algorithm/verl/__init__.py
+5
View File
@@ -0,0 +1,5 @@
# Copyright (c) Microsoft. All rights reserved.
from .interface import VERL
__all__ = ["VERL"]
agentlightning/algorithm/verl/interface.py
+39
View File
@@ -0,0 +1,39 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any, Optional
from hydra import initialize, compose
from omegaconf import OmegaConf
from agentlightning.verl.entrypoint import run_ppo
from agentlightning.client import AgentLightningClient
from agentlightning.algorithm.base import BaseAlgorithm
from agentlightning.types import Dataset
class VERL(BaseAlgorithm):
def __init__(self, config: dict):
super().__init__()
# Compose the base config exactly like your decorator:
with initialize(version_base=None, config_path="pkg://agentlightning/verl"):
base_cfg = compose(config_name="config")
# Merge your dict overrides
override_conf = OmegaConf.create(config)
self.config = OmegaConf.merge(base_cfg, override_conf)
def run(
self,
train_dataset: Optional[Dataset[Any]] = None,
validation_dataset: Optional[Dataset[Any]] = None,
dev_dataset: Optional[Dataset[Any]] = None,
) -> None:
if dev_dataset is not None:
raise ValueError("dev_dataset is not supported for VERL.")
run_ppo(self.config, train_dataset, validation_dataset)
def get_client(self) -> AgentLightningClient:
port = self.config.agentlightning.port
return AgentLightningClient(endpoint=f"http://localhost:{port}")
agentlightning/litagent.py
+311 -30
View File
@@ -2,11 +2,13 @@
from __future__ import annotations
import functools
import inspect
import logging
import weakref
from typing import Any, List, Dict, Union, Optional, TYPE_CHECKING
from typing import Any, Callable, Coroutine, List, Dict, Union, Optional, TYPE_CHECKING, TypeVar, Generic
from .types import NamedResources, Rollout, Task, TaskInput, Triplet, RolloutRawResult
from .types import LLM, NamedResources, Rollout, Task, TaskInput, Triplet, RolloutRawResult
if TYPE_CHECKING:
from .trainer import Trainer
@@ -16,8 +18,27 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
T = TypeVar("T")
class LitAgent:
__all__ = [
"LitAgent",
"LitAgentLLM",
"llm_rollout",
"rollout",
]
def is_v0_1_rollout_api(func: Callable) -> bool:
"""Check if the rollout API is v0.1.
Inspect the function signature to see if it has a rollout_id parameter.
Args:
func: The function to check.
"""
return "rollout_id" in inspect.signature(func).parameters
class LitAgent(Generic[T]):
"""Base class for the training and validation logic of an agent.
Developers should subclass this class and implement the rollout methods
@@ -38,6 +59,27 @@ class LitAgent:
self._trainer_ref: weakref.ReferenceType[Trainer] | None = None
self._runner_ref: weakref.ReferenceType[AgentRunner] | None = None
@property
def is_async(self) -> bool:
"""
Check if the agent implements asynchronous rollout methods.
Override this property for customized async detection logic.
Returns:
True if the agent has custom async rollout methods, False otherwise.
"""
return (
(
hasattr(self, "training_rollout_async")
and self.__class__.training_rollout_async is not LitAgent.training_rollout_async
)
or (
hasattr(self, "validation_rollout_async")
and self.__class__.validation_rollout_async is not LitAgent.validation_rollout_async
)
or (hasattr(self, "rollout_async") and self.__class__.rollout_async is not LitAgent.rollout_async)
)
def set_trainer(self, trainer: Trainer) -> None:
"""
Set the trainer for this agent.
@@ -122,20 +164,23 @@ class LitAgent:
logging. By default, this is a no-op.
"""
def training_rollout(self, task: TaskInput, rollout_id: str, resources: NamedResources) -> RolloutRawResult:
"""Defines the agent's behavior for a single training task.
def rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
"""Main entry point for executing a rollout.
This method should contain the logic for how the agent processes an
input, uses the provided resources (like LLMs or prompts), and
produces a result.
This method determines whether to call the synchronous or
asynchronous rollout method based on the agent's implementation.
If you don't wish to implement both training rollout and validation
rollout separately, you can just implement `rollout` which will work for both.
Args:
task: The task object received from the server, containing the
input data and metadata.
rollout_id: A unique identifier for the rollout, used for tracking
and reporting purposes.
resources: A dictionary of named resources (e.g., LLMs, prompt
templates) for the agent to use.
rollout: The full rollout object, please avoid from directly modifying it.
Most agents should only use `task` and `resources`. Use `rollout`
only if you need to access metadata like `rollout_id`.
Returns:
The result of the rollout, which can be one of:
@@ -146,9 +191,51 @@ class LitAgent:
- A list of dictionaries for any trace spans.
- A complete `Rollout` object for full control over reporting.
"""
raise NotImplementedError("Subclasses must implement the `training_rollout` method.")
raise NotImplementedError("Agents must implement the `rollout` method.")
def validation_rollout(self, task: TaskInput, rollout_id: str, resources: NamedResources) -> RolloutRawResult:
async def rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
"""Asynchronous version of the main rollout method.
This method determines whether to call the synchronous or
asynchronous rollout method based on the agent's implementation.
Args:
task: The task object received from the server, containing the
input data and metadata.
resources: A dictionary of named resources (e.g., LLMs, prompt
templates) for the agent to use.
rollout: The full rollout object, please avoid from directly modifying it.
Most agents should only use `task` and `resources`. Use `rollout`
only if you need to access metadata like `rollout_id`.
Returns:
The result of the rollout, which can be one of:
- None. The tracing should be handled by the agent runner.
- A float representing the final reward.
- A list of `Triplet` objects for detailed, step-by-step feedback.
- A list of `ReadableSpan` objects for OpenTelemetry tracing.
- A list of dictionaries for any trace spans.
- A complete `Rollout` object for full control over reporting.
"""
raise NotImplementedError("Agents must implement the `rollout_async` method for async operations.")
def training_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
"""Defines the agent's behavior for a single training task.
This method should contain the logic for how the agent processes an
input, uses the provided resources (like LLMs or prompts), and
produces a result.
Args:
task: The task object received from the server, containing the
input data and metadata.
resources: A dictionary of named resources (e.g., LLMs, prompt
templates) for the agent to use.
rollout: The full rollout object, please avoid from directly modifying it.
"""
return self.rollout(task, resources, rollout)
def validation_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
"""Defines the agent's behavior for a single validation task.
By default, this method redirects to `training_rollout`. Override it
@@ -157,19 +244,16 @@ class LitAgent:
Args:
task: The task object received from the server, containing the
input data and metadata.
rollout_id: A unique identifier for the validation rollout,
used for tracking and reporting purposes.
resources: A dictionary of named resources for the agent to use.
rollout: The full rollout object, avoid from modifying it.
Returns:
The result of the validation rollout. See `training_rollout` for
The result of the validation rollout. See `rollout` for
possible return types.
"""
return self.training_rollout(task, rollout_id, resources)
return self.rollout(task, resources, rollout)
async def training_rollout_async(
self, task: TaskInput, rollout_id: str, resources: NamedResources
) -> RolloutRawResult:
async def training_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
"""Asynchronous version of `training_rollout`.
This method should be implemented by agents that perform asynchronous
@@ -177,18 +261,16 @@ class LitAgent:
Args:
task: The task object received from the server.
rollout_id: A unique identifier for the training rollout,
used for tracking and reporting purposes.
resources: A dictionary of named resources for the agent to use.
rollout: The full rollout object, avoid from modifying it.
Returns:
The result of the asynchronous training rollout.
The result of the asynchronous training rollout. See `rollout` for
possible return types.
"""
raise NotImplementedError("Async agents must implement the `training_rollout_async` method.")
return await self.rollout_async(task, resources, rollout)
async def validation_rollout_async(
self, task: TaskInput, rollout_id: str, resources: NamedResources
) -> RolloutRawResult:
async def validation_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
"""Asynchronous version of `validation_rollout`.
By default, this method redirects to `training_rollout_async`.
@@ -196,11 +278,210 @@ class LitAgent:
Args:
task: The task object received from the server.
rollout_id: A unique identifier for the validation rollout,
used for tracking and reporting purposes.
resources: A dictionary of named resources for the agent to use.
rollout: The full rollout object, avoid from modifying it.
Returns:
The result of the asynchronous validation rollout.
The result of the asynchronous validation rollout. See `rollout` for
possible return types.
"""
return await self.training_rollout_async(task, rollout_id, resources)
return await self.rollout_async(task, resources, rollout)
LlmRolloutFunc = Union[
Callable[[T, LLM, Rollout], RolloutRawResult],
Callable[[T, LLM], RolloutRawResult],
Callable[[T, LLM, Rollout], Coroutine[Any, Any, RolloutRawResult]],
Callable[[T, LLM], Coroutine[Any, Any, RolloutRawResult]],
]
class LitAgentLLM(LitAgent[T]):
"""A specialized LitAgent that wraps a function-based rollout that accepts
dynamically a task input and a configured LLM.
This class allows users to define agent behavior using a simple function
that takes task input and an LLM resource, rather than implementing a full
LitAgent subclass.
"""
def __init__(self, llm_rollout_func: LlmRolloutFunc[T], *, trained_agents: Optional[str] = None) -> None:
"""
Initialize the LitAgentLLM with an LLM rollout function.
Args:
llm_rollout_func: A function that defines the agent's behavior.
Can be sync or async, and can optionally accept a Rollout parameter.
trained_agents: Optional string representing the trained agents.
This can be used to track which agents have been trained by this instance.
"""
super().__init__(trained_agents=trained_agents)
self.llm_rollout_func = llm_rollout_func
self._is_async = inspect.iscoroutinefunction(llm_rollout_func)
self._accepts_rollout = "rollout" in inspect.signature(llm_rollout_func).parameters
# Copy function metadata to preserve type hints and other attributes
functools.update_wrapper(self, llm_rollout_func)
def __call__(self, *args, **kwargs):
"""Make the agent instance callable, preserving the original function behavior."""
return self.llm_rollout_func(*args, **kwargs)
@property
def is_async(self) -> bool:
return self._is_async
def rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
"""Execute a synchronous rollout using the wrapped function.
Args:
task: The task input data.
resources: Dictionary of named resources including LLMs.
rollout: The rollout object with metadata.
Returns:
The result from the wrapped rollout function.
"""
if self._is_async:
raise RuntimeError("This LitAgentLLM uses an async function. Use rollout_async instead.")
# Find the first LLM resource
llm = self._get_llm_resource(resources)
if self._accepts_rollout:
return self.llm_rollout_func(task, llm=llm, rollout=rollout) # type: ignore
else:
return self.llm_rollout_func(task, llm=llm) # type: ignore
async def rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
"""Execute an asynchronous rollout using the wrapped function.
Args:
task: The task input data.
resources: Dictionary of named resources including LLMs.
rollout: The rollout object with metadata.
Returns:
The result from the wrapped rollout function.
"""
if not self._is_async:
raise RuntimeError("This LitAgentLLM uses a sync function. Use rollout instead.")
# Find the first LLM resource
llm = self._get_llm_resource(resources)
if self._accepts_rollout:
return await self.llm_rollout_func(task, llm=llm, rollout=rollout) # type: ignore
else:
return await self.llm_rollout_func(task, llm=llm) # type: ignore
def _get_llm_resource(self, resources: NamedResources) -> LLM:
"""Extract the first LLM resource from the resources dictionary.
Args:
resources: Dictionary of named resources.
Returns:
The first LLM resource found.
Raises:
ValueError: If no LLM resource is found.
"""
resource_found: LLM | None = None
for name, resource in resources.items():
if isinstance(resource, LLM):
if resource_found is not None:
logger.warning(f"Multiple LLM resources found in resources. Using the first one: '{name}'.")
break
resource_found = resource
if resource_found is None:
raise ValueError("No LLM resource found in the provided resources.")
return resource_found
def llm_rollout(func: LlmRolloutFunc[T], *, trained_agents: Optional[str] = None) -> LitAgentLLM[T]:
"""Create a LitAgentLLM from a function that takes (task, llm[, rollout]).
This decorator allows you to define an agent using a simple function
instead of creating a full LitAgent subclass. The returned LitAgentLLM
instance is callable, preserving the original function's behavior.
Args:
func: A function that defines the agent's behavior. Can be:
- sync: (task, llm) -> result
- sync with rollout: (task, llm, rollout) -> result
- async: async (task, llm) -> result
- async with rollout: async (task, llm, rollout) -> result
trained_agents: Optional string representing trained agents.
Returns:
A callable LitAgentLLM instance that preserves the original function's
type hints and behavior while providing all agent functionality.
Example:
@llm_rollout
def my_agent(task, llm):
# Agent logic here
return response
# Function is still callable with original behavior
result = my_agent(task, llm)
# Agent methods are also available
result = my_agent.rollout(task, resources, rollout)
"""
return LitAgentLLM(func, trained_agents=trained_agents)
def rollout(func: Union[LlmRolloutFunc[T], Callable], *, trained_agents: Optional[str] = None) -> LitAgent[T]:
"""Create a LitAgent from a function, automatically detecting the appropriate type.
This function inspects the provided callable and creates the appropriate
agent type based on its signature. The returned agent instance is callable,
preserving the original function's behavior and type hints.
Args:
func: A function that defines the agent's behavior.
trained_agents: Optional string representing trained agents.
Returns:
A callable LitAgent subclass instance that preserves the original function's
type hints and behavior while providing all agent functionality.
Example:
@rollout
def my_agent(task, llm):
client = OpenAI(base_url=llm.endpoint)
response = client.chat.completions.create(
model=llm.model,
messages=[{"role": "user", "content": task.input}],
)
# Function is still callable with original behavior
result = my_agent(task, llm)
# Agent methods are also available
result = my_agent.rollout(task, resources, rollout)
Raises:
NotImplementedError: If the function signature doesn't match any known patterns.
"""
sig = inspect.signature(func)
params = list(sig.parameters.keys())
# Check if it matches the LLM rollout API pattern
# Should have at least 2 params, with the second one being 'llm' or typed as LLM
if len(params) >= 2:
second_param = sig.parameters[params[1]]
# Check if the second parameter is named 'llm' or has LLM type annotation
if second_param.name == "llm" or (
second_param.annotation != inspect.Parameter.empty
and (second_param.annotation == LLM or str(second_param.annotation).endswith("LLM"))
):
return llm_rollout(func, trained_agents=trained_agents)
raise NotImplementedError(
f"Function signature {sig} does not match any known agent patterns. "
"Expected signatures: (task, llm[, rollout]) or async (task, llm[, rollout])"
)
agentlightning/runner.py
+15 -5
View File
@@ -12,7 +12,7 @@ import agentops
from opentelemetry.sdk.trace import ReadableSpan
from .client import AgentLightningClient
from .litagent import LitAgent
from .litagent import LitAgent, is_v0_1_rollout_api
from .types import Rollout, Task, Triplet, RolloutRawResult
from .types import ParallelWorkerBase
from .tracer.base import BaseTracer
@@ -157,7 +157,7 @@ class AgentRunner(ParallelWorkerBase):
logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
return False
rollout_obj = Rollout(rollout_id=task.rollout_id) # Default empty rollout
rollout_obj = Rollout(rollout_id=task.rollout_id, task=task) # Default empty rollout
try:
try:
@@ -169,7 +169,12 @@ class AgentRunner(ParallelWorkerBase):
start_time = time.time()
rollout_method = self.agent.training_rollout if task.mode == "train" else self.agent.validation_rollout
# Pass the task input, not the whole task object
result = rollout_method(task.input, task.rollout_id, resources_update.resources)
if is_v0_1_rollout_api(rollout_method):
result = rollout_method(
task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources # type: ignore
) # type: ignore
else:
result = rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj)
rollout_obj = self._to_rollout_object(result, task.rollout_id)
end_time = time.time()
logger.info(
@@ -226,7 +231,7 @@ class AgentRunner(ParallelWorkerBase):
logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
return False
rollout_obj = Rollout(rollout_id=task.rollout_id) # Default empty rollout
rollout_obj = Rollout(rollout_id=task.rollout_id, task=task) # Default empty rollout
try:
try:
@@ -240,7 +245,12 @@ class AgentRunner(ParallelWorkerBase):
self.agent.training_rollout_async if task.mode == "train" else self.agent.validation_rollout_async
)
# Pass the task input, not the whole task object
result = await rollout_method(task.input, task.rollout_id, resources_update.resources)
if is_v0_1_rollout_api(rollout_method):
result = await rollout_method(
task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources # type: ignore
) # type: ignore
else:
result = await rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj)
rollout_obj = self._to_rollout_object(result, task.rollout_id)
end_time = time.time()
logger.info(
agentlightning/trainer.py
+179 -36
View File
@@ -8,13 +8,13 @@ import signal
import time
from typing import List, Optional, Union
import importlib
import agentops
import warnings
from .client import AgentLightningClient
from .litagent import LitAgent
from .runner import AgentRunner
from .types import ParallelWorkerBase
from .types import Dataset, ParallelWorkerBase
from .algorithm.base import BaseAlgorithm
from .tracer.base import BaseTracer
from .tracer.agentops import AgentOpsTracer
from .tracer.triplet import TripletExporter
@@ -43,6 +43,7 @@ class Trainer(ParallelWorkerBase):
If None, a default `AgentOpsTracer` will be created with the current settings.
triplet_exporter: An instance of `TripletExporter` to export triplets from traces,
or a dictionary with the initialization parameters for the exporter.
algorithm: An instance of `BaseAlgorithm` to use for training.
"""
def __init__(
@@ -54,6 +55,7 @@ class Trainer(ParallelWorkerBase):
daemon: bool = True,
tracer: Union[BaseTracer, str, dict, None] = None,
triplet_exporter: Union[TripletExporter, dict, None] = None,
algorithm: Union[BaseAlgorithm, str, dict, None] = None,
):
super().__init__()
self.n_workers = n_workers
@@ -74,6 +76,8 @@ class Trainer(ParallelWorkerBase):
f"Invalid triplet_exporter type: {type(triplet_exporter)}. Expected TripletExporter, dict, or None."
)
self.algorithm = self._make_algorithm(algorithm)
if not self.daemon:
logger.warning(
"daemon=False. Worker processes are non-daemonic. "
@@ -104,6 +108,73 @@ class Trainer(ParallelWorkerBase):
return AgentOpsTracer(agentops_managed=True, instrument_managed=True, daemon=self.daemon)
raise ValueError(f"Invalid tracer type: {type(tracer)}. Expected BaseTracer, str, dict, or None.")
def _make_algorithm(self, algorithm: Union[BaseAlgorithm, str, dict, None]) -> Optional[BaseAlgorithm]:
"""Creates an algorithm instance based on the provided configuration."""
if isinstance(algorithm, BaseAlgorithm):
return algorithm
if isinstance(algorithm, str):
module_name, class_name = algorithm.rsplit(".", 1)
module = importlib.import_module(module_name)
algorithm_cls = getattr(module, class_name)
return algorithm_cls()
if isinstance(algorithm, dict):
algorithm_type = algorithm.get("type")
if algorithm_type is None:
raise ValueError("algorithm dict must have a 'type' key with the class full name")
module_name, class_name = algorithm_type.rsplit(".", 1)
module = importlib.import_module(module_name)
algorithm_cls = getattr(module, class_name)
# Remove 'type' key and pass remaining keys as kwargs
algorithm_kwargs = {k: v for k, v in algorithm.items() if k != "type"}
return algorithm_cls(**algorithm_kwargs)
if algorithm is None:
return None
raise ValueError(f"Invalid algorithm type: {type(algorithm)}. Expected BaseAlgorithm, str, dict, or None.")
def _extract_client_from_data(
self, data: Union[str, AgentLightningClient, Dataset]
) -> Optional[AgentLightningClient]:
"""Extract client from data if it's a string URL or AgentLightningClient."""
if isinstance(data, str):
if not data.startswith("http://") and not data.startswith("https://"):
raise ValueError("String data must be a valid URL starting with http:// or https://")
return AgentLightningClient(endpoint=data)
elif isinstance(data, AgentLightningClient):
return data
return None
def _extract_dataset_from_data(self, data: Union[str, AgentLightningClient, Dataset]) -> Optional[Dataset]:
"""Extract dataset from data if it's a Dataset."""
if isinstance(data, str) or isinstance(data, AgentLightningClient):
return None
return data
def _determine_backend(
self,
train_data: Union[str, AgentLightningClient, Dataset],
dev_data: Union[str, AgentLightningClient, Dataset, None] = None,
) -> Union[str, AgentLightningClient]:
"""Determine which backend to use for initialization."""
if self.dev:
if dev_data is None:
raise ValueError("dev_data must be provided when dev=True.")
client = self._extract_client_from_data(dev_data)
if client is None:
raise ValueError("dev_data must be a string URL or AgentLightningClient when dev=True.")
return client
else:
client = self._extract_client_from_data(train_data)
if client is None and self.algorithm is None:
raise ValueError(
"train_data must be a string URL or AgentLightningClient when no algorithm is provided."
)
elif client is None and self.algorithm is not None:
# Algorithm will be responsible for creating the client
client = self.algorithm.get_client()
logger.info(f"Algorithm created client: {client}")
return client
return client
def init(self, backend: Union[str, AgentLightningClient]) -> None:
logger.info(f"Initializing Trainer...")
@@ -218,43 +289,107 @@ class Trainer(ParallelWorkerBase):
if proc.name().startswith("AgentLightning-"):
proc.kill()
def _terminate_processes(self, processes: List[multiprocessing.Process]) -> None:
if self.n_workers > 1 and len(processes) > 0:
for i, p in enumerate(processes):
if p.is_alive():
logger.info(f"Terminating worker {i} (name: {p.name}, PID: {p.pid})...")
p.terminate()
else:
logger.info(f"Worker {i} (name: {p.name}, PID: {p.pid}) is not alive or has already terminated.")
for i, p in enumerate(processes):
if p.is_alive():
p.join(timeout=10) # Give some time to terminate
if p.is_alive(): # If still alive, kill
logger.warning(
f"Worker {i} (name: {p.name}, PID: {p.pid}) did not terminate gracefully, killing..."
)
p.kill()
p.join(timeout=10) # Ensure it's reaped
def fit(
self,
agent: LitAgent,
backend: Union[str, AgentLightningClient],
train_data: Union[str, AgentLightningClient, Dataset],
*,
val_data: Union[str, AgentLightningClient, Dataset, None] = None,
dev_data: Union[str, AgentLightningClient, Dataset, None] = None,
dev_backend: Union[str, AgentLightningClient, None] = None,
):
"""Train the agent using the provided data.
Each data argument can be a string URL connecting to a agent-lightning server,
or an AgentLightningClient instance connecting to a server (or mock server), or a dataset.
If no algorithm is provided when instantiating the trainer, the data must be
provided to connecting a server. Otherwise, dataset is also allowed and will be
passed to the algorithm.
If the algorithm is instantiated and there is no URL/client provided,
the algorithm will be responsible for creating a client that will connect to itself.
It can also create a mock client if the algorithm does not require a server.
"""
if dev_backend is not None:
warnings.warn("dev_backend is deprecated. Use dev_data instead.")
if dev_data is not None:
raise ValueError("dev_data and dev_backend cannot be provided at the same time.")
dev_data = dev_backend
# Extract datasets for algorithm if available
train_dataset = self._extract_dataset_from_data(train_data)
val_dataset = self._extract_dataset_from_data(val_data) if val_data else None
dev_dataset = self._extract_dataset_from_data(dev_data) if dev_data else None
# Initialize the algorithm with trainer if provided
if self.algorithm is not None:
self.algorithm.set_trainer(self)
# DO NOT RUN TRAINING HERE. Need to spawn the worker first.
# Determine the backend to use for client-server mode
backend = self._determine_backend(train_data, dev_data)
if self.dev:
if dev_backend is None:
raise ValueError("dev_backend must be provided when dev=True.")
logger.warning(f"Running in dev mode. Using dev backend: {dev_backend}")
self.init(dev_backend)
logger.warning(f"Running in dev mode. Using dev backend: {backend}")
else:
logger.debug(f"Running in non-dev mode. Using backend: {backend}")
self.init(backend)
self.init(backend)
processes: List[multiprocessing.Process] = []
# Determine if the agent is asynchronous.
is_async = (
hasattr(agent, "training_rollout_async")
and agent.__class__.training_rollout_async is not LitAgent.training_rollout_async
)
# Determine if the agent is asynchronous
mode = "asynchronous" if is_async else "synchronous"
mode = "asynchronous" if agent.is_async else "synchronous"
try:
if self.n_workers == 1:
logger.info(f"Running with n_workers=1 ({mode} in main process).")
num_tasks = self._worker_main_loop(agent, 0, is_async)
# Warn if algorithm is set with single worker mode
if self.algorithm is not None:
logger.warning(
"Algorithm is set but using single worker mode. Algorithm will never get the chance to run."
)
# Ideally the single worker should be run in a separate thread or process.
num_tasks = self._worker_main_loop(agent, 0, agent.is_async)
logger.info(f"Single worker mode finished. Tasks processed: {num_tasks}")
# If algorithm is provided and we have datasets, run algorithm after worker completes
if self.algorithm is not None and train_dataset is not None:
logger.info("Running algorithm training after worker completion.")
self.algorithm.run(
train_dataset=train_dataset,
validation_dataset=val_dataset,
dev_dataset=dev_dataset,
)
else:
logger.info(f"Running with n_workers={self.n_workers} ({mode} multiprocessing).")
for i in range(self.n_workers):
process_name = f"AgentLightning-Worker-{i}"
p = multiprocessing.Process(
target=self._worker_main_loop,
args=(agent, i, is_async),
args=(agent, i, agent.is_async),
daemon=self.daemon,
name=process_name,
)
@@ -263,6 +398,17 @@ class Trainer(ParallelWorkerBase):
p.start()
if self.daemon:
# If algorithm is provided and we have datasets, pass them to the algorithm
if self.algorithm is not None:
logger.info("All workers have been spawned. Running algorithm training with provided datasets.")
self.algorithm.run(
train_dataset=train_dataset,
validation_dataset=val_dataset,
dev_dataset=dev_dataset,
)
logger.info("Algorithm exits. Killing the workers.")
self._terminate_processes(processes)
for i, p in enumerate(processes):
p.join() # Wait for the process to complete
logger.info(
@@ -283,29 +429,26 @@ class Trainer(ParallelWorkerBase):
multiprocessing_process._children.clear() # type: ignore
if self.algorithm is not None:
logger.info("Main process continues to run algorithm.")
self.algorithm.run(
train_dataset=train_dataset,
validation_dataset=val_dataset,
dev_dataset=dev_dataset,
)
logger.info("Algorithm exits. Killing the workers.")
self._terminate_processes(processes)
except KeyboardInterrupt:
if self.n_workers > 1 and len(processes) > 0:
logger.info(f"KeyboardInterrupt received. Terminating workers...")
for i, p in enumerate(processes):
if p.is_alive():
logger.info(f"Terminating worker {i} (name: {p.name}, PID: {p.pid})...")
p.terminate()
else:
logger.info(
f"Worker {i} (name: {p.name}, PID: {p.pid}) is not alive or has already terminated."
)
for i, p in enumerate(processes):
if p.is_alive():
p.join(timeout=10) # Give some time to terminate
if p.is_alive(): # If still alive, kill
logger.warning(
f"Worker {i} (name: {p.name}, PID: {p.pid}) did not terminate gracefully, killing..."
)
p.kill()
p.join(timeout=10) # Ensure it's reaped
logger.info("KeyboardInterrupt received. Killing the workers.")
self._terminate_processes(processes)
logger.info(f"Workers terminated or single worker interrupted.")
raise
except Exception as e:
logger.exception(f"Unhandled exception in fit method.")
self._terminate_processes(processes)
logger.info(f"Workers terminated or single worker interrupted.")
raise
finally:
if self.daemon:
self.teardown()
agentlightning/types.py
+21 -1
View File
@@ -1,6 +1,8 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any, Dict, List, Optional, Union, Literal, Annotated
from __future__ import annotations
from typing import Any, Dict, Generic, List, Optional, Protocol, TypeVar, Union, Literal, Annotated
from pydantic import BaseModel, Field, Discriminator
from opentelemetry.sdk.trace import ReadableSpan
@@ -20,8 +22,11 @@ __all__ = [
"ResourcesUpdate",
"GenericResponse",
"ParallelWorkerBase",
"Dataset",
]
T = TypeVar("T")
class Triplet(BaseModel):
"""A standard structure for a single turn in a trajectory."""
@@ -37,6 +42,9 @@ class Rollout(BaseModel):
rollout_id: str
# Echoing the input task
task: Optional[Task] = None
# Primary, high-level feedback
final_reward: Optional[float] = None
@@ -204,3 +212,15 @@ class ParallelWorkerBase:
def teardown(self, *args: Any, **kwargs: Any) -> None:
pass
class Dataset(Protocol, Generic[T]):
"""The general interface for a dataset.
It's currently implemented as a protocol, having a similar interface to torch.utils.data.Dataset.
You don't have to inherit from this class; you can use a simple list if you want to.
"""
def __getitem__(self, index: int) -> T: ...
def __len__(self) -> int: ...
agentlightning/verl/daemon.py
+1 -1
View File
@@ -274,7 +274,7 @@ class AgentModeDaemon:
llm_resource = LLM(
endpoint=f"http://127.0.0.1:{self.proxy_port}/v1",
model=self.train_information.get("model", "default-model"),
sampling_parameters={"temperature": self.train_information.get("temperature", 0.7)},
sampling_parameters={"temperature": self.train_information.get("temperature", 0.7 if is_train else 0.0)},
)
resources: NamedResources = {"main_llm": llm_resource}
resources_id = await self.server.update_resources(resources)
agentlightning/verl/dataset.py
+15
View File
@@ -1,8 +1,12 @@
# Copyright (c) Microsoft. All rights reserved.
import torch
from datasets import Dataset as HuggingFaceDataset
from omegaconf import DictConfig
from verl.utils.dataset.rl_dataset import RLHFDataset
from agentlightning.types import Dataset
class AgentDataset(RLHFDataset):
@@ -20,3 +24,14 @@ class AgentDataset(RLHFDataset):
# Workaround for data proto. At least one tensor is needed.
row_dict["fake_ids"] = torch.ones(1, dtype=torch.int)
return row_dict
class LoadedDataset(AgentDataset):
def __init__(self, dataset: Dataset):
super().__init__([], None, DictConfig({})) # type: ignore
dataset_copy = [dataset[i] for i in range(len(dataset))]
self.dataframe = HuggingFaceDataset.from_list(dataset_copy)
def _read_files_and_tokenize(self):
pass
agentlightning/verl/entrypoint.py
+28 -17
View File
@@ -1,9 +1,12 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any
import hydra
import ray
from .dataset import AgentDataset
from agentlightning.types import Dataset
from .dataset import AgentDataset, LoadedDataset
from .trainer import AgentLightningTrainer
from verl.trainer.ppo.reward import load_reward_manager
from verl.trainer.main_ppo import create_rl_sampler
@@ -11,10 +14,10 @@ from verl.trainer.main_ppo import create_rl_sampler
@hydra.main(config_path="pkg://agentlightning/verl", config_name="config", version_base=None)
def main(config):
run_ppo(config)
run_ppo(config, None, None)
def run_ppo(config) -> None:
def run_ppo(config: Any, train_dataset: Dataset | None, val_dataset: Dataset | None) -> None:
if not ray.is_initialized():
# this is for local ray cluster
ray.init(
@@ -25,12 +28,12 @@ def run_ppo(config) -> None:
)
runner = TaskRunner.remote()
ray.get(runner.run.remote(config))
ray.get(runner.run.remote(config, train_dataset, val_dataset))
@ray.remote(num_cpus=1) # please make sure main_task is not scheduled on head
class TaskRunner:
def run(self, config):
def run(self, config: Any, train_dataset: Dataset | None, val_dataset: Dataset | None):
# print initial config
from pprint import pprint
@@ -123,18 +126,26 @@ class TaskRunner:
from verl.utils.dataset.rl_dataset import collate_fn
# Use our special dataset
train_dataset = AgentDataset(
data_files=config.data.train_files,
tokenizer=tokenizer,
processor=processor,
config=config.data,
)
val_dataset = AgentDataset(
data_files=config.data.val_files,
tokenizer=tokenizer,
processor=processor,
config=config.data,
)
if train_dataset is None:
train_dataset = AgentDataset(
data_files=config.data.train_files,
tokenizer=tokenizer,
processor=processor,
config=config.data,
)
else:
train_dataset = LoadedDataset(train_dataset)
if val_dataset is None:
val_dataset = AgentDataset(
data_files=config.data.val_files,
tokenizer=tokenizer,
processor=processor,
config=config.data,
)
else:
val_dataset = LoadedDataset(val_dataset)
train_sampler = create_rl_sampler(config.data, train_dataset)
trainer = AgentLightningTrainer(
config=config,
docs/quickstart/getting-started.md
+1 -1
View File
@@ -66,7 +66,7 @@ from agentlightning.trainer import Trainer
agent = SimpleAgent()
trainer = Trainer(n_workers=2) # Create 2 parallel workers
trainer.fit(agent, backend="http://127.0.0.1:9997")
trainer.fit(agent, "http://127.0.0.1:9997")
```
The trainer creates separate processes for each worker, allowing them to execute tasks independently. This parallelization significantly speeds up the optimization process - with 2 workers, you can test prompts twice as fast.
examples/apo/apo_client.py
+1 -1
View File
@@ -38,4 +38,4 @@ if __name__ == "__main__":
dotenv.load_dotenv()
agent = SimpleAgent()
trainer = Trainer(n_workers=2)
trainer.fit(agent, backend="http://127.0.0.1:9997")
trainer.fit(agent, "http://127.0.0.1:9997")
examples/calc_x/calc_agent_dev.py
+1 -1
View File
@@ -30,4 +30,4 @@ def dev_task_loader() -> DevTaskLoader:
if __name__ == "__main__":
Trainer(n_workers=1, dev=True, max_tasks=2).fit(CalcAgent(), "http://localhost:9999/", dev_task_loader())
Trainer(n_workers=1, dev=True, max_tasks=2).fit(CalcAgent(), "http://localhost:9999/", dev_data=dev_task_loader())
examples/calc_x/calc_agent_v0_2.py
+121
View File
@@ -0,0 +1,121 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any
import re
from datasets import Dataset
from agentlightning import rollout, Trainer, LLM
from agentlightning.algorithm.verl import VERL
from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams
from calc_agent import get_agent, eval
calculator_mcp_server = StdioServerParams(command="uvx", args=["mcp-server-calculator"])
@rollout
async def calc_agent(task: Any, llm: LLM) -> Any:
async with McpWorkbench(calculator_mcp_server) as workbench:
calc_agent = get_agent(
llm.model,
llm.endpoint,
llm.sampling_parameters.get("temperature", 0.7),
workbench,
)
try:
output_format = "Output the answer when you are ready. The answer should be surrounded by three sharps (`###`), in the form of ### ANSWER: <answer> ###."
prompt = task["question"] + " " + output_format
result = await calc_agent.run(task=prompt)
# evaluate
answer = re.search(r"###\s*ANSWER:\s*(.+?)(\s*###|$)", result.messages[-1].content)
if answer:
answer = answer.group(1)
else:
answer = result.messages[-1].content
except Exception as e:
print("Failure:", str(e))
answer = "None"
reward = await eval(answer, str(task["result"])) # reward is tracked with the decorator
print("answer: {} ground_truth: {} reward: {}".format(answer, task["result"], reward))
def main():
rl_training_config = {
"agentlightning": {
"port": 9999,
},
"algorithm": {
"adv_estimator": "grpo",
"use_kl_in_reward": False,
},
"data": {
"train_files": "data/train.parquet",
"val_files": "data/test_mini.parquet",
"train_batch_size": 32,
"max_prompt_length": 4096,
"max_response_length": 2048,
"truncation": "error",
},
"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-0.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"],
"project_name": "AgentLightningDebug",
"experiment_name": "train_verl",
"nnodes": 1,
"save_freq": 256,
"test_freq": 6,
"total_epochs": 1,
"total_training_steps": 6,
},
}
train_dataset = Dataset.from_parquet("data/train.parquet").to_list()
val_dataset = Dataset.from_parquet("data/test_mini.parquet").to_list()
print("First 5 rows of train dataset:")
print(train_dataset[:5])
print("First 5 rows of val dataset:")
print(val_dataset[:5])
trainer = Trainer(algorithm=VERL(rl_training_config), n_workers=2)
trainer.fit(calc_agent, train_dataset, val_data=val_dataset)
if __name__ == "__main__":
main()
examples/spider/sql_agent.py
+1 -1
View File
@@ -551,4 +551,4 @@ def spider_dev_data():
if __name__ == "__main__":
dotenv.load_dotenv()
agent, trainer = agentlightning.lightning_cli(LitSQLAgent, agentlightning.Trainer)
trainer.fit(agent, os.environ["VERL_API_BASE"], spider_dev_data())
trainer.fit(agent, os.environ["VERL_API_BASE"], dev_data=spider_dev_data())
tests/test_decorator_executability.py
+143
View File
@@ -0,0 +1,143 @@
# Copyright (c) Microsoft. All rights reserved.
"""Test that @llm_rollout and @rollout decorators preserve function executability."""
import inspect
import pytest
from agentlightning.litagent import llm_rollout, rollout, LitAgentLLM
@llm_rollout
def sample_llm_rollout_func(task, llm):
"""A test function with llm_rollout decorator."""
return f"Processed task: {task} with LLM: {llm}"
@rollout
def sample_rollout_func(task, llm):
"""A test function with rollout decorator."""
return f"Processed task: {task} with LLM: {llm}"
def test_llm_rollout_preserves_executability():
"""Test that @llm_rollout decorated functions remain executable."""
test_task = "Hello World"
test_llm = "gpt-4"
# Function should be callable
assert callable(sample_llm_rollout_func)
# Function should execute and return expected result
result = sample_llm_rollout_func(test_task, test_llm)
expected = f"Processed task: {test_task} with LLM: {test_llm}"
assert result == expected
def test_llm_rollout_preserves_metadata():
"""Test that @llm_rollout preserves function metadata."""
# Function name should be preserved
assert sample_llm_rollout_func.__name__ == "sample_llm_rollout_func"
# Docstring should be preserved
assert sample_llm_rollout_func.__doc__ == "A test function with llm_rollout decorator."
def test_llm_rollout_returns_litagent_instance():
"""Test that @llm_rollout returns a LitAgentLLM instance."""
assert isinstance(sample_llm_rollout_func, LitAgentLLM)
# Should have agent methods
assert hasattr(sample_llm_rollout_func, "rollout")
assert hasattr(sample_llm_rollout_func, "rollout_async")
assert hasattr(sample_llm_rollout_func, "training_rollout")
def test_llm_rollout_preserves_signature():
"""Test that @llm_rollout preserves function signature."""
sig = inspect.signature(sample_llm_rollout_func)
params = list(sig.parameters.keys())
# Should have the expected parameters
assert params == ["task", "llm"]
def test_rollout_preserves_executability():
"""Test that @rollout decorated functions remain executable."""
test_task = "Hello World"
test_llm = "gpt-4"
# Function should be callable
assert callable(sample_rollout_func)
# Function should execute and return expected result
result = sample_rollout_func(test_task, test_llm)
expected = f"Processed task: {test_task} with LLM: {test_llm}"
assert result == expected
def test_rollout_preserves_metadata():
"""Test that @rollout preserves function metadata."""
# Function name should be preserved
assert sample_rollout_func.__name__ == "sample_rollout_func"
# Docstring should be preserved
assert sample_rollout_func.__doc__ == "A test function with rollout decorator."
def test_rollout_returns_litagent_instance():
"""Test that @rollout returns a LitAgent instance (actually LitAgentLLM for this pattern)."""
assert isinstance(sample_rollout_func, LitAgentLLM)
# Should have agent methods
assert hasattr(sample_rollout_func, "rollout")
assert hasattr(sample_rollout_func, "rollout_async")
assert hasattr(sample_rollout_func, "training_rollout")
def test_rollout_preserves_signature():
"""Test that @rollout preserves function signature."""
sig = inspect.signature(sample_rollout_func)
params = list(sig.parameters.keys())
# Should have the expected parameters
assert params == ["task", "llm"]
@pytest.mark.asyncio
async def test_async_function_with_llm_rollout():
"""Test that async functions work with @llm_rollout decorator."""
@llm_rollout
async def async_agent(task, llm):
"""An async test function."""
return f"Async processed: {task} with {llm}"
# Should be callable
assert callable(async_agent)
# Should preserve async nature when called directly
result = await async_agent("test", "llm")
assert result == "Async processed: test with llm"
# Should be marked as async
assert async_agent.is_async
@pytest.mark.asyncio
async def test_async_function_with_rollout():
"""Test that async functions work with @rollout decorator."""
@rollout
async def async_agent(task, llm):
"""An async test function."""
return f"Async processed: {task} with {llm}"
# Should be callable
assert callable(async_agent)
# Should preserve async nature when called directly
result = await async_agent("test", "llm")
assert result == "Async processed: test with llm"
# Should be marked as async
assert async_agent.is_async
tests/test_litagent_hooks.py
+2 -2
View File
@@ -61,10 +61,10 @@ class HookAgent(LitAgent):
self.end_called = False
self.end_rollout = None
def training_rollout(self, task, rollout_id, resources):
def training_rollout(self, task, resources, rollout):
return 0.5
async def training_rollout_async(self, task, rollout_id, resources):
async def training_rollout_async(self, task, resources, rollout):
return 0.5
def on_rollout_start(self, task, runner, tracer):