[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>
aoshen02 committed
Aug 8, 2026 at 09:19 UTC
f7ef489e93cf92b8d6ce7403b49f1db867bcc35e
2 files changed
+68
-4
tests/distributed/test_shm_broadcast.py
+65
@@ -713,3 +713,68 @@ def test_shm_ring_buffer_creation_checks_free_space():
713
pytest.raises(RuntimeError, match="Insufficient space"),
714
):
715
ShmRingBuffer(n_reader=1, max_chunk_bytes=24 * 1024 * 1024, max_chunks=10)
716
+
717
+
718
+def test_remote_subscribe_addr_unique_concurrent_writers(
719
+ monkeypatch: pytest.MonkeyPatch,
720
+):
721
+ """Writers bind the remote socket to port 0 (kernel-assigned), so
722
+ concurrent writers never race for the same probed port and the
723
+ announced address is connectable.
724
+
725
+ Pre-fix, the writer probed a port with get_open_port() and bound it
726
+ afterwards; pinning the probe to one free port makes every writer
727
+ bind the same port and fail deterministically on that code path,
728
+ while the late-binding implementation never consults the probe."""
729
+ from vllm.distributed.device_communicators import shm_broadcast
730
+
731
+ colliding_port = get_open_port()
732
+ monkeypatch.setattr(
733
+ shm_broadcast, "get_open_port", lambda: colliding_port, raising=False
734
+ )
735
+
736
+ n_writers = 32
737
+ queues: list[MessageQueue] = []
738
+ lock = threading.Lock()
739
+
740
+ def make_writer():
741
+ q = MessageQueue(
742
+ n_reader=1,
743
+ n_local_reader=0,
744
+ max_chunk_bytes=4096,
745
+ max_chunks=2,
746
+ connect_ip="127.0.0.1",
747
+ )
748
+ with lock:
749
+ queues.append(q)
750
+
751
+ threads = [threading.Thread(target=make_writer) for _ in range(n_writers)]
752
+ for t in threads:
753
+ t.start()
754
+ for t in threads:
755
+ t.join()
756
+
757
+ assert len(queues) == n_writers
758
+ addrs = [q.export_handle().remote_subscribe_addr for q in queues]
759
+ assert all(addr and addr.startswith("tcp://") for addr in addrs)
760
+ assert len(set(addrs)) == n_writers
761
+
762
+ writer = queues[0]
763
+ received = []
764
+
765
+ def reader_main():
766
+ reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0)
767
+ reader.wait_until_ready()
768
+ received.append(reader.dequeue())
769
+ reader.remote_socket.close(linger=0)
770
+
771
+ reader_thread = threading.Thread(target=reader_main)
772
+ reader_thread.start()
773
+ writer.wait_until_ready()
774
+ writer.enqueue("ping")
775
+ reader_thread.join(timeout=30)
776
+ assert not reader_thread.is_alive()
777
+ assert received == ["ping"]
778
+
779
+ for q in queues:
780
+ q.remote_socket.close(linger=0)
vllm/distributed/device_communicators/shm_broadcast.py
+3
-4
@@ -36,7 +36,6 @@ from vllm.logger import init_logger
36
from vllm.platforms import current_platform
37
from vllm.utils.network_utils import (
38
get_ip,
39
- get_open_port,
39
get_open_zmq_inproc_path,
40
get_open_zmq_ipc_path,
41
is_valid_ipv6_address,
@@ -525,13 +524,13 @@ class MessageQueue:
524
connect_ip = get_ip()
525
self.remote_socket = context.socket(XPUB)
526
self.remote_socket.setsockopt(XPUB_VERBOSE, True)
528
- remote_subscribe_port = get_open_port()
527
if is_valid_ipv6_address(connect_ip):
528
self.remote_socket.setsockopt(IPV6, 1)
529
remote_addr_ipv6 = True
530
connect_ip = f"[{connect_ip}]"
533
- socket_addr = f"tcp://{connect_ip}:{remote_subscribe_port}"
534
- self.remote_socket.bind(socket_addr)
531
+ self.remote_socket.bind(f"tcp://{connect_ip}:0")
532
+ last_endpoint = self.remote_socket.getsockopt(zmq.LAST_ENDPOINT)
533
+ remote_subscribe_port = last_endpoint.decode().rsplit(":", 1)[1]
534
remote_subscribe_addr = f"tcp://{connect_ip}:{remote_subscribe_port}"
535
else:
536
remote_subscribe_addr = None