# Copyright (c) Microsoft. All rights reserved. from __future__ import annotations import asyncio import logging from datetime import datetime from typing import TYPE_CHECKING, Any, List, Literal, Optional from agentlightning.types import Attempt, Dataset, Rollout, RolloutStatus, Span from .base import Algorithm from .utils import with_llm_proxy, with_store if TYPE_CHECKING: from agentlightning.llm_proxy import LLMProxy from agentlightning.store.base import LightningStore logger = logging.getLogger(__name__) __all__ = ["FastAlgorithm", "Baseline"] class FastAlgorithm(Algorithm): """Base class for lightweight algorithms optimised for developer workflows. Fast algorithms prioritise short feedback loops so an agent developer can run small-scale experiments without waiting for long-running training jobs to finish. """ def _timestamp_to_iso_str(timestamp: float) -> str: return datetime.fromtimestamp(timestamp).isoformat() class Baseline(FastAlgorithm): """Reference implementation that streams the full dataset through the rollout queue. The baseline algorithm batches task submissions, waits for each rollout to finish, and logs every collected span and reward. It is primarily useful as a smoke test for the platform plumbing rather than a performant trainer. The baseline algorithm will auto-start a LLM proxy if one is provided and not yet started. Args: n_epochs: Number of dataset passes to execute for both the train and val splits during developer experiments. train_split: Fraction of the concatenated dataset to treat as training data. Must be strictly between 0 and 1. polling_interval: Interval, in seconds, to poll the store for queue depth and rollout completion. max_queue_length: Number of rollouts allowed to wait in the queue before throttling additional submissions. span_verbosity: Level of detail to include when logging span metadata. Raises: ValueError: If `train_split` falls outside the `(0, 1)` interval. Examples: ```python from agentlightning.algorithm.fast import Baseline algorithm = Baseline(n_epochs=2, train_split=0.8, span_verbosity="key_values") trainer.fit(algorithm, train_dataset=my_train, val_dataset=my_val) ``` """ def __init__( self, *, n_epochs: int = 1, train_split: float = 0.5, polling_interval: float = 5.0, max_queue_length: int = 4, span_verbosity: Literal["keys", "key_values", "none"] = "keys", ) -> None: super().__init__() self.n_epochs = n_epochs self.train_split = train_split self.polling_interval = polling_interval self.max_queue_length = max_queue_length self.span_verbosity = span_verbosity if not (0.0 < self.train_split < 1.0): raise ValueError("train_split must be between 0 and 1.") self._finished_rollout_count = 0 def _span_to_string(self, rollout_id: str, attempt: Attempt, span: Span) -> str: """Format a span for logging based on the configured verbosity.""" if self.span_verbosity == "none": return "" prefix_msg = f"[Rollout {rollout_id} | Attempt {attempt.attempt_id} | Span {span.span_id}] #{span.sequence_id} ({span.name}) " elapsed = f"{span.end_time - span.start_time:.2f}" if span.start_time and span.end_time else "unknown" msg = ( prefix_msg + f"From {_timestamp_to_iso_str(span.start_time) if span.start_time else 'unknown'}, " + f"to {_timestamp_to_iso_str(span.end_time) if span.end_time else 'unknown'}, " + f"{elapsed} seconds. " ) if self.span_verbosity == "key_values": msg += f"Attributes: {span.attributes}" else: msg += f"Attribute keys: {list(span.attributes.keys())}" return msg async def _handle_rollout_finish(self, rollout: Rollout) -> None: """Log attempt metadata and emit adapted traces when a rollout ends.""" store = self.get_store() rollout_id = rollout.rollout_id rollout_end_time = rollout.end_time or asyncio.get_event_loop().time() logger.info( f"[Rollout {rollout_id}] Finished with status {rollout.status} in {rollout_end_time - rollout.start_time:.2f} seconds." ) # Logs all the attempts and their corresponding spans attempts = await store.query_attempts(rollout_id) for attempt in attempts: logger.info( "[Rollout %s | Attempt %s] ID: %s. Status: %s. Worker: %s", rollout_id, attempt.sequence_id, attempt.attempt_id, attempt.status, attempt.worker_id, ) spans = await store.query_spans(rollout_id=rollout_id) for span in spans: if self.span_verbosity != "none": logger.info(self._span_to_string(rollout.rollout_id, attempt, span)) # Attempts to adapt the spans using the adapter if provided try: adapter = self.get_adapter() except ValueError: logger.warning("No adapter set for MockAlgorithm. Skipping trace adaptation.") adapter = None if adapter is not None: spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest") transformed_data = adapter.adapt(spans) logger.info(f"[Rollout {rollout_id}] Adapted data: {transformed_data}") async def _enqueue_rollouts( self, dataset: Dataset[Any], train_indices: List[int], val_indices: List[int], resources_id: str ) -> None: """Submit rollouts while respecting the maximum queue length.""" store = self.get_store() for index in train_indices + val_indices: queuing_rollouts = await store.query_rollouts(status_in=["queuing", "requeuing"]) if len(queuing_rollouts) <= 1: # Only enqueue a new rollout when there is at most 1 rollout in the queue. sample = dataset[index] mode = "train" if index in train_indices else "val" rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id) logger.info(f"[Rollout {rollout.rollout_id}] Enqueued in {mode} mode with sample: {sample}") await asyncio.sleep(self.polling_interval) async def _harvest_rollout_spans(self, rollout_id: str): """Poll rollout status updates until completion and log transitions.""" store = self.get_store() last_status: Optional[RolloutStatus] = None while True: rollout = await store.get_rollout_by_id(rollout_id) if rollout is not None: if rollout.status in ["succeeded", "failed", "cancelled"]: # Rollout is finished, log all the data. await self._handle_rollout_finish(rollout) # We are done here. self._finished_rollout_count += 1 logger.info(f"Finished {self._finished_rollout_count} rollouts.") break if last_status != rollout.status: if last_status is not None: logger.info(f"[Rollout {rollout_id}] Status changed to {rollout.status}.") else: logger.info(f"[Rollout {rollout_id}] Status is initialized to {rollout.status}.") last_status = rollout.status else: logger.debug(f"[Rollout {rollout_id}] Status is still {rollout.status}.") await asyncio.sleep(self.polling_interval) @with_llm_proxy() @with_store async def run( self, store: LightningStore, # Injected by decorator - callers should not provide this parameter llm_proxy: Optional[LLMProxy], # Injected by decorator - callers should not provide this parameter train_dataset: Optional[Dataset[Any]] = None, val_dataset: Optional[Dataset[Any]] = None, ) -> None: """Execute the baseline loop across the provided datasets.""" train_dataset_length = len(train_dataset) if train_dataset is not None else 0 val_dataset_length = len(val_dataset) if val_dataset is not None else 0 if train_dataset_length == 0 and val_dataset_length == 0: logger.error( "MockAlgorithm requires at least one dataset. Provide train_dataset or val_dataset before running." ) return concatenated_dataset = [train_dataset[i] for i in range(train_dataset_length) if train_dataset is not None] + [ val_dataset[i] for i in range(val_dataset_length) if val_dataset is not None ] train_indices = list(range(0, train_dataset_length)) val_indices = list(range(train_dataset_length, train_dataset_length + val_dataset_length)) logger.debug(f"Train indices: {train_indices}") logger.debug(f"Val indices: {val_indices}") # Currently we only supports a single resource update at the start. initial_resources = self.get_initial_resources() if initial_resources is not None: resource_update = await store.update_resources("default", initial_resources) resources_id = resource_update.resources_id logger.info(f"Initial resources set: {initial_resources}") else: logger.warning("No initial resources provided. Skip initializing resources.") resources_id = None for epoch in range(self.n_epochs): harvest_tasks: List[asyncio.Task[None]] = [] logger.info(f"Proceeding epoch {epoch + 1}/{self.n_epochs}.") for index in train_indices + val_indices: logger.info( f"Processing index {index}. {len(train_indices)} train indices and {len(val_indices)} val indices in total." ) while True: queuing_rollouts = await store.query_rollouts(status_in=["queuing", "requeuing"]) if len(queuing_rollouts) <= self.max_queue_length: # Only enqueue a new rollout when there is at most "max_queue_length" rollout in the queue. sample = concatenated_dataset[index] mode = "train" if index in train_indices else "val" rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id) harvest_tasks.append(asyncio.create_task(self._harvest_rollout_spans(rollout.rollout_id))) logger.info(f"Enqueued rollout {rollout.rollout_id} in {mode} mode with sample: {sample}") break else: # Sleep a bit and try again later. await asyncio.sleep(self.polling_interval) # Wait for all harvest tasks to complete logger.info(f"Waiting for {len(harvest_tasks)} harvest tasks to complete...") if len(harvest_tasks) > 0: await asyncio.gather(*harvest_tasks)