[Bugfix] Fix ZMQ port TOCTOU race in shm_broadcast MessageQueue (#50960)

Signed-off-by: RTCartist <wangshengb@buaa.edu.cn>
Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: RTCartist <wangshengb@buaa.edu.cn>
Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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)

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