nemo_rl.models.generation.vllm.vllm_generation#

Module Contents#

Classes#

VllmGeneration

Functions#

http_server_api_key_and_worker_config

Generate the exposed HTTP server’s bearer token and the worker config carrying it.

_record_vllm_generation_metrics

Record vLLM token-usage metrics to nemo-lens (no-op unless exporting).

Data#

logger

API#

nemo_rl.models.generation.vllm.vllm_generation.logger#

‘getLogger(…)’

nemo_rl.models.generation.vllm.vllm_generation.http_server_api_key_and_worker_config(
config: nemo_rl.models.generation.vllm.config.VllmConfig,
) → tuple[Optional[str], nemo_rl.models.generation.vllm.config.VllmConfig]#

Generate the exposed HTTP server’s bearer token and the worker config carrying it.

Reads vllm_cfg.http_server_api_key_required and vllm_cfg.expose_http_server from config. When the flag is off, no key exists and config is returned as the worker config. When it is on, a random URL-safe key is generated and placed as vllm_cfg.http_server_api_key in a shallow copy of config; the workers receive that copy and refuse OpenAI-route requests without the key. config itself is never modified, so the job’s saved and logged configuration does not carry the key. The flag without an exposed server is a configuration error: there would be no server to protect and no key to hand to Gym.

Returns:

(api_key, worker_config); api_key is None when no key is required.

nemo_rl.models.generation.vllm.vllm_generation._record_vllm_generation_metrics(
model_name: str | None,
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict,
combined: nemo_rl.distributed.batched_data_dict.BatchedDataDict,
batch_duration_s: float | None = None,
) → None#

Record vLLM token-usage metrics to nemo-lens (no-op unless exporting).

class nemo_rl.models.generation.vllm.vllm_generation.VllmGeneration(
cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster,
config: nemo_rl.models.generation.vllm.config.VllmConfig,
name_prefix: str = 'vllm_policy',
workers_per_node: Optional[Union[int, list[int]]] = None,
defer_model_load: bool = False,
)#

Bases: nemo_rl.models.generation.interfaces.GenerationInterface

classmethod validate_settings(
master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig,
) → None#

Reject pure-config vLLM settings the SC entrypoint cannot honor.

static init_cluster_placement_groups(
cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster,
config: nemo_rl.models.generation.vllm.config.VllmConfig,
) → None#

Pre-initialize placement groups matching the strategy VllmGeneration expects.

Call this before constructing VllmGeneration when other components compete for the same Ray resources and you need deterministic ordering — topology-constrained inference PGs should be created before unconstrained ones so they claim domain-aligned nodes first.

VllmGeneration.__init__ calls _init_placement_groups internally, but that call early-returns when PGs already exist, so calling this method first is safe.

_get_tied_worker_bundle_indices(
cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster,
) → list[tuple[int, list[int]]]#

Calculate bundle indices for tensor and pipeline parallel workers.

Handles both unified placement groups (for cross-node model parallelism) and per-node placement groups (for node-local model parallelism).

_report_device_id() → list[list[str]]#

Report the device ID of vllm workers.

_report_dp_openai_server_base_urls() → list[Optional[str]]#

Report the data parallel OpenAI server base URLs of vLLM workers, only populated if it is async vLLM engine and the HTTP server is active.

_collect_reserved_urls() → list[Optional[str]]#

Collect reserved URLs from DP leaders before model loading.

Only called when defer_model_load=True. Workers have bound ports during init and can report their reserved URLs immediately.

load_and_start() → None#

Load models on all workers and start HTTP servers.

Called after a deferred init (defer_model_load=True) to perform the heavy model loading. Updates dp_openai_server_base_urls with the actual running server URLs and populates device_uuids.

_post_init()#
setup_token_capture(
dp_cfg: nemo_rl.data_plane.interfaces.DataPlaneConfig,
staging_partition: str,
*,
capture_media: bool = False,
) → MediaColumnSpec | None#

Install ledger-authoritative token capture in every DP-leader worker.

Called once at setup when token_capture.enabled; each async worker builds its in-worker data-plane client + TQTokenSink and makes the single Gym install_capture call. Returns the media column spec the leaders pinned (None for text-only capture).

set_rollout_weight_version(version: int) → None#

Rotate the weight version workers stamp on captured model calls.

_get_raw_spec_counters() → dict[str | tuple[str, int], float]#

Collect raw spec decode counters from workers.

snapshot_step_metrics() → None#

Snapshot current spec decode counters to begin tracking a training step.

Call this before generation to establish a baseline for metrics delta.

Raises:

RuntimeWarning – If called twice without get_step_metrics() in between.

get_step_metrics() → dict[str, float]#

Get the vLLM engine metrics delta since snapshot_step_metrics().

Covers the speculative-decoding family plus the engine’s token, sequence-length and request-outcome series, all derived from the same pair of snapshots so the wider coverage costs no extra RPC.

Returns:

Dictionary of delta metrics with ‘vllm/’ prefix. Returns empty dict if snapshot_step_metrics() was not called.

Raises:

RuntimeWarning – If called without snapshot_step_metrics() first.

init_collective(
ip: str,
port: int,
world_size: int,
*,
train_world_size: int,
) → list[ray.ObjectRef]#

Initialize the collective communication.

shard_liveness_ref(shard_idx: int) → ray.ObjectRef#

Liveness of the worker leading this shard. The caller need not know the layout.

log_shard_gpu_state(
shard_idx: int,
*,
label: str,
timeout_s: float = 30.0,
) → None#

Read the shard leader’s GPU from its own node. Never raises.

The leader is enough: every worker of a shard is on the same bundle set and it is the leader’s device that the replacement engine allocates on first.

restart_shard(shard_idx: int) → Optional[str]#

Rebuild one data-parallel shard’s workers and bring its engine back up.

Blocking and slow – it reloads the model – so callers run it off the control loop. Only this shard’s workers are touched; the rest of the fleet keeps serving throughout.

Mirrors the startup sequence (create workers -> load_model -> post_init -> report URL) rather than inventing a shorter one, because anything the deferred- load path does at startup is equally required for a replacement.

Returns:

The replacement’s OpenAI base URL, or None for a sync engine that exposes no HTTP server. The URL is expected to differ from the old one: the new engine binds its own port, so callers must publish it rather than assume the fleet’s URL list is still accurate.

set_refit_membership(
membership: nemo_rl.weight_sync.membership.RefitMembership,
) → None#

Record which shards take part in refits from now on.

Rebuilding the communicator is not enough on its own. Every refit dispatch – update_weights_from_collective, nccl_reshard_refit – goes through run_all_workers_*, which walks the whole worker group. Left alone they would keep calling the dead shard’s Ray actor after the rebuild and fail the refit with RayActorError, so the run would still die, just differently.

_refit_leaders_by_shard() → list[tuple[int, int, Any]]#

Snapshot (shard index, leader index, worker) in refit rank order.

Falls back to every leader when no membership has been recorded, which is the state for the entire life of a run that never loses a shard.

_refit_leader_workers() → list[Any]#

DP leaders that should receive refit calls, in rank order.

rebuild_collective(
membership: nemo_rl.weight_sync.membership.RefitMembership,
ip: str,
port: int,
) → list[ray.ObjectRef]#

Re-init the collective over the surviving shards only.

Deliberately not init_collective with a filter. run_all_workers_multiple_data walks every worker in the group, so it would dispatch to the shard we are rebuilding because it is gone – and calling into a dead Ray actor is the hang this is meant to end. Here the surviving DP leaders are addressed directly.

Only leaders are called: each one collective_rpcs into its own TP/PP workers, so one Ray call per shard reaches every rank in it.

generate(
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
greedy: bool = False,
) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]#

Generate a batch of data using vLLM.

generate_text(
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
greedy: bool = False,
) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]#

Generate text responses using vLLM.

async _async_generate_base(
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
method_name: str,
data_validation_fn,
greedy: bool = False,
) → AsyncGenerator[tuple[int, nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]], None]#

Base async generation method that handles common worker management logic.

Parameters:

  • data – Input data for generation

  • method_name – Name of the worker method to call (‘generate_async’ or ‘generate_text_async’)

  • data_validation_fn – Function to validate input data

  • greedy – Whether to use greedy decoding

Yields:

Tuple of (original_index, BatchedDataDict containing generation result)

attach_fleet_health(
monitor: nemo_rl.models.generation.fleet_health.GenerationFleetHealth,
selector: nemo_rl.models.generation.fleet_health.HealthyShardSelector,
) → None#

Route generation through fleet health from now on.

Parameters:

  • monitor – Owns shard eligibility and receives observed failures.

  • selector – Picks among the shards the monitor considers serving.

_next_dp_shard_idx() → int#

Return the data-parallel shard that should serve the next request.

async _generate_on_shard(
*,
data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],
method_name: str,
greedy: bool,
dp_shard_idx: int,
leader_worker_idx: int,
) → AsyncGenerator[tuple[int, nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]], None]#

Run one generation on a chosen shard, reporting its failures to the fleet.

A dead worker surfaces here as a Ray actor error. Reporting it is what lets the next request skip this shard instead of rediscovering the same corpse, and re-raising it as GenerationUnavailable is what tells the rollout retry policy the prompt is fine and worth re-dispatching.

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 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 responses asynchronously, yielding individual samples as they complete.

This method provides per-sample streaming across all workers, yielding each sample result as soon as it’s ready, regardless of which worker processed it.

prepare_for_generation(
*args: Any,
**kwargs: Any,
) → bool#

Wake workers up for colocated inference.

finish_generation(*args: Any, **kwargs: Any) → bool#

Sleep workers and reset prefix cache.

shutdown() → bool#

Shut down all vLLM workers and clean up resources.

prepare_refit_info(
state_dict_info: Optional[dict[str, Any]],
) → Optional[list[str]]#

Prepare the info for refit.

Returns:

When MXFP8 trainer-side pre-quantization is enabled (vllm_cfg.refit_prequantize), the parameter names the engine wants quantized on the trainer before streaming. None otherwise.

update_weights_via_ipc_zmq() → list[ray.ObjectRef]#

Update weights of the policy using IPC handles via ZMQ socket.

update_weights_from_collective(
refit_timeout_s: Optional[float] = None,
) → list[ray.ObjectRef]#

Update weights of the policy using collective communication.

init_nccl_reshard_comm_group(
pp_ips: list[str],
pp_ports: list[int],
pp_size: int,
train_ranks_per_stage: int,
sub_world_size: int,
) → list[ray.ObjectRef]#

Initialize the nccl_reshard bulk-path comm group(s) on all gen workers.

One group per PP stage (non-PP = pp_size 1).

prepare_nccl_reshard_refit_info(refit_info: dict) → None#

Forward per-layer param metadata to vLLM workers for nccl_reshard refit.

rebuild_nccl_reshard_comm_group(
membership: nemo_rl.weight_sync.membership.RefitMembership,
pp_ips: list[str],
pp_ports: list[int],
pp_size: int,
train_ranks_per_stage: int,
sub_world_size: int,
) → list[ray.ObjectRef]#

Re-init the bulk-path comm groups over the surviving shards only.

The bulk groups are sized train_ranks_per_stage + inference_world_size, so losing a shard changes their world size as well as the shared model_update_group’s – both families have to be rebuilt together or the two disagree about who is present.

nccl_reshard_refit(
refit_timeout_s: Optional[float] = None,
) → list[ray.ObjectRef]#

Receive weights from training workers via nccl_reshard (xferdtensor).

start_gpu_profiling() → None#

Start GPU profiling.

stop_gpu_profiling() → None#

Stop GPU profiling.

_collect_vllm_logger_metrics(
worker_method_name: str,
) → dict[str, Any]#

Collect one logger payload from every model-owner vLLM worker.

get_vllm_logger_metrics() → dict[str, Any]#

Collect vLLM metric histories for step-level performance reports.

drain_latest_vllm_logger_metrics() → dict[str, Any]#

Consume bounded latest-value snapshots for frequent telemetry polls.

clear_vllm_logger_metrics() → None#
clear_logger_metrics() → None#

Clear logger metrics for performance reporting.

get_logger_metrics() → dict[str, Any]#

Get logger metrics for performance reporting.

drain_latest_logger_metrics() → dict[str, Any]#

Consume latest values without transferring full worker histories.

__del__() → None#

Shuts down the worker groups when the object is deleted or is garbage collected.

This is an extra safety net in case the user forgets to call shutdown() and the pointer to the object is lost due to leaving a function scope. It’s always recommended that the user calls shutdown().

invalidate_kv_cache() → bool#

Invalidate reusable caches in vLLM (e.g., prefix/KV cache) after weight updates.

For async_engine, calls reset_prefix_cache_async on workers. For sync, calls reset_prefix_cache. Returns True if all workers report success.

pause_generation_for_refit(
*,
clear_cache: bool,
timeout_s: Optional[float] = None,
) → bool#

Pause every async vLLM engine while preserving in-flight requests.

resume_generation_after_refit(
*,
timeout_s: Optional[float] = None,
) → bool#

Resume surviving engines; fleet health excludes confirmed actor deaths.

property requires_kv_scale_sync: bool#

Check if KV cache scales should be synchronized during refit.

Only traditional per-tensor FP8 caches expose separately refittable k_scale/v_scale parameters.