From f7ef489e93cf92b8d6ce7403b49f1db867bcc35e Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sat, 8 Aug 2026 09:19:02 +0800 Subject: [PATCH] [Bugfix] Fix ZMQ port TOCTOU race in shm_broadcast MessageQueue (#50960) Signed-off-by: RTCartist Signed-off-by: aoshen02 Co-authored-by: RTCartist Co-authored-by: aoshen02 Co-authored-by: Claude Fable 5 --- tests/distributed/test_shm_broadcast.py | 65 +++++++++++++++++++ .../device_communicators/shm_broadcast.py | 7 +- 2 files changed, 68 insertions(+), 4 deletions(-) diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index 0b41396503..c3833b6b9f 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -713,3 +713,68 @@ def test_shm_ring_buffer_creation_checks_free_space(): pytest.raises(RuntimeError, match="Insufficient space"), ): ShmRingBuffer(n_reader=1, max_chunk_bytes=24 * 1024 * 1024, max_chunks=10) + + +def test_remote_subscribe_addr_unique_concurrent_writers( + monkeypatch: pytest.MonkeyPatch, +): + """Writers bind the remote socket to port 0 (kernel-assigned), so + concurrent writers never race for the same probed port and the + announced address is connectable. + + Pre-fix, the writer probed a port with get_open_port() and bound it + afterwards; pinning the probe to one free port makes every writer + bind the same port and fail deterministically on that code path, + while the late-binding implementation never consults the probe.""" + from vllm.distributed.device_communicators import shm_broadcast + + colliding_port = get_open_port() + monkeypatch.setattr( + shm_broadcast, "get_open_port", lambda: colliding_port, raising=False + ) + + n_writers = 32 + queues: list[MessageQueue] = [] + lock = threading.Lock() + + def make_writer(): + q = MessageQueue( + n_reader=1, + n_local_reader=0, + max_chunk_bytes=4096, + max_chunks=2, + connect_ip="127.0.0.1", + ) + with lock: + queues.append(q) + + threads = [threading.Thread(target=make_writer) for _ in range(n_writers)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(queues) == n_writers + addrs = [q.export_handle().remote_subscribe_addr for q in queues] + assert all(addr and addr.startswith("tcp://") for addr in addrs) + assert len(set(addrs)) == n_writers + + writer = queues[0] + received = [] + + def reader_main(): + reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0) + reader.wait_until_ready() + received.append(reader.dequeue()) + reader.remote_socket.close(linger=0) + + reader_thread = threading.Thread(target=reader_main) + reader_thread.start() + writer.wait_until_ready() + writer.enqueue("ping") + reader_thread.join(timeout=30) + assert not reader_thread.is_alive() + assert received == ["ping"] + + for q in queues: + q.remote_socket.close(linger=0) diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index f6d2a24755..6caacea6ec 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -36,7 +36,6 @@ from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.network_utils import ( get_ip, - get_open_port, get_open_zmq_inproc_path, get_open_zmq_ipc_path, is_valid_ipv6_address, @@ -525,13 +524,13 @@ class MessageQueue: connect_ip = get_ip() self.remote_socket = context.socket(XPUB) self.remote_socket.setsockopt(XPUB_VERBOSE, True) - remote_subscribe_port = get_open_port() if is_valid_ipv6_address(connect_ip): self.remote_socket.setsockopt(IPV6, 1) remote_addr_ipv6 = True connect_ip = f"[{connect_ip}]" - socket_addr = f"tcp://{connect_ip}:{remote_subscribe_port}" - self.remote_socket.bind(socket_addr) + self.remote_socket.bind(f"tcp://{connect_ip}:0") + last_endpoint = self.remote_socket.getsockopt(zmq.LAST_ENDPOINT) + remote_subscribe_port = last_endpoint.decode().rsplit(":", 1)[1] remote_subscribe_addr = f"tcp://{connect_ip}:{remote_subscribe_port}" else: remote_subscribe_addr = None