nemo_rl.models.generation.vllm.vllm_worker_async

Module Contents

Classes

VllmAsyncGenerationWorkerImpl
VllmAsyncGenerationWorker

Data

LOGGER

API

nemo_rl.models.generation.vllm.vllm_worker_async.LOGGER

‘getLogger(…)’

class nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorkerImpl(

config,

bundle_indices=None,

fraction_of_gpus: float = 1.0,

seed=None,

extra_env_vars: Optional[list[str]] = None,

defer_model_load: bool = False,

)

Bases: nemo_rl.models.generation.vllm.checkpoint_engine.VllmAsyncCheckpointEngineRpcMixin, nemo_rl.models.generation.vllm.vllm_worker.BaseVllmGenerationWorker

_return_routed_experts_enabled() → bool

_reserve_port() → None

Bind and listen on a TCP socket to reserve a free port from the OS.

The socket is held open in LISTENING state and later passed directly to uvicorn via the sockets= parameter in server.serve(). The socket is never closed and re-opened, so there is zero gap where another process could steal the port.

load_model() → None

Load the vLLM model and create the engine.

Called after a deferred init to perform the heavy model loading.

_create_engine(llm_kwargs: dict[str, Any]) → None

_install_engine_input_socket_lock() → None

Serialise sends on AsyncMPClient.input_socket across OS threads to prevent race conditions that block the vLLM engine (e.g. during in flight weight updates in async grpo).

_start_vllm_metrics_logger() → None

Start a background thread that periodically collects vLLM logger metrics.

Controlled by vllm_metrics_logger_interval (default: 0.5) in vllm_cfg. Runs only on the model-owner actor.

get_vllm_logger_metrics() → dict[str, Any]

clear_vllm_logger_metrics() → None

async post_init_async()

async get_reserved_url() → Optional[str]

Return the URL from the reserved socket, available before model loading.

async report_dp_openai_server_base_url() → Optional[str]

_setup_vllm_openai_api_server(app: fastapi.FastAPI) → fastapi.FastAPI

_setup_vllm_server() → tuple[threading.Thread, str, uvicorn.Server]

async init_collective_async(

rank_prefix: int,

ip: str,

port: int,

world_size: int,

train_world_size: int,

) → None

async generate_async(

data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],

greedy: bool = False,

) → AsyncGenerator[tuple[int, nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]], None]

Generate a batch of data using vLLM’s AsyncLLMEngine, yielding results as they are ready.

Parameters:

• data – BatchedDataDict with input_ids and input_lengths

• greedy – Whether to use greedy decoding instead of sampling

Yields:

Tuple of (original_index, BatchedDataDict conforming to GenerationOutputSpec for the single sequence)

async generate_text_async(

data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],

greedy: bool = False,

) → AsyncGenerator[tuple[int, nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]], None]

Generate text responses asynchronously, yielding results as they are ready.

Parameters:

• data – BatchedDataDict containing prompts with text strings

• greedy – Whether to use greedy decoding instead of sampling

Yields:

Tuple of (original_index, BatchedDataDict containing single text response)

async report_device_id_async() → list[str]

Async version of report_device_id.

async prepare_refit_info_async(

state_dict_info: dict[str, Any],

) → None

Async version of prepare_refit_info.

async update_weights_via_ipc_zmq_async() → bool

Async version of update_weights_via_ipc_zmq.

async update_weights_from_collective_async() → bool

Async version of update_weights_from_collective.

async init_nccl_reshard_comm_group_async(

rank_prefix: int,

pp_ips: list[str],

pp_ports: list[int],

pp_size: int,

train_ranks_per_stage: int,

sub_world_size: int,

) → None

Async version of init_nccl_reshard_comm_group.

async prepare_nccl_reshard_refit_info_async(refit_info: dict) → None

Async version of prepare_nccl_reshard_refit_info.

async nccl_reshard_refit_async() → bool

Async version of nccl_reshard_refit.

async reset_prefix_cache_async()

Async version of reset_prefix_cache.

async sleep_async()

Async version of sleep.

async wake_up_async(**kwargs)

Async version of wake_up.

async shutdown() → bool

Clean up vLLM resources.

class nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker(

config,

bundle_indices=None,

fraction_of_gpus: float = 1.0,

seed=None,

extra_env_vars: Optional[list[str]] = None,

defer_model_load: bool = False,

)

Bases: nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorkerImpl