The implementation of RolloutController

[garrett4wade]

Overview

This plan implements the RolloutController following the single-controller architecture design. The implementation is broken down into 4 major phases.

---

Phase 1: Implement Local SGLang/vLLM Engine

Step 1.1: Create LocalSGLangEngine class

File: areal/engine/sglang_local.py

Create a local inference engine that runs SGLang within the same process:

class LocalSGLangEngine(InferenceEngine):
    def __init__(self, config: InferenceEngineConfig):
        self.config = config
        self._version = 0
        self.lock = Lock()

    def initialize(self, engine_id: Optional[str] = None):
        # Import sglang runtime
        import sglang as sgl

        # Initialize local SGLang runtime
        self.engine = sgl.Engine(**self.engine_args)
        ...

    async def agenerate(self, req: ModelRequest) -> ModelResponse:
        # Use local runtime.generate() instead of HTTP requests
        # Similar logic to RemoteSGLangEngine but calling runtime directly
        pass

Notes:

1. A single LocalSGLangEngine may occupy multiple GPUs or even span multiple nodes.
2. There is an old experimental implementation in areal/experimental/sglang_engine.py, which requires significant refactoring before use.

Step 1.2: Create LocalvLLMEngine class

File: areal/engine/vllm_local.py

Similar structure to LocalSGLangEngine but using vLLM's AsyncLLMEngine.

---

Phase 2: Refactor WorkflowExecutor to Isolate Async Thread Logic

Step 2.1: Create generic AsyncTaskRunner

File: areal/core/async_task_runner.py

Create a generic, reusable async task executor with NO AReaL-specific logic:

from typing import Callable, Awaitable, TypeVar, Dict, List, Tuple, Any
import asyncio
import queue
import threading
import time
import traceback
import uvloop
from dataclasses import dataclass

T = TypeVar('T')

@dataclass
class _TaskResult:
    task_id: str
    result: Any
    timestamp: float

class AsyncTaskRunner:
    """Generic async task executor - runs any async function in a background thread

    This is a pure utility with NO capacity management or domain-specific logic.
    """

    def __init__(
        self,
        max_concurrent_tasks: int,
        queue_size: int | None = None,
        logger=None,
    ):
        self.max_concurrent_tasks = max_concurrent_tasks
        self.logger = logger

        # Queues
        qsize = queue_size or max_concurrent_tasks * 16
        self.input_queue = queue.Queue(maxsize=qsize)
        self.output_queue = queue.Queue(maxsize=qsize)

        # Control
        self.exiting = threading.Event()
        self.paused = threading.Event()

        # Task tracking
        self.running_tasks: Dict[str, asyncio.Task] = {}
        self.task_id_counter = 0

    def start(self):
        """Start the async execution thread"""
        self.thread = threading.Thread(target=self._run_thread, daemon=True)
        self.thread.start()

    def stop(self):
        """Stop the async execution thread"""
        self.exiting.set()
        self.thread.join()

    def submit(self, async_fn: Callable[[], Awaitable[T]]) -> str:
        """Submit an async function for execution

        Parameters
        ----------
        async_fn : Callable[[], Awaitable[T]]
            Async function to execute (no arguments)

        Returns
        -------
        str
            Task ID for tracking
        """
        task_id = str(self.task_id_counter)
        self.task_id_counter += 1

        try:
            self.input_queue.put_nowait((task_id, async_fn))
        except queue.Full:
            raise RuntimeError("Task queue is full")

        return task_id

    def wait(self, count: int, timeout: float | None = None) -> List[Tuple[str, Any]]:
        """Wait for task results

        Parameters
        ----------
        count : int
            Number of results to wait for
        timeout : float | None
            Timeout in seconds

        Returns
        -------
        List[Tuple[str, Any]]
            List of (task_id, result) pairs
        """
        tik = time.perf_counter()
        timeout = timeout or float(7 * 24 * 3600)
        results = []

        while not self.exiting.is_set() and time.perf_counter() - tik < timeout:
            # Drain output queue
            while True:
                try:
                    task_result = self.output_queue.get_nowait()
                    results.append((task_result.task_id, task_result.result))
                except queue.Empty:
                    break

            if len(results) >= count:
                break
            else:
                time.sleep(0.05)

        if self.exiting.is_set():
            raise RuntimeError("Task runner is exiting")
        if len(results) < count:
            raise TimeoutError(f"Timed out waiting for {count} results, got {len(results)}")

        return results[:count]

    def pause(self):
        """Pause task submission"""
        self.paused.set()

    def resume(self):
        """Resume task submission"""
        self.paused.clear()

    def _run_thread(self):
        """Thread entry point"""
        try:
            uvloop.run(self._async_loop())
        except Exception:
            traceback.print_exc()

    async def _async_loop(self):
        # Similar to WorkflowExecutor._rollout_thread_async
        ...

Step 2.2: Refactor WorkflowExecutor to use AsyncTaskRunner

File: areal/api/workflow_api.py

WorkflowExecutor composes AsyncTaskRunner and handles all AReaL-specific logic externally. Key changes:

1. Use AsyncTaskRunner: Generic async task executor (no AReaL logic)
2. Remote workflow execution: Support passing workflow class path + config dict for dynamic import/instantiation in the async task (enables workflow execution on remote inference engine process)

from areal.core.async_task_runner import AsyncTaskRunner

class WorkflowExecutor:
    def __init__(
        self,
        config: InferenceEngineConfig,
        staleness_manager: StalenessManager | None = None,
    ):
        self.config = config
        # No inference_engine attribute - it's injected on-demand

        # Staleness manager will be initialized in initialize()
        self.staleness_manager = staleness_manager

        # Generic async task runner (no AReaL logic)
        self.task_runner: Optional[AsyncTaskRunner] = None

    def initialize(self, logger=None, train_data_parallel_size: int | None = None):
        if logger is None:
            logger = logging.getLogger("WorkflowExecutor")
        self.logger = logger

        # Initialize staleness manager if not provided
        if self.staleness_manager is None:
            # Same staleness manager initialization as current code
            self.staleness_manager = StalenessManager(...)

        # Create generic task runner
        self.task_runner = AsyncTaskRunner(
            max_concurrent_tasks=self.staleness_manager.max_concurrent_rollouts,
            queue_size=self.config.queue_size,
            logger=self.logger,
        )
        self.task_runner.start()

    def submit(
        self,
        data: Dict[str, Any],
        workflow: "RolloutWorkflow" | None = None,
        workflow_builder: Callable | None = None,
        workflow_class_path: str | None = None,  # NEW: for remote execution
        workflow_args: Dict[str, Any] | None = None,  # NEW: for remote execution
        should_accept: Callable | None = None,
    ) -> None:
        """Submit a rollout request

        Parameters
        ----------
        data : Dict[str, Any]
            Input data for the workflow
        workflow : RolloutWorkflow | None
            Workflow instance (for local execution)
        workflow_builder : Callable | None
            Workflow builder function (for local execution)
        workflow_class_path : str | None
            Workflow class import path (e.g., "areal.workflow.rlvr.RLVRWorkflow").
            For remote execution - will be imported dynamically in async task.
        workflow_args : Dict[str, Any] | None
            Config dict for workflow initialization (used with workflow_class_path)
        should_accept : Callable | None
            Filter function for results
        """
        # Validate workflow specification
        has_instance = workflow is not None or workflow_builder is not None
        has_class_path = workflow_class_path is not None

        if not has_instance and not has_class_path:
            raise ValueError("Must provide either workflow/workflow_builder or workflow_class_path")
        if has_instance and has_class_path:
            raise ValueError("Cannot provide both workflow instance and workflow_class_path")

        # Create async function that captures workflow.arun_episode
        async def run_episode():
            # Option 1: Local execution with workflow instance
            if has_instance:
                if workflow is None:
                    wf = workflow_builder()
                else:
                    wf = workflow
                result = await wf.arun_episode(inference_engine, data)

            # Option 2: Remote execution with dynamic import
            else:
                # Dynamically import workflow class
                module_path, class_name = workflow_class_path.rsplit('.', 1)
                module = __import__(module_path, fromlist=[class_name])
                workflow_class = getattr(module, class_name)

                # Instantiate with config
                wf = workflow_class(**(workflow_args or {}))
                result = await wf.arun_episode(inference_engine, data)

            # Check trajectory format if enabled
            if self.config.check_trajectory_format and result is not None:
                check_trajectory_format(
                    result, expected_keys=self._expected_trajectory_keys
                )
                # Track expected keys
                if isinstance(result, dict) and "input_ids" in result:
                    if self._expected_trajectory_keys is None:
                        self._expected_trajectory_keys = set(result.keys())

            # Convert CompletionWithTokenLogpReward to tensor dict
            if isinstance(result, dict) and all(
                isinstance(v, CompletionWithTokenLogpReward) for v in result.values()
            ):
                result = concat_padded_tensors(
                    [v.to_tensor_dict() for v in result.values()]
                )

            # Apply should_accept filter
            if result is not None and should_accept is not None:
                if not should_accept(result):
                    return None

            return result

        # Submit to task runner and notify staleness manager
        task_id = self.task_runner.submit(run_episode)
        self.staleness_manager.on_rollout_submitted()

        if self.config.enable_rollout_tracing:
            stat = self.staleness_manager.get_stats()
            self.logger.info(
                f"Submit rollout {task_id}. "
                f"Submit: {stat.submitted}, running: {stat.running}, accepted: {stat.accepted}."
            )

    def wait(self, count: int, timeout: float | None = None) -> Dict[str, Any]:
        """Wait for workflow results

        Handles staleness callbacks and result caching.
        """
        tik = time.perf_counter()
        timeout = timeout or float(7 * 24 * 3600)

        while time.perf_counter() - tik < timeout:
            # Try to get results from task runner
            try:
                results = self.task_runner.wait(count=1, timeout=0.05)
                for task_id, result in results:
                    # Handle accepted/rejected rollouts
                    if result is not None:
                        # Accepted
                        self.staleness_manager.on_rollout_accepted()
                        self.result_cache.append(_TimedResult(
                            t=time.monotonic_ns(),
                            data=result
                        ))
                        if self.config.enable_rollout_tracing:
                            stat = self.staleness_manager.get_stats()
                            self.logger.info(
                                f"Accept rollout {task_id}. "
                                f"Submit: {stat.submitted}, running: {stat.running}, "
                                f"accepted: {stat.accepted}."
                            )
                    else:
                        # Rejected
                        self.staleness_manager.on_rollout_rejected()
                        if self.config.enable_rollout_tracing:
                            stat = self.staleness_manager.get_stats()
                            self.logger.info(
                                f"Reject rollout {task_id}. "
                                f"Submit: {stat.submitted}, running: {stat.running}, "
                                f"accepted: {stat.accepted}."
                            )
            except TimeoutError:
                pass

            # Check if we have enough results
            if len(self.result_cache) >= count:
                break

        if len(self.result_cache) < count:
            raise TimeoutError(
                f"Timed out waiting for {count} rollouts, only received {len(self.result_cache)}."
            )

        # Return results
        self.result_cache.sort(key=lambda x: x.t)
        results, self.result_cache = (
            self.result_cache[:count],
            self.result_cache[count:],
        )
        random.shuffle(results)
        return concat_padded_tensors([r.data for r in results])

    def pause(self):
        """Pause request submission"""
        if self.task_runner:
            self.task_runner.pause()

    def resume(self):
        """Resume request submission"""
        if self.task_runner:
            self.task_runner.resume()

    # rollout_batch() and prepare_batch() also accept workflow_class_path + workflow_args (for remote execution)

Usage Examples:

# Example 1: Local execution with workflow instance
executor.submit(
    data={"prompt": "What is 2+2?"},
    workflow=RLVRWorkflow(n_samples=4),
)

# Example 2: Remote execution with class path (workflow runs on remote process)
executor.submit(
    data={"prompt": "What is 2+2?"},
    workflow_class_path="areal.workflow.rlvr.RLVRWorkflow",
    workflow_args={"n_samples": 4, "temperature": 0.8},
)

---

Phase 3: Extend RPCClient/RPCServer

1. Support async calls.
2. Add a new endpoint in RPCServer called /exec_workflow, which instantiates a workflow object, runs workflow.arun_episode(self.engine, data), stores the produced data locally, and returns the metadata to the client.

---

Phase 4: Implement RolloutController

RolloutController is a pure composition of refactored components.

Step 4.1: Create RolloutController class

File: areal/controller/rollout_controller.py

from typing import List, Dict, Any, Callable, Optional
from areal.api.scheduler_api import Scheduler, Worker
from areal.api.engine_api import InferenceEngine
from areal.api.cli_args import InferenceEngineConfig
from areal.api.workflow_api import WorkflowExecutor, RolloutWorkflow
from areal.scheduler.rpc.rpc_client import RPCClient

class RolloutController:
    """Controller for orchestrating M independent inference engine workers

    This is a pure composition layer that:
    1. Launches local inference engines on workers via scheduler
    2. Schedules requests to specific engines (e.g., via round-robin)
    3. Delegates actual execution to AsyncTaskRunner
    """

    def __init__(
        self,
        engine_class: type[InferenceEngine],  # e.g., LocalSGLangEngine
        config: InferenceEngineConfig,
        scheduler: Scheduler,
    ):
        self.engine_class = engine_class
        self.config = config
        self.scheduler = scheduler

        # Workers and RPC
        self.workers: List[Worker] = []
        self.rpc_client = RPCClient()

        # Round-robin scheduling
        self.current_engine_idx = 0  # Round-robin index

        # WorkflowExecutor (manages async execution + staleness)
        self.workflow_executor: Optional[WorkflowExecutor] = None

    def initialize(
        self,
        engine_id: Optional[str] = None,
        train_data_parallel_size: int | None = None
    ):
        """Initialize workers that host local inference engines"""

        if engine_id is None:
            engine_id = uuid.uuid4().hex
        self.engine_id = engine_id
        self.logger = logging.getLogger(f"[RolloutController {engine_id}]")

        # Step 1: Create workers via scheduler
        worker_key = f"rollout_controller_{self.engine_id}"
        scheduling_config = self._build_scheduling_config()

        job_id = self.scheduler.create_workers(
            worker_key=worker_key,
            scheduler_config=scheduling_config,
        )

        # Step 2: Get worker endpoints
        self.workers = self.scheduler.get_workers(
            worker_key=worker_key,
            timeout=self.config.setup_timeout
        )
        self.logger.info(f"Got {len(self.workers)} workers: {self.workers}")

        # Step 3: Register workers with RPC client
        for worker in self.workers:
            self.rpc_client.register(
                worker_id=worker.id,
                ip=worker.ip,
                port=int(worker.ports[0])  # Use first port for RPC
            )

        # Step 4: Create local inference engines on each worker via RPC
        for worker in self.workers:
            self.rpc_client.create_engine(
                worker_id=worker.id,
                engine_obj=self.engine_class(self.config),
                init_config=self.config,
            )

        # Step 5: Create AsyncTaskRunner (delegates async execution)
        self.task_runner = AsyncTaskRunner(...)
        self.task_runner.initialize(...)

    # === Request scheduling with load balancing ===

    def _choose_worker(self) -> Worker:
        """Choose a worker via round-robin scheduling"""
        worker = self.workers[self.current_engine_idx]
        self.current_engine_idx = (self.current_engine_idx + 1) % len(self.workers)
        return worker

    def submit(
        self,
        data: Dict[str, Any],
        workflow_class_path: str | None = None,
        workflow_args: Dict[str, Any] | None = None,
        should_accept_path: str | None = None
    ) -> None:
        worker = self._choose_worker()

        # Create async function that executes workflow on remote worker
        async def run_episode():
            meta: DistributedBatch = await self.scheduler.call_engine(
                worker,
                "/exec_workflow",
                json=dict(
                    data=data,
                    workflow_class_path=workflow_class_path,
                    workflow_args=workflow_args,
                    should_accept_path=should_accept_path
                )
            )
            return meta

        # Submit to task runner and notify staleness manager
        task_id = self.task_runner.submit(run_episode)
        self.staleness_manager.on_rollout_submitted()

        if self.config.enable_rollout_tracing:
            stat = self.staleness_manager.get_stats()
            self.logger.info(
                f"Submit rollout {task_id}. "
                f"Submit: {stat.submitted}, running: {stat.running}, accepted: {stat.accepted}."
            )

    def wait(self, count: int, timeout: float | None = None) -> DistributedBatch:
        # Similar to WorkflowExecutor.wait
        ...

[nuzant]

In general LGTM. However, I have two questions about detailed design and implementation:

1. What if in the future we need to use a pre-deployed inference resource pool for serverless inference? Can we just replace LocalInferenceEngine with RemoteInferenceEngine in this design to use remote inference server?
2. What is the point of task_runner in RolloutController? I did not really understand what /exec_workflow does. Since running a workflow requires a WorkflowExecutor, it seems /exec_workflow should be a RPC call to WorkflowExecutor.submit. In this scenario, it does not need a task runner to execute async call engine.

[garrett4wade]

1. LocalInferenceEngine Architecture

The LocalInferenceEngine functions as a server with two primary responsibilities:

1. Executing the user's workflow code (executing it in the controller process would be computationally expensive)
2. Generating responses on GPU

Implementing a LocalInferenceEngine is the simplest way to achieve this goal. While RemoteInferenceEngine is viable, it requires launching actual SGLang/vLLM servers in advance, which is inconvenient. For serverless inference, we'll replace LocalInferenceEngine with a ServerlessInferenceEngine that uses the inference service to generate responses rather than generating them locally.

2. RolloutController Implementation

RolloutController should implement logic similar to WorkflowExecutor. Instead of running workflow.arun_episode directly, the task_runner in RolloutController submits HTTP requests to run the workflow remotely. The server side (i.e., LocalInferenceEngine) exposes async HTTP endpoints (/exec_workflow) to RolloutController. This is a simple RPC without complex task submission logic. The metadata will be returned immediately when a trajectory is complete.

What is the role of WorkflowExecutor in LocalInferenceEngine?

It is unnecessary in single-controller mode. We've moved its core logic to RolloutController. LocalInferenceEngine is simply an RPC server without task submission logic.

The reason why we need WorkflowExecutor in LocalInferenceEngine is that we hope this atomic module can be used independently in users' code snippet, without launching, scheduling, and control plane logic.

TL;DR

• Single-controller mode: (top layer, controller) HTTP task runner → (/exec_workflow) workflow.arun_episode(engine, data) → (bottom layer, engine) SGLang/vLLM engine agenerate
• SPMD mode: SGLang/vLLM remote engine → workflow executor/task runner → workflow.arun_episode(engine, data) → SGLang/vLLM engine agenerate

[garrett4wade]

Implementation Plan Update

About LocalInfEngine

While logically clean, local SGLang/vLLM engines can conflict with AReaL's internal workflow executor thread. We've decided to transition to an alternative approach as outlined in: #476

About RPCClient/RPCServer and Workflow Execution

We initially planned to implement a FastAPI HTTP server with an asynchronous /run_workflow endpoint, allowing the ASGI server to handle concurrent workflow execution. However, uvicorn spawns subthread workers to run endpoint methods, and these worker threads are incompatible with AReaL's internal threading. This forces us to switch to a synchronous server like Flask. With a synchronous server, async endpoints run within the same thread using asyncio.run, which is blocking—meaning we can't use an async /run_workflow endpoint to leverage the HTTP server for concurrent workflow execution anyway.

As an alternative, we'll use submit and wait engine APIs to let the workflow executor's internal thread handle workflow execution. The rollout controller's submit and wait methods are roughly implemented as follows:

    def submit(
        self,
        data: dict[str, Any],
        workflow_path: str,
        workflow_kwargs: dict[str, Any],
        should_accept_path: str | None = None,
    ) -> None:
        # Add to pending queue (will be submitted when capacity allows)
        self._pending_inputs.append(
            _RemoteRolloutTaskInput(
                data=data,
                workflow_kwargs=workflow_kwargs,
                workflow_path=workflow_path,
                should_accept_path=should_accept_path,
            )
        )

    async def _wait_callback(self, worker: Worker):
        # Wait for a generation to return
        result = "NO_RESULT"
        tik = time.time()
        while result == "NO_RESULT" and time.time() - tik < self.config.request_timeout:
            result = await self.scheduler.async_call_engine(
                worker.id, "wait_quiet", count=1, timeout=1, max_retries=1
            )

        # The RPCServer will return None if the
        # trajectory is rejected.
        if result is not None:
            self.staleness_manager.on_rollout_accepted()
            if self.config.enable_rollout_tracing:
                stat = self.staleness_manager.get_stats()
                self.logger.info(
                    f"Finish and accept rollout. "
                    f"Submit: {stat.submitted}, "
                    f"running: {stat.running}, "
                    f"accepted: {stat.accepted}."
                )
            return result
        else:
            self.staleness_manager.on_rollout_rejected()
            if self.config.enable_rollout_tracing:
                stat = self.staleness_manager.get_stats()
                self.logger.info(
                    f"Finish but reject rollout. "
                    f"Submit: {stat.submitted}, "
                    f"running: {stat.running}, "
                    f"accepted: {stat.accepted}."
                )
            return None

    def _commit_one_to_runner(self):
        """Commit one pending input to task runner with staleness tracking."""
        task_input = self._pending_inputs.pop(0)

        # Choose worker via round-robin
        worker = self._choose_worker()

        self.scheduler.call_engine(
            worker.id,
            "submit",
            data=task_input.data,
            workflow_path=task_input.workflow_path,
            workflow_kwargs=task_input.workflow_kwargs,
            should_accept_path=task_input.should_accept_path,
        )

        # Submit a wait callback to AsyncTaskRunner
        try:
            self.runner.submit(
                self._wait_callback,
                worker,
            )
        except TaskQueueFullError:
            raise queue.Full("Input queue full")

        # Notify staleness manager AFTER successful submission
        self.staleness_manager.on_rollout_submitted()
        if self.config.enable_rollout_tracing:
            stat = self.staleness_manager.get_stats()
            self.logger.info(
                f"Submit rollout. "
                f"Submit: {stat.submitted}, "
                f"running: {stat.running}, "
                f"accepted: {stat.accepted}."
            )

    def wait(self, count: int, timeout: float | None = None) -> DistributedBatch:
        #######################################################
        # The following logic is copied from WorkflowExecutor #
        #######################################################
        start_time = time.perf_counter()
        timeout = timeout or float(7 * 24 * 3600)

        # Keep requesting results from runner until we have enough accepted
        # (non-None) results. Use short timeout (1 second) for each wait call
        # to allow periodic checking. This matches original behavior where
        # wait() would poll and allow prepare_batch() to continue
        while True:
            # Submit pending inputs
            # Check capacity before submitting
            capacity = self.get_capacity()
            # Submit pending tasks
            for _ in range(capacity):
                if len(self._pending_inputs) == 0:
                    break
                self._commit_one_to_runner()

            if len(self._pending_results) >= count:
                break

            elapsed = time.perf_counter() - start_time
            remaining_timeout = timeout - elapsed

            if remaining_timeout <= 0:
                raise TimeoutError(
                    f"Timed out waiting for {count} rollouts, only received "
                    f"{len(self._pending_results)}."
                )

            # Try to get at least the number we still need, but request at least 1
            # Note: runner.wait() might return fewer due to rejections (None results)
            needed = max(1, count - len(self._pending_results))

            try:
                # Use short timeout to allow periodic returns (matches original
                # polling behavior)
                batch = self.runner.wait(
                    count=needed, timeout=min(0.1, remaining_timeout)
                )

                # Filter out None results (rejected trajectories)
                # runner.wait() returns List[T] where T can be None for
                # rejected rollouts
                accepted = [result for result in batch if result is not None]
                self._pending_results.extend(accepted)
            except TimeoutError:
                pass

        if self.config.enable_rollout_tracing:
            self.logger.info("Rollout results are ready!")

        # Extract requested number of results
        results = self._pending_results[:count]
        self._pending_results = self._pending_results[count:]

        # Shuffle for randomness (helps with data diversity)
        random.shuffle(results)

        # Convert to DistributedBatch
        if len(results) == 0:
            return DistributedBatchMemory.from_dict({})

        return DistributedBatchMemory.concat(
            [DistributedBatchMemory.from_dict(r) for r in results]
        )

In summary, we call the submit method via RPC to submit a request to a specific worker based on workloads. We then register a wait callback and add it to the AsyncTaskRunner. This task continuously calls the wait method via RPC to fetch results. Note that results may not return in submission order. This approach allows us to reuse the wait logic from WorkflowExecutor while maintaining global staleness and capacity management.