NVIDIA-NeMo / NVIDIA-NeMo/RL

mcore inference HTTP server port pick is TOCTOU-prone; mirror vLLM's held-socket reservation

Open
#3,775 3 comments 0 reactions 1 assignee Assigned to @tdene View on GitHub
Automation bug
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

The mcore inference HTTP server picks its port by scanning, releasing, and later re-binding. Another process on the node can take the port in that gap. vLLM already solved this; mcore should do the same.

Pre-existing — not introduced by any particular PR. Surfaced while reviewing #3731.

## The gap

Permalinks pinned at `dfdf6894d5380f13f8b9c285315aa0f2ee99dafb`.

1. [`megatron_worker.py:317`](https://github.com/NVIDIA-NeMo/RL/blob/dfdf6894d5380f13f8b9c285315aa0f2ee99dafb/nemo_rl/models/generation/megatron/megatron_worker.py#L317) calls `_get_free_port_local()`.
2. That helper binds, listens, then **closes** the socket and returns only the integer — the `with` block at [`virtual_cluster.py:235-245`](https://github.com/NVIDIA-NeMo/RL/blob/dfdf6894d5380f13f8b9c285315aa0f2ee99dafb/nemo_rl/distributed/virtual_cluster.py#L235-L245) exits before `return port`.
3. The integer is passed as [`server_port=free_port`](https://github.com/NVIDIA-NeMo/RL/blob/dfdf6894d5380f13f8b9c285315aa0f2ee99dafb/nemo_rl/models/generation/megatron/megatron_worker.py#L323).
4. Upstream re-binds it at [`text_generation_server.py:184`](https://github.com/NVIDIA/Megatron-LM/blob/14346b65a2d0790e451919858f7771078105c5f0/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py#L184).

Between 2 and 4 the port is free for anyone to take.

The pool is `[1400, 1999)` (`DEFAULT_MASTER_PORT_RANGE_*`, [`virtual_cluster.py:128-129`](https://github.com/NVIDIA-NeMo/RL/blob/dfdf6894d5380f13f8b9c285315aa0f2ee99dafb/nemo_rl/distributed/virtual_cluster.py#L128-L129)) — shared with Ray master ports, so it is not private to the HTTP server.

## Failure modes

- **Usual:** an uncaught `OSError: [Errno 98] Address already in use` from the re-bind at step 4.
- **Quieter and worse:** upstream sets `SO_REUSEPORT` when available ([`text_generation_server.py:177-181`](https://github.com/NVIDIA/Megatron-LM/blob/14346b65a2d0790e451919858f7771078105c5f0/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py#L177-L181)). If the process that took the port also set `SO_REUSEPORT` under the same UID, **both sockets bind and the kernel splits incoming connections between them** — no error, just some requests going to the wrong server.

Only global rank 0 starts a server ([`megatron_worker.py:362-366`](https://github.com/NVIDIA-NeMo/RL/blob/dfdf6894d5380f13f8b9c285315aa0f2ee99dafb/nemo_rl/models/generation/megatron/megatron_worker.py#L362-L366)), so there is no contention between mcore servers inside one job. The exposure is co-tenants on the same node — which is exactly the contention vLLM hit.

## How vLLM solved it

vLLM holds the socket open from reservation through serving, so there is no gap:

- [`vllm_worker_async.py:132-153`](https://github.com/NVIDIA-NeMo/RL/blob/dfdf6894d5380f13f8b9c285315aa0f2ee99dafb/nemo_rl/models/generation/vllm/vllm_worker_async.py#L132-L153) `_reserve_port()` binds and keeps a listening socket on the worker.
- [`vllm_worker_async.py:921-928`](https://github.com/NVIDIA-NeMo/RL/blob/dfdf6894d5380f13f8b9c285315aa0f2ee99dafb/nemo_rl/models/generation/vllm/vllm_worker_async.py#L921-L928) hands that same socket to uvicorn via `sockets=`.

Its docstring states the invariant directly: the socket *"is never closed and re-opened, so there is zero gap where another process could steal the port."*

Related prior work on vLLM port contention: #2380, #3103.

## Suggested fix

Upstream already supports the handoff — no Megatron-LM change needed. [`start_text_gen_server`](https://github.com/NVIDIA/Megatron-LM/blob/14346b65a2d0790e451919858f7771078105c5f0/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py#L144-L153) accepts `sock: Optional[socket.socket]`, documented at [L163](https://github.com/NVIDIA/Megatron-LM/blob/14346b65a2d0790e451919858f7771078105c5f0/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py#L163) as *"The caller may pass in a socket it has already bound ahead of time."* It validates the socket is bound to a real port and reuses it.

So in `_setup_openai_api_server`: bind and hold the socket (as `_reserve_port` does), then pass it through as `sock=` instead of `server_port=`.

There is a second benefit. Doing the reservation at construction time also makes the URL available before the engine is loaded, which is how vLLM populates `dp_openai_server_base_urls` under `defer_model_load` ([`vllm_generation.py:276-279`](https://github.com/NVIDIA-NeMo/RL/blob/dfdf6894d5380f13f8b9c285315aa0f2ee99dafb/nemo_rl/models/generation/vllm/vllm_generation.py#L276-L279)). That is the same change #3731 needs for its deferred-init path, so the two can land together.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.