101 lines
3.9 KiB
Python
101 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import time
|
|
from contextlib import contextmanager
|
|
|
|
from vllm.distributed.device_communicators import shm_broadcast
|
|
|
|
MessageQueue = shm_broadcast.MessageQueue
|
|
|
|
# Cap on how long an idle reader parks before re-reading the authoritative SHM
|
|
# written-flag. Bounds lost-notify recovery latency to ~5s while the periodic
|
|
# wakeup stays negligible (one flag check per reader every 5s).
|
|
SHM_READER_RECHECK_INTERVAL_MS = 5000
|
|
|
|
|
|
def timeout_ms(self) -> int:
|
|
"""Returns a timeout, capped at the recheck interval, that is:
|
|
- min(time to deadline, time to next warning) if we're logging warnings
|
|
- time to deadline, if we're not logging warnings
|
|
- recheck interval if the timeout is None and we're not logging warnings
|
|
- raise TimeoutError if we are past the deadline
|
|
"""
|
|
wait_ms = SHM_READER_RECHECK_INTERVAL_MS
|
|
if self.warning_wait_time_ms is not None:
|
|
wait_ms = min(wait_ms, self.warning_wait_time_ms)
|
|
if self.timeout is None:
|
|
return wait_ms
|
|
time_left_ms = int((self.deadline - time.monotonic()) * 1000)
|
|
if time_left_ms <= 0:
|
|
raise TimeoutError
|
|
return min(wait_ms, time_left_ms)
|
|
|
|
|
|
@contextmanager
|
|
def acquire_read(
|
|
self,
|
|
timeout: float | None = None,
|
|
indefinite: bool = False,
|
|
):
|
|
assert self._is_local_reader, "Only readers can acquire read"
|
|
read_timeout = self.ReadTimeoutWithWarnings(timeout=timeout, should_warn=not indefinite)
|
|
with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
|
|
while True:
|
|
|
|
def check():
|
|
shm_broadcast.memory_fence()
|
|
read_flag = metadata_buffer[self.local_reader_rank + 1]
|
|
written_flag = metadata_buffer[0]
|
|
return not (not written_flag or read_flag)
|
|
|
|
if shm_broadcast.SPINLOOP_EXT_ENABLED and not check():
|
|
shm_broadcast.spinloop(
|
|
metadata_buffer[0 : self.local_reader_rank + 1],
|
|
check,
|
|
timeout=shm_broadcast.SPINLOOP_TIMEOUT_SECONDS,
|
|
)
|
|
|
|
if not check():
|
|
# this block is either
|
|
# (1) not written
|
|
# (2) already read by this reader
|
|
|
|
# for readers, `self.current_idx` is the next block to read
|
|
# if this block is not ready,
|
|
# we need to wait until it is written
|
|
self._spin_condition.wait(timeout_ms=read_timeout.timeout_ms())
|
|
|
|
if self.shutting_down:
|
|
raise RuntimeError("cancelled")
|
|
|
|
# if we wait for a long time, log a message
|
|
if read_timeout.should_warn():
|
|
shm_broadcast.logger.info(
|
|
shm_broadcast.LONG_WAIT_TIME_LOG_MSG,
|
|
shm_broadcast.VLLM_RINGBUFFER_WARNING_INTERVAL,
|
|
)
|
|
|
|
continue
|
|
|
|
# found a block that is not read by this reader
|
|
# let caller read from the buffer
|
|
with self.buffer.get_data(self.current_idx) as buf:
|
|
try:
|
|
yield buf
|
|
finally:
|
|
# caller has read from the buffer; set the read flag.
|
|
metadata_buffer[self.local_reader_rank + 1] = 1
|
|
# Memory fence ensures the read flag is visible to the writer.
|
|
# Without this, writer may not see our read completion and
|
|
# could wait indefinitely for all readers to finish.
|
|
shm_broadcast.memory_fence()
|
|
next_idx = self.current_idx + 1
|
|
self.current_idx = next_idx % self.buffer.max_chunks
|
|
self._spin_condition.record_read()
|
|
break
|
|
|
|
|
|
MessageQueue.ReadTimeoutWithWarnings.timeout_ms = timeout_ms
|
|
MessageQueue.acquire_read = acquire_read
|