# Copyright (c) Microsoft. All rights reserved. """Agent runner implementation for executing agent rollouts. This module provides the concrete implementation of the runner interface, handling the execution of agent rollouts with support for tracing, hooks, and distributed worker coordination. """ from __future__ import annotations import asyncio import logging import random import threading import time from contextlib import suppress from typing import ( TYPE_CHECKING, Any, Awaitable, Callable, List, Literal, Optional, Sequence, TypeVar, cast, ) from opentelemetry.sdk.trace import ReadableSpan from agentlightning.litagent import LitAgent from agentlightning.reward import emit_reward, find_final_reward from agentlightning.store.base import LightningStore from agentlightning.tracer.base import Tracer from agentlightning.tracer.otel import OtelTracer from agentlightning.types import ( AttemptedRollout, Hook, NamedResources, Rollout, RolloutMode, RolloutRawResult, Span, SpanCoreFields, ) from agentlightning.utils.system_snapshot import system_snapshot if TYPE_CHECKING: from agentlightning.execution.events import ExecutionEvent from .base import Runner T_task = TypeVar("T_task") logger = logging.getLogger(__name__) class LitAgentRunner(Runner[T_task]): """Execute [`LitAgent`][agentlightning.LitAgent] tasks with tracing support. This runner manages the complete lifecycle of agent rollout execution, including task polling, resource management, tracing, and hooks. It supports both continuous iteration over tasks from the store and single-step execution. Attributes: worker_id: Identifier for the active worker process, if any. """ def __init__( self, tracer: Tracer, max_rollouts: Optional[int] = None, poll_interval: float = 5.0, heartbeat_interval: float = 10.0, interval_jitter: float = 0.5, heartbeat_launch_mode: Literal["asyncio", "thread"] = "thread", heartbeat_include_gpu: bool = False, ) -> None: """Initialize the agent runner. Args: tracer: [`Tracer`][agentlightning.Tracer] used for rollout spans. max_rollouts: Optional cap on iterations processed by [`iter`][agentlightning.LitAgentRunner.iter]. poll_interval: Seconds to wait between store polls when no work is available. heartbeat_interval: Seconds to wait between sending heartbeats to the store. interval_jitter: Jitter factor for the poll interval. The actual interval will be between poll_interval - interval_jitter and poll_interval + interval_jitter. This is to avoid the overload caused by the synchronization of the runners. heartbeat_launch_mode: Launch mode for the heartbeat loop. Can be "asyncio" or "thread". "thread" is the default and recommended mode as it prevents blocking the event loop under load. Use "asyncio" for simpler deployments with low worker counts. heartbeat_include_gpu: Whether to include GPU stats in heartbeat snapshots. Querying GPU stats can be slow under load, so this is disabled by default. """ super().__init__() self._tracer = tracer self._max_rollouts = max_rollouts self._poll_interval = poll_interval self._heartbeat_interval = heartbeat_interval self._interval_jitter = interval_jitter self._heartbeat_launch_mode = heartbeat_launch_mode self._heartbeat_include_gpu = heartbeat_include_gpu self._random_state = random.Random() # Set later self._agent: Optional[LitAgent[T_task]] = None self._hooks: Sequence[Hook] = [] self._store: Optional[LightningStore] = None self.worker_id: Optional[int] = None def init(self, agent: LitAgent[T_task], *, hooks: Optional[Sequence[Hook]] = None, **kwargs: Any) -> None: """Initialize the runner with the agent. This sets up the agent-runner relationship, registers hooks, and initializes the tracer. Args: agent: [`LitAgent`][agentlightning.LitAgent] instance executed by the runner. hooks: Optional sequence of [`Hook`][agentlightning.Hook] callbacks invoked around tracing and rollout boundaries. **kwargs: Additional initialization arguments (currently unused). """ self._agent = agent self._agent.set_runner(self) self._hooks = [*hooks] if hooks is not None else [] self._tracer.init() def init_worker(self, worker_id: int, store: LightningStore, **kwargs: Any) -> None: """Initialize the runner for each worker with worker_id and store. This method is called once per worker in a distributed setup to provide the worker with its ID and store connection. Args: worker_id: Unique identifier for this worker process. store: [`LightningStore`][agentlightning.LightningStore] used for task coordination and persistence. **kwargs: Additional worker-specific initialization arguments (currently unused). """ self._store = store self.worker_id = worker_id self._tracer.init_worker(worker_id, store) def teardown(self, *args: Any, **kwargs: Any) -> None: """Teardown the runner and clean up all resources. This method resets all internal state including the agent, store, hooks, and worker ID, and calls the tracer's teardown method. Args: *args: Additional teardown arguments (currently unused). **kwargs: Additional teardown keyword arguments (currently unused). """ self._agent = None self._store = None self.worker_id = None self._hooks = [] self._tracer.teardown() def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None: """Teardown the runner for a specific worker. This method cleans up worker-specific resources and resets the worker ID. Args: worker_id: Unique identifier of the worker being torn down. *args: Additional teardown arguments (currently unused). **kwargs: Additional teardown keyword arguments (currently unused). """ self.worker_id = None self._tracer.teardown_worker(worker_id) @property def tracer(self) -> Tracer: """Get the tracer instance. Returns: The Tracer instance used by this runner. """ return self._tracer def get_agent(self) -> LitAgent[T_task]: """Get the agent instance. Returns: The LitAgent instance managed by this runner. Raises: ValueError: If the agent has not been initialized via [`init`][agentlightning.LitAgentRunner.init]. """ if self._agent is None: raise ValueError("Agent not initialized. Call init() first.") return self._agent def get_store(self) -> LightningStore: """Get the store instance. Returns: The LightningStore instance for this worker. Raises: ValueError: If the store has not been initialized via [`init_worker`][agentlightning.LitAgentRunner.init_worker]. """ if self._store is None: raise ValueError("Store not initialized. Call init_worker() first.") return self._store def get_worker_id(self) -> str: """Get the formatted worker ID string. Returns: A formatted string like "Worker-0" if initialized, or "Worker-Unknown" if the worker ID has not been set. """ return f"Worker-{self.worker_id}" if self.worker_id is not None else "Worker-Unknown" def _log_prefix(self, rollout_id: Optional[str] = None) -> str: """Generate a standardized log prefix for the current worker. This creates a consistent prefix format for log messages to identify which worker and rollout the message is associated with. Args: rollout_id: Optional rollout ID to include in the prefix. Returns: A formatted log prefix string like "[Worker 0 | Rollout xyz]", "[Worker 0]", "[Rollout xyz]", or "[Default Worker]". """ if self.worker_id is not None: if rollout_id: return f"[Worker {self.worker_id} | Rollout {rollout_id}]" else: return f"[Worker {self.worker_id}]" if rollout_id: return f"[Rollout {rollout_id}]" return "[Default Worker]" async def _trigger_hooks( self, hook_type: Literal["on_trace_start", "on_trace_end", "on_rollout_start", "on_rollout_end"], *args: Any, **kwargs: Any, ) -> None: """Trigger all registered hooks of a specific type. This method calls the specified hook method on all registered hooks, catching and logging any exceptions that occur during hook execution to prevent them from disrupting the main execution flow. Args: hook_type: The type of hook to trigger. Valid values are: "on_trace_start", "on_trace_end", "on_rollout_start", "on_rollout_end". *args: Positional arguments to pass to the hook methods. **kwargs: Keyword arguments to pass to the hook methods. """ for hook in self._hooks: try: await getattr(hook, hook_type)(*args, **kwargs) except Exception: logger.exception(f"{self._log_prefix()} Exception during {hook_type} hook {hook}.") async def _post_process_rollout_result( self, rollout: AttemptedRollout, raw_result: RolloutRawResult ) -> List[ReadableSpan] | List[Span]: """Standardizes the agent's return value and report what's needed to report to the store. Args: rollout: The rollout object for the current task. raw_result: The output from the agent's rollout method. Returns: The spans that are assumed to be added to the store. This only serves as an estimation for logging purposes. For precise tracking, use the store directly. """ store = self.get_store() trace_spans: list[Span] = [] result_recognized: bool = False # Case 0: result is None if raw_result is None: trace_spans = self._tracer.get_last_trace() result_recognized = True # Case 1: result is a float (final reward) if isinstance(raw_result, (bool, int, float)): if isinstance(raw_result, (bool, int)): logger.warning( f"{self._log_prefix(rollout.rollout_id)} Reward is not a number, got: {type(raw_result)}. " "Auto converting to float." ) raw_result = float(raw_result) # Preserve the existing spans before another span is emitted trace_spans = list(self._tracer.get_last_trace()) # This will NOT emit another span to the tracer reward_span_core_fields = emit_reward(raw_result, propagate=False) # We add it to the store manually sequence_id = await store.get_next_span_sequence_id(rollout.rollout_id, rollout.attempt.attempt_id) reward_span = Span.from_core_fields( reward_span_core_fields, rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id, sequence_id=sequence_id, ) await store.add_span(reward_span) result_recognized = True # Case 2-4: result is a list if isinstance(raw_result, list): # Case 2: result is a list of ReadableSpan (OpenTelemetry spans) if len(raw_result) > 0 and all(isinstance(t, ReadableSpan) for t in raw_result): if isinstance(self._tracer, OtelTracer): logger.warning( f"{self._log_prefix(rollout.rollout_id)} Tracer is already an OpenTelemetry tracer. " "The traces should have already been added to the store. " "Returning the traces from the rollout will result in duplicate spans." ) for span in raw_result: added_span = await store.add_otel_span( rollout.rollout_id, rollout.attempt.attempt_id, cast(ReadableSpan, span) ) if added_span is not None: trace_spans.append(added_span) else: logger.error( f"{self._log_prefix(rollout.rollout_id)} Failed to add OpenTelemetry span to the store: {span}" ) result_recognized = True # Case 3: result is a list of Span (agentlightning spans) elif len(raw_result) > 0 and all(isinstance(t, Span) for t in raw_result): # Add the spans directly to the store for span in raw_result: await store.add_span(cast(Span, span)) trace_spans = [cast(Span, span) for span in raw_result] result_recognized = True # Case 4: result is a list of SpanCoreFields (agentlightning spans) elif len(raw_result) > 0 and all(isinstance(t, SpanCoreFields) for t in raw_result): # Add the spans directly to the store too, but needs to get sequence id first sequence_ids = await store.get_many_span_sequence_ids( [(rollout.rollout_id, rollout.attempt.attempt_id) for _ in range(len(raw_result))] ) trace_spans = [ Span.from_core_fields( cast(SpanCoreFields, span_core_fields), rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id, sequence_id=sequence_id, ) for span_core_fields, sequence_id in zip(raw_result, sequence_ids, strict=True) ] await store.add_many_spans(trace_spans) result_recognized = True # Left over cases for list elif len(raw_result) == 0: logger.warning( f"{self._log_prefix(rollout.rollout_id)} The rollout returns an empty list. " "Please check your rollout implementation." ) trace_spans = [] result_recognized = True else: types = [type(t).__name__ for t in raw_result][:10] raise ValueError( f"Invalid raw result type. It's expected to be a list of ReadableSpan or Span, " f"but got: {', '.join(types)}..." ) if not result_recognized: raise TypeError( f"Invalid raw result type. It's expected to be none, float, or a list of ReadableSpan or Span, " f"but got: {type(raw_result).__name__}..." ) return trace_spans async def _emit_heartbeat(self, store: LightningStore) -> None: """Send a heartbeat tick to the store. Args: store: The lightning store to update. """ logger.debug(f"{self._log_prefix()} Preparing to emit heartbeat.") worker_id = self.get_worker_id() try: snapshot = await asyncio.wait_for( asyncio.to_thread(system_snapshot, self._heartbeat_include_gpu), timeout=self._heartbeat_interval, ) logger.debug(f"{self._log_prefix()} Heartbeat snapshot acquired.") except asyncio.TimeoutError: logger.warning( "%s Heartbeat snapshot acquisition timed out after %.1fs, skipping.", self._log_prefix(), self._heartbeat_interval, ) return except asyncio.CancelledError: # bypass the exception raise except Exception: logger.exception("%s Unable to acquire heartbeat snapshot.", self._log_prefix()) return try: await asyncio.wait_for(store.update_worker(worker_id, snapshot), timeout=self._heartbeat_interval) logger.debug(f"{self._log_prefix()} Heartbeat updated successfully.") except asyncio.CancelledError: # bypass the exception raise except asyncio.TimeoutError: logger.warning( "%s update worker heartbeat timed out after %.1fs, skipping.", self._log_prefix(), self._heartbeat_interval, ) except Exception: logger.exception("%s Unable to update worker heartbeat.", self._log_prefix()) def _start_heartbeat_loop(self, store: LightningStore) -> Optional[Callable[[], Awaitable[None]]]: """Start a background heartbeat loop and return an async stopper.""" if self._heartbeat_interval <= 0: return None if self.worker_id is None: logger.warning("%s Cannot start heartbeat loop without worker_id.", self._log_prefix()) return None if self._heartbeat_launch_mode == "asyncio": return self._start_heartbeat_asyncio_loop(store) if self._heartbeat_launch_mode == "thread": return self._start_heartbeat_thread_loop(store) raise ValueError(f"Unsupported heartbeat launch mode: {self._heartbeat_launch_mode}") def _start_heartbeat_asyncio_loop(self, store: LightningStore) -> Optional[Callable[[], Awaitable[None]]]: """Start a background heartbeat loop using asyncio. Args: store: The lightning store to update. Returns: An async stopper function that can be used to stop the heartbeat loop. """ stop_event = asyncio.Event() async def heartbeat_loop() -> None: while not stop_event.is_set(): try: # Run _emit_heartbeat in thread pool to avoid blocking the event loop. # Timeout at the interval - if it takes longer, the data is stale anyway. await self._emit_heartbeat(store) except Exception: logger.exception("%s Heartbeat failed.", self._log_prefix()) with suppress(asyncio.TimeoutError): interval = self._heartbeat_interval + self._random_state.uniform( -self._interval_jitter, self._interval_jitter ) interval = max(interval, 0.01) await asyncio.wait_for(stop_event.wait(), timeout=interval) task = asyncio.create_task(heartbeat_loop(), name=f"{self.get_worker_id()}-heartbeat") async def stop() -> None: stop_event.set() with suppress(asyncio.CancelledError): await task return stop def _start_heartbeat_thread_loop(self, store: LightningStore) -> Optional[Callable[[], Awaitable[None]]]: """Start a background heartbeat loop using threading. It uses two threads: one to produce the snapshot and one to consume it, to avoid either of them blocking the event loop. Args: store: The lightning store to update. Returns: An async stopper function that can be used to stop the heartbeat loop. """ stop_evt = threading.Event() lock = threading.Lock() latest_snapshot = None latest_ts = 0.0 # time.monotonic() when snapshot was captured # Consider snapshot stale after ~1 interval plus jitter slack. stale_after = self._heartbeat_interval + self._interval_jitter + 1.0 worker_id = self.get_worker_id() def producer() -> None: nonlocal latest_snapshot, latest_ts while not stop_evt.is_set(): try: logger.debug(f"{self._log_prefix()} Heartbeat producer: acquiring snapshot.") snap = system_snapshot(self._heartbeat_include_gpu) # sync logger.debug(f"{self._log_prefix()} Heartbeat producer: snapshot acquired.") ts = time.monotonic() with lock: latest_snapshot = snap latest_ts = ts except Exception: logger.warning("%s Heartbeat producer: system_snapshot failed.", self._log_prefix(), exc_info=True) interval = self._heartbeat_interval + self._random_state.uniform( -self._interval_jitter, self._interval_jitter ) stop_evt.wait(max(interval, 0.01)) def consumer() -> None: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) last_warned_ts = None # Track which snapshot we've already warned about try: while not stop_evt.is_set(): with lock: snap = latest_snapshot ts = latest_ts wait_interval = max( self._heartbeat_interval + self._random_state.uniform(-self._interval_jitter, self._interval_jitter), 0.01, ) if snap is None: # probably just started logger.debug("%s Heartbeat consumer: no snapshot yet; skipping update.", self._log_prefix()) stop_evt.wait(wait_interval) continue age = time.monotonic() - ts if age > stale_after: # Only warn once per stale snapshot (check if we haven't warned about this timestamp yet) if last_warned_ts != ts: logger.warning( "%s Heartbeat consumer: snapshot stale (age=%.2fs > %.2fs); skipping update.", self._log_prefix(), age, stale_after, ) last_warned_ts = ts stop_evt.wait(wait_interval) continue try: logger.debug(f"{self._log_prefix()} Heartbeat consumer: updating worker.") loop.run_until_complete( asyncio.wait_for( store.update_worker(worker_id, snap), timeout=self._heartbeat_interval, ) ) logger.debug(f"{self._log_prefix()} Heartbeat consumer: worker updated.") except asyncio.TimeoutError: logger.warning( "%s Heartbeat consumer: update timed out after %.1fs.", self._log_prefix(), self._heartbeat_interval, ) except Exception: logger.warning("%s Heartbeat consumer: update failed.", self._log_prefix(), exc_info=True) stop_evt.wait(wait_interval) finally: with suppress(Exception): loop.stop() with suppress(Exception): loop.close() t_prod = threading.Thread(target=producer, name=f"{worker_id}-heartbeat-producer", daemon=True) t_cons = threading.Thread(target=consumer, name=f"{worker_id}-heartbeat-consumer", daemon=True) t_prod.start() t_cons.start() async def stop() -> None: stop_evt.set() await asyncio.to_thread(t_prod.join) await asyncio.to_thread(t_cons.join) return stop async def _sleep_until_next_poll(self, event: Optional[ExecutionEvent] = None) -> None: """Sleep until the next poll interval, with optional event-based interruption. If an event is provided, the method will check it periodically (every 0.1s) and return early if the event is set. Args: event: Optional [`ExecutionEvent`][agentlightning.ExecutionEvent] object that can be used to interrupt the sleep. If set during the sleep period, the method returns immediately. """ interval = self._poll_interval + self._random_state.uniform(-self._interval_jitter, self._interval_jitter) interval = max(interval, 0.01) if event is None: await asyncio.sleep(interval) return current_time = time.time() next_time = current_time + interval while time.time() < next_time: await asyncio.sleep(0.1) if event.is_set(): return async def _step_impl(self, next_rollout: AttemptedRollout, raise_on_exception: bool = False) -> str: """Execute a single rollout implementation. This is the core method that handles the execution of a single rollout, including resource fetching, hook triggering, agent invocation, tracing, and result processing. Args: next_rollout: The rollout to execute, containing input data, mode, and resources information. raise_on_exception: If True, exceptions during rollout execution will be re-raised. If False, exceptions are logged but not propagated. """ store = self.get_store() agent = self.get_agent() rollout_id = next_rollout.rollout_id resources_id = next_rollout.resources_id resources_update = None if resources_id: resources_update = await store.get_resources_by_id(resources_id) else: logger.debug(f"{self._log_prefix(rollout_id)} No 'resources_id'. Fetching latest resources.") resources_update = await store.get_latest_resources() if not resources_update: if raise_on_exception: raise RuntimeError(f"{self._log_prefix(rollout_id)} Failed to fetch resources") else: logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.") return rollout_id logger.debug(f"{self._log_prefix(rollout_id)} Resources fetched (id={resources_update.resources_id}).") trace_spans: List[ReadableSpan] | List[Span] = [] has_exception: bool = False try: await self._trigger_hooks(hook_type="on_rollout_start", agent=agent, runner=self, rollout=next_rollout) start_time = time.time() logger.debug(f"{self._log_prefix(rollout_id)} Prepared for trace context.") async with self._tracer.trace_context( name=rollout_id, rollout_id=rollout_id, attempt_id=next_rollout.attempt.attempt_id ): logger.debug(f"{self._log_prefix(rollout_id)} Entered trace context.") await self._trigger_hooks( hook_type="on_trace_start", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout ) # NOTE: This is the most costly step in the whole function # If the rollout method becomes unresponsive or timeouts, there is nothing we can do within the runner. # We might need some mechanisms in execution strategy to restart the runner. But that's a future work. if agent.is_async(): rollout_method = ( agent.training_rollout_async if next_rollout.mode == "train" else agent.validation_rollout_async ) logger.debug(f"{self._log_prefix(rollout_id)} Starting async rollout method.") result = await rollout_method( next_rollout.input, resources=resources_update.resources, rollout=next_rollout ) logger.debug(f"{self._log_prefix(rollout_id)} Async rollout method completed.") else: rollout_method = ( agent.training_rollout if next_rollout.mode == "train" else agent.validation_rollout ) logger.debug(f"{self._log_prefix(rollout_id)} Starting sync rollout method.") result = rollout_method( next_rollout.input, resources=resources_update.resources, rollout=next_rollout ) logger.debug(f"{self._log_prefix(rollout_id)} Sync rollout method completed.") await self._trigger_hooks( hook_type="on_trace_end", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout ) logger.debug(f"{self._log_prefix(rollout_id)} Trace context exited.") # Possible exceptions in post_process will be caught in the overall exception handler trace_spans = await self._post_process_rollout_result(next_rollout, result) last_reward = find_final_reward(trace_spans) end_time = time.time() logger.info( f"{self._log_prefix(rollout_id)} Completed in " f"{end_time - start_time:.2f}s. Collected {len(trace_spans)} span(s). " f"Final reward: {last_reward}" ) except Exception: logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.") has_exception = True if raise_on_exception: raise finally: try: await self._trigger_hooks( hook_type="on_rollout_end", agent=agent, runner=self, rollout=next_rollout, spans=trace_spans ) except Exception: logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.") try: if has_exception: # possibly timed out and cancelled? await store.update_attempt(rollout_id, next_rollout.attempt.attempt_id, status="failed") else: await store.update_attempt(rollout_id, next_rollout.attempt.attempt_id, status="succeeded") except Exception: logger.exception( f"{self._log_prefix(rollout_id)} Exception during update_attempt. Giving up the update." ) return rollout_id async def iter(self, *, event: Optional[ExecutionEvent] = None) -> None: """Run the runner, continuously iterating over tasks in the store. This method polls the store for new rollouts and executes them until: - The event is set (if provided) - The max_rollouts limit is reached (if configured) - No more tasks are available All exceptions during rollout execution are caught and logged but not propagated, allowing the runner to continue processing subsequent tasks. Args: event: Optional ExecutionEvent object to signal the runner to stop. The runner will check this event periodically and stop gracefully when set. """ num_tasks_processed = 0 logger.info(f"{self._log_prefix()} Started async rollouts (max: {self._max_rollouts or 'unlimited'}).") store = self.get_store() stop_heartbeat = self._start_heartbeat_loop(store) try: while not (event is not None and event.is_set()) and ( self._max_rollouts is None or num_tasks_processed < self._max_rollouts ): # Retrieve the next rollout next_rollout: Optional[Rollout] = None while not (event is not None and event.is_set()): logger.debug(f"{self._log_prefix()} Try to poll for next rollout.") next_rollout = await store.dequeue_rollout(worker_id=self.get_worker_id()) logger.debug(f"{self._log_prefix()} Next rollout retrieved: {next_rollout}") if next_rollout is None: logger.debug( f"{self._log_prefix()} No rollout to poll. Waiting for {self._poll_interval} seconds." ) await self._sleep_until_next_poll(event) else: break if next_rollout is None: return # Execute the step await self._step_impl(next_rollout) num_tasks_processed += 1 if num_tasks_processed % 10 == 0 or num_tasks_processed == 1: logger.info( f"{self._log_prefix()} Progress: {num_tasks_processed}/{self._max_rollouts or 'unlimited'}" ) finally: if stop_heartbeat is not None: await stop_heartbeat() logger.info(f"{self._log_prefix()} Finished async rollouts. Processed {num_tasks_processed} tasks.") async def step( self, input: T_task, *, resources: Optional[NamedResources] = None, mode: Optional[RolloutMode] = None, event: Optional[ExecutionEvent] = None, ) -> Rollout: """Execute a single task directly, bypassing the task queue. This method creates a new rollout for the given input and executes it immediately. Unlike [`iter()`][agentlightning.LitAgentRunner.iter], exceptions are propagated to the caller. Args: input: The task input to be processed by the agent. resources: Optional named resources to be used for this specific task. If provided, a new resources entry will be created in the store. If not provided, the latest resources from the store will be used. mode: Optional rollout mode ("train" or "validation"). If not provided, the agent's default mode will be used. event: Optional ExecutionEvent object to signal interruption (currently unused but included for interface consistency). Returns: The completed rollout. Raises: Exception: Any exception that occurs during rollout execution will be re-raised to the caller. """ store = self.get_store() if resources is not None: resources_update = await store.add_resources(resources) resources_id = resources_update.resources_id else: resources_id = None attempted_rollout = await self.get_store().start_rollout( input=input, mode=mode, resources_id=resources_id, worker_id=self.get_worker_id() ) rollout_id = await self._step_impl(attempted_rollout, raise_on_exception=True) completed_rollout = await store.get_rollout_by_id(rollout_id) if completed_rollout is None: raise RuntimeError(f"{self._log_prefix()} Failed to fetch completed rollout by id after step: {rollout_id}") return completed_rollout