Rollout Architecture Overview

Relevant source files

The following files were used as context for generating this wiki page:

• osmosis_ai/rollout/__init__.py
• osmosis_ai/rollout/core/base.py
• osmosis_ai/rollout/core/schemas.py
• osmosis_ai/rollout/server/app.py
• osmosis_ai/rollout/config/settings.py
• osmosis_ai/rollout/validator.py
• osmosis_ai/rollout/core/llm_client.py

The Osmosis Remote Rollout SDK provides a framework for integrating autonomous agent loops with the Osmosis training ecosystem. Its primary purpose is to facilitate the execution of multi-turn agent trajectories, allowing agent frameworks to communicate with "TrainGate" for LLM completions and result reporting.

The architecture is built around a provider-agnostic protocol where users implement a specific agent logic in a RolloutAgentLoop, which is then hosted via a FastAPI-based RolloutServer. The SDK manages the complexities of asynchronous communication, idempotency, metrics collection, and protocol validation.
Sources: osmosis_ai/rollout/init.py:1-20, osmosis_ai/rollout/core/base.py:1-30

Core Components and Classes

The system is composed of several architectural layers, ranging from core abstract base classes to specialized Pydantic schemas for protocol enforcement.

RolloutAgentLoop

The RolloutAgentLoop is the primary extension point for developers. It is an abstract base class that requires the implementation of get_tools and an asynchronous run method.
Sources: osmosis_ai/rollout/core/base.py:254-340

RolloutContext

The RolloutContext object is passed to the agent's run method. It acts as an execution sandbox, providing access to the initial RolloutRequest, the list of available tools, and a pre-configured LLMClientProtocol for making chat completion calls.
Sources: osmosis_ai/rollout/core/base.py:84-251

RolloutResult

The final state of an agent's trajectory is encapsulated in a RolloutResult. This object includes the final conversation history, a status indicator (COMPLETED or ERROR), and optional metadata such as rewards or performance metrics.
Sources: osmosis_ai/rollout/core/base.py:44-81

Class Hierarchy Diagram

The following diagram illustrates the relationship between the core execution components:

Lifecycle and Protocol Flow

The rollout process follows a specific sequence of events initiated by a "TrainGate" server.

1. Initialization: TrainGate sends a RolloutRequest to the server's /v1/rollout/init endpoint.
2. Tool Discovery: The server calls the agent's get_tools() method and returns an InitResponse (202 Accepted) immediately.
3. Background Execution: The server spawns a background task to execute the run() method of the RolloutAgentLoop.
4. LLM Interaction: During execution, the agent uses ctx.chat() to request completions from TrainGate.
5. Completion: Once the loop finishes, the server POSTs a RolloutResponse back to the TrainGate's /v1/rollout/completed endpoint.

Sequence Diagram: Rollout Execution

Server and API Architecture

The RolloutServer is implemented using FastAPI and includes built-in state management for handling concurrency and idempotency.

Endpoint Summary

Endpoint	Method	Description	Auth Required
/v1/rollout/init	POST	Starts a new rollout session or returns existing session tools.	Optional API Key
/health	GET	Basic health check returning active and completed rollout counts.	No
/platform/health	GET	Authenticated health check for Osmosis Platform validation.	Bearer Token
Sources: osmosis_ai/rollout/server/app.py:192-275

Concurrency and Idempotency

The server utilizes an AppState object to track active tasks. Idempotency is handled by caching the InitResponse for a given rollout_id or idempotency_key. If a duplicate request is received while a rollout is running, the server awaits the cached future instead of spawning a new task.
Sources: osmosis_ai/rollout/server/app.py:277-310

Configuration and Settings

The SDK uses Pydantic-based settings that can be configured via environment variables with specific prefixes.

Category	Prefix	Key Settings	Default
Server	OSMOSIS_ROLLOUT_SERVER_	max_concurrent_rollouts	100
		record_ttl_seconds	3600.0
Client	OSMOSIS_ROLLOUT_CLIENT_	timeout_seconds	300.0
		max_retries	3
Global	OSMOSIS_ROLLOUT_	max_metadata_size_bytes	1MB
Sources: osmosis_ai/rollout/config/settings.py:46-155

Validation Framework

To ensure protocol compliance, the SDK includes a validator module. It performs static and dynamic checks on RolloutAgentLoop implementations before they are served.

Validation Checks

• Name Verification: Ensures the name attribute is a non-empty string.
• Method Signature: Checks that run() is an async function and get_tools() is callable.
• Schema Compliance: Validates that the tools returned by get_tools() conform to the OpenAI-compatible function calling format.
  Sources: osmosis_ai/rollout/validator.py:100-260

Data Models (Schemas)

The protocol relies on strict Pydantic models to define the communication interface between the server and TrainGate.

RolloutRequest Fields

Field	Type	Description
rollout_id	str	Unique ID (1-256 chars).
server_url	str	Callback URL for TrainGate.
messages	List[Dict]	Initial conversation history.
completion_params	Dict	Temperature, top_p, etc.
metadata	Dict	Opaque configuration data (max 1MB).
Sources: osmosis_ai/rollout/core/schemas.py:321-383

RolloutMetrics Fields

Metrics are collected throughout the trajectory to monitor performance:

• total_latency_ms: Total wall-clock time.
• llm_latency_ms: Cumulative time spent in LLM calls.
• tool_latency_ms: Cumulative time spent in tool execution.
• prompt_tokens / response_tokens: Token usage statistics.
  Sources: osmosis_ai/rollout/core/schemas.py:236-261

Conclusion

The Rollout Architecture is a robust, asynchronous framework designed to decouple agent logic from infrastructure. By leveraging a centralized RolloutContext and standardized OpenAI-compatible schemas, it provides a consistent interface for agent training workflows while maintaining high performance through concurrency controls and idempotency handling.