[Bugfix][Core] shm_broadcast: bound idle reader waits and release read slots (#45224)
Signed-off-by: Chaemin Lim <chaemin.lim@mangoboost.io> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Edwin Lim <edwin.lim@mangoboost.io> Co-authored-by: Jaeyoun Kim <jaeyoun.kim@mangoboost.io> Co-authored-by: Nick Hill <nickhill123@gmail.com>
chaeminlim-mb committed
Jul 23, 2026 at 19:13 UTC
10c75477b07c2f1a361f54b7357af1019bba5fd8
2 files changed
+155
-23
tests/distributed/test_shm_broadcast.py
+130
@@ -348,6 +348,136 @@ def test_message_queue_busy_to_idle():
348
distributed_run(worker_fn_test_busy_to_idle, 4)
349
350
351
+@pytest.mark.parametrize("should_warn", [False, True])
352
+def test_reader_timeout_caps_indefinite_waits(should_warn):
353
+ with (
354
+ mock.patch(
355
+ "vllm.distributed.device_communicators.shm_broadcast."
356
+ "SHM_READER_RECHECK_INTERVAL_MS",
357
+ new=7,
358
+ ),
359
+ mock.patch(
360
+ "vllm.distributed.device_communicators.shm_broadcast."
361
+ "VLLM_RINGBUFFER_WARNING_INTERVAL",
362
+ new=60,
363
+ ),
364
+ ):
365
+ timeout = MessageQueue.ReadTimeoutWithWarnings(
366
+ timeout=None, should_warn=should_warn
367
+ )
368
+ assert timeout.timeout_ms() == 7
369
+
370
+
371
+def test_reader_rechecks_shm_after_idle_wait_timeout_without_notify():
372
+ writer = MessageQueue(
373
+ n_reader=1,
374
+ n_local_reader=1,
375
+ max_chunk_bytes=1024 * 1024,
376
+ max_chunks=1,
377
+ )
378
+ reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0)
379
+ payload = 123
380
+ poll_started = threading.Event()
381
+ allow_timeout = threading.Event()
382
+ result = {}
383
+
384
+ def acquire_read_in_thread():
385
+ try:
386
+ with reader.acquire_read(indefinite=True) as buf:
387
+ result["value"] = buf[0]
388
+ except Exception as exc:
389
+ result["exc"] = exc
390
+
391
+ def poll_timeout(*, timeout: int | None = None):
392
+ poll_started.set()
393
+ assert allow_timeout.wait(timeout=5)
394
+ return []
395
+
396
+ try:
397
+ writer.wait_until_ready()
398
+ reader.wait_until_ready()
399
+ reader._spin_condition.last_read = 0
400
+ reader._spin_condition.busy_loop_s = 0
401
+
402
+ with (
403
+ mock.patch(
404
+ "vllm.distributed.device_communicators.shm_broadcast."
405
+ "SHM_READER_RECHECK_INTERVAL_MS",
406
+ new=50,
407
+ ),
408
+ mock.patch(
409
+ "vllm.distributed.device_communicators.shm_broadcast."
410
+ "VLLM_RINGBUFFER_WARNING_INTERVAL",
411
+ new=60,
412
+ ),
413
+ mock.patch.object(
414
+ reader._spin_condition.poller,
415
+ "poll",
416
+ side_effect=poll_timeout,
417
+ ) as poll,
418
+ ):
419
+ read_thread = threading.Thread(target=acquire_read_in_thread, daemon=True)
420
+ read_thread.start()
421
+ assert poll_started.wait(timeout=5)
422
+ with writer.acquire_write(timeout=0.1) as buf:
423
+ buf[0] = payload
424
+ allow_timeout.set()
425
+ read_thread.join(timeout=5)
426
+
427
+ assert not read_thread.is_alive()
428
+ poll.assert_called_once_with(timeout=50)
429
+
430
+ if "exc" in result:
431
+ raise result["exc"]
432
+ assert result["value"] == payload
433
+ with writer.buffer.get_metadata(0) as metadata_buffer:
434
+ assert metadata_buffer[0] == 1
435
+ assert metadata_buffer[1] == 1
436
+ finally:
437
+ writer.shutdown()
438
+ reader.shutdown()
439
+ for socket in (
440
+ writer.local_socket,
441
+ writer._spin_condition.local_notify_socket,
442
+ reader.local_socket,
443
+ reader._spin_condition.local_notify_socket,
444
+ reader._spin_condition.read_cancel_socket,
445
+ reader._spin_condition.write_cancel_socket,
446
+ ):
447
+ socket.close(linger=0)
448
+
449
+
450
+def test_acquire_read_releases_slot_when_reader_raises():
451
+ writer = MessageQueue(
452
+ n_reader=1,
453
+ n_local_reader=1,
454
+ max_chunk_bytes=1024 * 1024,
455
+ max_chunks=1,
456
+ )
457
+ reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0)
458
+ try:
459
+ writer.wait_until_ready()
460
+ reader.wait_until_ready()
461
+
462
+ writer.enqueue({"payload": "first"})
463
+
464
+ with (
465
+ pytest.raises(RuntimeError, match="reader failed"),
466
+ reader.acquire_read(timeout=0.1),
467
+ ):
468
+ raise RuntimeError("reader failed")
469
+
470
+ with writer.buffer.get_metadata(0) as metadata_buffer:
471
+ assert metadata_buffer[0] == 1
472
+ assert metadata_buffer[1] == 1
473
+
474
+ with writer.acquire_write(timeout=0.1) as buf:
475
+ buf[0] = 0
476
+ finally:
477
+ writer.shutdown()
478
+ reader.shutdown()
479
+
480
+
481
def test_warning_logs(caplog_vllm):
482
"""
483
Test that warning logs are emitted at VLLM_RINGBUFFER_WARNING_INTERVAL intervals
vllm/distributed/device_communicators/shm_broadcast.py
+25
-23
@@ -57,6 +57,11 @@ if TYPE_CHECKING:
57
from _typeshed import SizedBuffer
58
59
VLLM_RINGBUFFER_WARNING_INTERVAL = envs.VLLM_RINGBUFFER_WARNING_INTERVAL
60
+# Cap on how long an idle reader parks before re-reading the authoritative SHM
61
+# written-flag. Bounds lost-notify recovery latency to ~5s while the periodic
62
+# wakeup stays negligible (one flag check per reader every 5s).
63
+SHM_READER_RECHECK_INTERVAL_MS = 5000
64
+
65
66
from_bytes_big = functools.partial(int.from_bytes, byteorder="big")
67
@@ -631,25 +636,22 @@ class MessageQueue:
636
self.n_warning = 1
637
self.timeout = timeout
638
634
- def timeout_ms(self) -> int | None:
635
- """Returns a timeout that is:
639
+ def timeout_ms(self) -> int:
640
+ """Returns a timeout, capped at the recheck interval, that is:
641
- min(time to deadline, time to next warning) if we're logging warnings
642
- time to deadline, if we're not logging warnings
638
- - None if the timeout is None and we're not logging warnings
643
+ - recheck interval if the timeout is None and we're not logging warnings
644
- raise TimeoutError if we are past the deadline
645
"""
641
- warning_wait_time = self.warning_wait_time_ms
646
+ wait_ms = SHM_READER_RECHECK_INTERVAL_MS
647
+ if self.warning_wait_time_ms is not None:
648
+ wait_ms = min(wait_ms, self.warning_wait_time_ms)
649
if self.timeout is None:
643
- return warning_wait_time
644
-
650
+ return wait_ms
651
time_left_ms = int((self.deadline - time.monotonic()) * 1000)
652
if time_left_ms <= 0:
653
raise TimeoutError
648
-
649
- if warning_wait_time and warning_wait_time < time_left_ms:
650
- return warning_wait_time
651
-
652
- return time_left_ms
654
+ return min(wait_ms, time_left_ms)
655
656
def should_warn(self) -> bool:
657
"""Returns true if it's time to log a warning for a timeout that is not
@@ -710,18 +712,18 @@ class MessageQueue:
712
# found a block that is not read by this reader
713
# let caller read from the buffer
714
with self.buffer.get_data(self.current_idx) as buf:
713
- yield buf
714
-
715
- # caller has read from the buffer
716
- # set the read flag
717
- metadata_buffer[self.local_reader_rank + 1] = 1
718
- # Memory fence ensures the read flag is visible to the writer.
719
- # Without this, writer may not see our read completion and
720
- # could wait indefinitely for all readers to finish.
721
- memory_fence()
722
- self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks
723
-
724
- self._spin_condition.record_read()
715
+ try:
716
+ yield buf
717
+ finally:
718
+ # caller has read from the buffer; set the read flag.
719
+ metadata_buffer[self.local_reader_rank + 1] = 1
720
+ # Memory fence ensures the read flag is visible to the writer.
721
+ # Without this, writer may not see our read completion and
722
+ # could wait indefinitely for all readers to finish.
723
+ memory_fence()
724
+ next_idx = self.current_idx + 1
725
+ self.current_idx = next_idx % self.buffer.max_chunks
726
+ self._spin_condition.record_read()
727
break
728
729
def enqueue(self, obj, timeout: float | None = None):