`nemo_rl.models.generation.vllm.patches`#

Module Contents#

Functions#

`_get_vllm_file`	Return absolute path to a vLLM file or raise if it cannot be found.
`_locked_file_patch`	Yield (content, writer) under an exclusive file lock.
`_patch_vllm_init_workers_ray`	Patch vLLM’s Ray executor env propagation and worker runtime_env.
`_patch_vllm_llama_eagle3_own_lm_head`	Patch LlamaEagle3 to keep truncated draft lm_head ownership.
`_patch_vllm_tool_parser_namespace_tool`	Guard vLLM’s NamespaceTool import for openai < 2.25.
`_patch_vllm_ray_executor_v2_tcpstore_port`	Keep RayExecutorV2’s TCPStore port out of the MessageQueue’s scan range.
`_patch_vllm_shm_broadcast_bind_retry`	Make MessageQueue’s remote socket survive losing a port race.
`_patch_vllm_radio_layerscale_loader`	Load explicit RADIO LayerScale weights and initialize folded weights.
`_patch_vllm_glm_decoder_sequence_parallel_moe`	Restore the vLLM 0.24 decoder boundary for GLM DSA models.
`ensure_vllm_source_compat`	Apply interpreter-independent vLLM source-compat patches.
`_apply_vllm_patches`

API#

nemo_rl.models.generation.vllm.patches._get_vllm_file(relative_path: str) → str#

Return absolute path to a vLLM file or raise if it cannot be found.

The relative_path should be a POSIX-style path under the vllm package root, e.g. “v1/executor/ray_executor.py” or “attention/layer.py”.

nemo_rl.models.generation.vllm.patches._locked_file_patch(file_path: str)#: Yield (content, writer) under an exclusive file lock.

nemo_rl.models.generation.vllm.patches._patch_vllm_init_workers_ray( py_executable: str, extra_env_vars: list[str] | None, ) → bool#

Patch vLLM’s Ray executor env propagation and worker runtime_env.

Pass custom runtime_env in _init_workers_ray call (file patch).
- This allows passing custom py_executable to worker initialization.
Forward extra env vars to the Ray workers via vLLM’s additive VLLM_RAY_EXTRA_ENV_VARS_TO_COPY hook (vLLM >= 0.25). NCCL_, HF_, and HUGGING_FACE_* vars are already copied by vLLM’s default prefix list (this includes the NCCL_CUMEM_ENABLE/NCCL_NVLS_ENABLE workaround from https://github.com/NVIDIA-NeMo/RL/pull/898).

.. note:: Step 1 patches the v1 Ray executor, which vLLM 0.25 no longer selects by default: VLLM_USE_RAY_V2_EXECUTOR_BACKEND flipped from "0" (0.20) to "1" (0.25), so Executor.get_class returns RayExecutorV2 for ray-backed engines. RayExecutorV2 has no _init_workers_ray at all – it creates workers inline, and its _build_runtime_env never sets py_executable.

The patch is kept because it is still load-bearing when
``VLLM_USE_RAY_V2_EXECUTOR_BACKEND=0`` selects the v1 executor. Under
the 0.25 default it is inert, and workers get the right interpreter
from Ray's per-field ``runtime_env`` inheritance instead: the parent
NeMo-RL actor sets ``py_executable``, and a child created with a
``runtime_env`` that omits it inherits the parent's value.

So a ``True`` return means "the anchor is in place", not "this is what
put the workers on the right interpreter". The caller logs
accordingly.

Returns:: Whether the v1 runtime_env source patch is in place. The env-var merge in step 2 cannot fail, but step 1 is anchored on a call-site string; if that moves upstream the py_executable injection silently stops happening, so the caller must not report success unconditionally.

nemo_rl.models.generation.vllm.patches._patch_vllm_llama_eagle3_own_lm_head(logger) → None#: Patch LlamaEagle3 to keep truncated draft lm_head ownership.

nemo_rl.models.generation.vllm.patches._patch_vllm_tool_parser_namespace_tool(logger) → None#

Guard vLLM’s NamespaceTool import for openai < 2.25.

vLLM 0.25 imports openai.types.responses.NamespaceTool (added in openai 2.25.0) at the top of tool_parsers/utils.py, but nemo-gym pins openai<=2.7.2 and its child server venvs must match the parent’s openai version exactly. NamespaceTool is only used in isinstance checks for Responses-API namespace tools, which cannot be constructed by an openai client that predates the feature, so a never-matching stub is a faithful fallback.

nemo_rl.models.generation.vllm.patches._patch_vllm_ray_executor_v2_tcpstore_port(logger) → None#

Keep RayExecutorV2’s TCPStore port out of the MessageQueue’s scan range.

vLLM 0.25’s RayExecutorV2._init_executor picks the torch.distributed TCPStore port with a bind-probe (Step 3) but only binds it much later, in the rank-0 worker’s init_process_group. In between, Step 4 builds the broadcast MessageQueue; when the engine spans nodes that queue needs a real TCP socket, so it calls get_open_port() and binds and holds the result (shm_broadcast.py: remote_subscribe_port = get_open_port() then remote_socket.bind(...)). Both searches start at VLLM_PORT, so the queue deterministically takes the very port the probe just released and engine startup dies with EADDRINUSE (DeepSeek-V3 generation TP=32, observed on port 7000). Engines that fit on one node use a shm/ipc socket instead and never allocate a TCP port here, which is why only node-spanning engines are affected.

Offsetting the TCPStore search past the queue’s scan range removes the collision while keeping both ports inside the engine’s 100-port window, and therefore below the OS ephemeral floor. That band is deliberate: leaving VLLM_PORT unset would send vLLM to kernel-assigned ephemeral ports and reintroduce the TOCTOU contention this layout exists to prevent (#2380, #3103).

The offset must be applied before the local_dp_rank is None test, not inside it. vLLM’s own disjoint-window branch below reads as if it only applies to DP engines, but ParallelConfig.__post_init__ takes the “offline SPMD” path for every engine NeMo-RL builds and assigns data_parallel_rank_local = envs.VLLM_DP_RANK_LOCAL (0 by default) and data_parallel_master_port = envs.VLLM_DP_MASTER_PORT (0 by default). So a plain non-DP engine arrives here with local_dp_rank=0, not None: the None branch is dead, and the DP branch searches from 0 + 100 + 0 * 32 = 100, fails all 32 attempts on the privileged range, and falls through to get_open_port() — straight back to VLLM_PORT. That is exactly the port the MessageQueue takes. See RL-1104.

Returns without raising when the snippet is missing, but logs at warning level so a silent no-op is visible in worker logs.

nemo_rl.models.generation.vllm.patches._patch_vllm_shm_broadcast_bind_retry(logger) → None#

Make MessageQueue’s remote socket survive losing a port race.

MessageQueue.__init__ picks the port for its remote (TCP) socket with remote_subscribe_port = get_open_port(), which probes a port and releases it, and only binds it with ZMQ several statements later (shm_broadcast.py: self.remote_socket.bind(socket_addr)). The window between the probe and the bind is a TOCTOU race.

On vLLM 0.25 that race is lost reliably, not occasionally. Every RayWorkerProc on a non-driver node takes n_local_reader=0 (ray_executor_v2.py::_init_message_queues), so every one of them needs a real TCP port, and they all scan from the same VLLM_PORT – 7000 for a node-spanning engine. _init_message_queues runs immediately after init_device(), whose process-group setup is a collective barrier, so all workers on the node arrive at the probe within microseconds of each other, all see the same port free, and all but one die with::

zmq.error.ZMQError: Address already in use (addr='tcp://10.65.1.9:7000')

Workers on the driver node take n_local_reader=1 and use an ipc:// socket instead, which is why only node-spanning engines are affected – and why no nightly test catches it (none runs an engine whose tensor_parallel_size * pipeline_parallel_size exceeds cluster.gpus_per_node). See RL-1111.

Fix the race at the bind rather than the probe: retry, advancing past the port that was lost. This is safe and terminating because a port a peer already holds with ZMQ is visible to the next _get_open_port probe (a plain bind(("", port)) on it fails with EADDRINUSE), so each retry makes forward progress.

Deliberately keeps the search anchored at VLLM_PORT instead of letting vLLM fall back to bind(("", 0)): kernel-assigned ephemeral ports are exactly the TOCTOU contention the reserved sub-ephemeral band exists to prevent (#2380, #3103).

Patching the bind (rather than handing each worker a private start port) also covers every other MessageQueue with a remote reader – notably the executor’s own rpc_broadcast_mq – instead of the one call site that happens to be failing today.

Returns without raising when the snippet is missing, but logs at warning level so a silent no-op is visible in worker logs.

nemo_rl.models.generation.vllm.patches._patch_vllm_radio_layerscale_loader(logger) → None#

Load explicit RADIO LayerScale weights and initialize folded weights.

vLLM 0.25.1 uses ls1 and ls2 in RadioVisionEncoderLayer but skips them in RadioModel.load_weights. Explicit checkpoint values are therefore ignored, while folded checkpoints leave the parameters at dummy initialization. Patch the loader so explicit values are loaded and absent values are initialized to RADIO’s configured identity factor.

nemo_rl.models.generation.vllm.patches._patch_vllm_glm_decoder_sequence_parallel_moe(logger) → None#

Restore the vLLM 0.24 decoder boundary for GLM DSA models.

vLLM 0.25.1 keeps hidden states sequence-parallel across attention and MoE decoder layers when TP, DP, and EP are all enabled. GLM-5.1/5.2 decode diverges on that new path: the first generated token is correct, while subsequent decode-token logprobs collapse. Keep vLLM’s existing MoE-local sequence parallelism, but disable the new decoder-level optimization for glm_moe_dsa so the MoE gathers its output as it did in vLLM 0.24.

The upstream bug and proposed fix are tracked at https://github.com/vllm-project/vllm/issues/50154 and https://github.com/vllm-project/vllm/pull/50155. Remove this patch after upgrading to a vLLM release containing the fix and validating iterative GLM-5.1/5.2 decode with TP, DP, and EP all enabled.

nemo_rl.models.generation.vllm.patches.ensure_vllm_source_compat() → None#

Apply interpreter-independent vLLM source-compat patches.

Safe to call from any process that imports vLLM directly (e.g. the tools/model_diagnostics scripts, which construct vllm.LLM without going through a NeMo-RL generation worker). Must be called BEFORE the first import vllm submodule that pulls in vllm.tool_parsers. Worker processes get this via _apply_vllm_patches at init.

nemo_rl.models.generation.vllm.patches._apply_vllm_patches( py_executable: str, *, extra_env_vars: list[str] | None = None, ) → None#

nemo_rl.models.generation.vllm.patches#

Module Contents#

Functions#

API#

`nemo_rl.models.generation.vllm.patches`#