Rollout 接口#

本节介绍 RLinf 框架中 Rollout 类的关键 API。 它包含基于 SGLang 和 Hugging Face 后端的实现。

SGLang#

class rlinf.workers.rollout.sglang.sglang_worker.SGLangWorker#
__init__(config, placement, weight_reload='sync', config_rollout=None)#

Initialize the Worker with the given parent address and world size.

Only non-Ray workers should provide parent_address, world_size and rank. For example, when a Worker is created via multiprocessing by another Worker, the parent address, world size and rank should be provided.

参数:

  • parent_address (Optional[WorkerAddress]) -- The address of the parent worker. This is used to set up the WorkerAddress for this worker.

  • world_size (Optional[int]) -- The total number of workers in the group. If not provided, it will be set to the environment variable WORLD_SIZE.

  • rank (Optional[int]) -- The rank of this worker in the group. If not provided, it will be set to the environment variable RANK.

  • config (DictConfig)

  • placement (ModelParallelComponentPlacement)

  • weight_reload (Literal['sync', 'cpu', None])

  • config_rollout (DictConfig)

get_sampling_param_from_config()#

Get sampling parameters from the configuration.

参数:

cfg_sampling_params (DictConfig)

返回类型:

dict

shutdown()#

Shutdown the SGLang task.

async async_generate(prompt=None, sampling_params=None, input_ids=None, image_data=None, return_logprob=False, request_info=None)#

Asynchronously generate text using the underlying SGLang engine and return the engine result together with the original input_ids, answers, and idx.

This wrapper calls self._engine.async_generate(...) and forwards the provided arguments. Because the SGLang engine does not include the original input_ids in its response, this method returns the input_ids alongside the engine result for downstream use.

参数:

  • prompt (List[str] | str | None) -- Same as SGLang engine's prompt argument.

  • sampling_params (List[Dict] | Dict | None) -- Same as SGLang engine's sampling_params argument.

  • input_ids (List[List[int]] | List[int] | None) -- Same as SGLang engine's input_ids argument.

  • return_logprob (List[bool] | bool | None) -- Same as SGLang engine's return_logprob argument.

  • request_info (Any | None) -- Any additional request info you wish to be associated with this generation request. This argument will not be passed to the SGLang engine and returned directly.

  • image_data (list | None)

返回:

A tuple containing the engine result and the original request_info.

返回类型:

Tuple[Dict, Any | None]

async offload_engine()#

Release the model weights from the SGLang engine.

async onload_engine()#

Onload the model weights from cpu to the SGLang engine.

async abort_generation()#

Abort the generation.

async sync_model_from_actor()#

Update the weights of the SGLang engine.

Huggingface#

class rlinf.workers.rollout.hf.huggingface_worker.MultiStepRolloutWorker#
__init__(cfg)#

Initialize the Worker with the given parent address and world size.

Only non-Ray workers should provide parent_address, world_size and rank. For example, when a Worker is created via multiprocessing by another Worker, the parent address, world size and rank should be provided.

参数:

  • parent_address (Optional[WorkerAddress]) -- The address of the parent worker. This is used to set up the WorkerAddress for this worker.

  • world_size (Optional[int]) -- The total number of workers in the group. If not provided, it will be set to the environment variable WORLD_SIZE.

  • rank (Optional[int]) -- The rank of this worker in the group. If not provided, it will be set to the environment variable RANK.

  • cfg (DictConfig)

async recv_from_and_record_batch_routes_with_timeout(group_name, channel, *, route_key=None, tag=None, batch_size=None, merge_fn=None, infer_batch_size_fn=None, timeout_time=0.02, recv_queue_size=0)#

Receive routed batch shards and record their return routes.

This method is used in env-decoupled mode. It builds a receive plan for the source worker group, receives shard messages from channel one by one, and stops when all planned items are received or timeout_time is reached.

Each received channel item must be a dict with batch_index (route metadata) and batch (payload shard).

The batch_index values are stored in self.batch_router[tag] so a later send call can split the response and send each shard back to the original source rank. The received payload shards are validated, then merged with merge_fn or the default merge_batches.

参数:

  • group_name (str) -- Source worker group name.

  • channel (Any | None) -- Channel used to receive routed batch shards.

  • route_key (Any) -- Optional key used to separate independent routed streams.

  • tag (str | None) -- Routing tag used to build receive keys and index recorded routes.

  • batch_size (int | None) -- Expected batch size for each planned receive entry.

  • merge_fn (Callable[[list[Any]], Any] | None) -- Optional custom function for merging received shards.

  • infer_batch_size_fn (Callable[[Any], int] | None) -- Optional function used to infer shard batch size during validation.

  • timeout_time (float) -- Maximum time in seconds to wait before finalizing partial results.

  • recv_queue_size (int) -- Number of receive queue entries used when building the receive plan.

返回:

A merged payload and its split sizes. If only one shard is received, the current implementation returns that shard directly.

send_to_recorded_batch_routes(group_name, channel, data, *, route_key=None, tag=None, split_fn=None, split_sizes)#

Send split batch results back using recorded batch routes.

This method is used after a previous receive call has populated self.batch_router[tag] with batch indices from incoming messages. The outgoing data is split according to split_sizes and each shard is sent to the rank encoded in the corresponding recorded batch index.

Each outgoing channel item is a dict with batch_index (recorded batch index) and batch (payload shard).

After all shards are queued, the recorded batch indices for tag are cleared to avoid reusing stale routes.

参数:

  • group_name (str) -- Destination worker group name.

  • channel (Any | None) -- Channel used to send the split payloads.

  • data (Any) -- Payload to split and send.

  • route_key (Any) -- Optional key used to separate independent routed streams.

  • tag (str | None) -- Routing tag whose recorded batch indices should be consumed.

  • split_fn (Callable[[Any, list[int]], list[Any]] | None) -- Optional custom splitter. If omitted, split_batch is used.

  • split_sizes (list[int]) -- Batch sizes used to split data. Must have the same length as self.batch_router[tag].

返回:

AsyncRouteWork wrapping the async channel put operations.

async sync_model_from_actor()#

Sync model parameters from the actor worker.