mcore inference HTTP server port pick is TOCTOU-prone; mirror vLLM's held-socket reservation
- 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
Assessment
This issue has not been assessed yet.