[Bugfix] Fix TP=1 PLE CPU-offload startup rendezvous deadlock at kernel warmup

vllm/v1/engine/utils.py

@@ -4,6 +4,7 @@
 import contextlib
 import os
 import threading
+import time
 import weakref
 from collections.abc import Callable, Iterator, Sequence
 from dataclasses import dataclass
@@ -65,6 +66,17 @@ def get_engine_process_shutdown_timeout(
         return ROCM_ENGINE_PROCESS_SHUTDOWN_TIMEOUT_S
     return process_timeout
 
+# Wall-clock backstop for a LIVE-but-hung EngineCore. wait_for_engine_startup()
+# below only raises when a core proc EXITS (a sentinel fires); a proc that is
+# alive but never sends its ready message -- e.g. the TP=1 PLE CPU/disk-offload
+# warmup rendezvous deadlock (vllm-project/vllm#53960) -- would otherwise loop
+# here forever. This bound sits WELL ABOVE a full cold boot (~13-14 min here) so
+# it never pre-empts a slow-but-healthy start; it converts an infinite wait into
+# a NAMED TimeoutError. Override with VLLM_ENGINE_CORE_STARTUP_TIMEOUT (seconds).
+ENGINE_CORE_STARTUP_TIMEOUT_S = float(
+    os.getenv("VLLM_ENGINE_CORE_STARTUP_TIMEOUT", "1800")
+)
+
 
 class CoreEngineState(Enum):
     NEW = auto()
@@ -1283,6 +1295,7 @@ def wait_for_engine_startup(
         frontend_process_by_fd[fd] = proc
         poller.register(fd, zmq.POLLIN)
 
+    startup_deadline = time.monotonic() + ENGINE_CORE_STARTUP_TIMEOUT_S
     while any(conn_pending) or any(start_pending):
         events = poller.poll(STARTUP_POLL_PERIOD_MS)
         if not events:
@@ -1296,6 +1309,24 @@ def wait_for_engine_startup(
                     "Waiting for %d local, %d remote core engine proc(s) to start.",
                     *start_pending,
                 )
+            # A proc EXIT is handled below (a sentinel fires and produces an
+            # event). Reaching the deadline with no event means the proc(s) are
+            # ALIVE but not ready -- a live hang, not an exit -- so name it
+            # rather than loop forever (vllm-project/vllm#53960).
+            if time.monotonic() >= startup_deadline:
+                raise TimeoutError(
+                    "EngineCore startup exceeded "
+                    f"{ENGINE_CORE_STARTUP_TIMEOUT_S:.0f}s with "
+                    f"{conn_pending[0]} local + {conn_pending[1]} remote proc(s) "
+                    f"still to connect and {start_pending[0]} local + "
+                    f"{start_pending[1]} remote proc(s) still to start. The core "
+                    "engine process(es) are ALIVE but have not sent a ready "
+                    "message -- this is a live hang, distinct from a process "
+                    "exit (reported separately with the exit code). A PLE "
+                    "CPU/disk-offload warmup rendezvous deadlock on TP=1 "
+                    "(vllm-project/vllm#53960) is one known cause. Adjust with "
+                    "VLLM_ENGINE_CORE_STARTUP_TIMEOUT."
+                )
             continue
         if len(events) > 1 or events[0][0] != handshake_socket:
             # One of the local core, coordinator, or watched frontend processes exited.

vllm/v1/ple_offload/connector.py

@@ -5,6 +5,7 @@
 import os
 import queue
 import threading
+import time
 from dataclasses import dataclass
 from multiprocessing.reduction import ForkingPickler
 from typing import Any
@@ -23,10 +24,16 @@
     PleOffloadLayer,
 )
 from vllm.v1.ple_offload.protocol import (
+    _PLE_OFFLOAD_ACK_DECODER,
     PleOffloadRegistration,
     PleOffloadRequest,
+    barrier_timeout_s,
 )
 
+# Poll granularity for the startup ACK wait. The wait itself is bounded by
+# barrier_timeout_s(); this only controls how often the deadline is rechecked.
+_ACK_POLL_MS = 1000
+
 logger = init_logger(__name__)
 
 
@@ -66,6 +73,13 @@ def __init__(
         self.device = device
         self.dp_rank = get_dp_group().rank_in_group
         self.tp_rank = get_tp_group().rank_in_group
+        # Stable per-worker id, matching the offload worker's expectation, plus
+        # the private endpoint the offload worker returns this worker's ACK on.
+        self._worker_id = (
+            self.dp_rank * vllm_config.parallel_config.world_size
+            + vllm_config.parallel_config.rank
+        )
+        self._ack_addr = f"{ipc_addr}.ack.{self._worker_id}"
         self._layers = self._setup_layers(vllm_config, model)
 
         # Both runner paths stage into the same shared buffers. TP0 registers
@@ -111,14 +125,27 @@ def __init__(
         self._request_thread_ready = threading.Event()
         self._zmq_ctx: zmq.Context | None = None
         self._registration_socket: zmq.Socket | None = None
+        self._ack_socket: zmq.Socket | None = None
         self._d2h_event_pool: queue.Queue[torch.cuda.Event] | None = None
 
         try:
             self._zmq_ctx = zmq.Context()
+            # Bind the ACK endpoint BEFORE registering so the offload worker's
+            # later PUSH always has a bound peer to reach.
+            self._ack_socket = self._zmq_ctx.socket(zmq.PULL)
+            self._ack_socket.bind(self._ack_addr)
             self._registration_socket = self._zmq_ctx.socket(zmq.PUSH)
             self._registration_socket.connect(ipc_addr)
             self._register_with_offload_worker(vllm_config, ipc_addr)
 
+            # STARTUP RENDEZVOUS BARRIER (vllm-project/vllm#53960): block here
+            # until the offload worker acknowledges this registration. This
+            # orders offload-ready BEFORE the first warmup forward (which runs
+            # only after this constructor returns) and turns a lost registration
+            # into a NAMED failure instead of an infinite GPU-stream wait. It is
+            # a one-time STARTUP wait; the steady-state decode path is untouched.
+            self._await_registration_ack(ipc_addr)
+
             if self.tp_rank == 0:
                 # ForkingPickler may replace CPU storage while converting its
                 # sharing strategy, so register only the final addresses.
@@ -211,10 +238,7 @@ def _register_with_offload_worker(
         # Each GPU worker owns distinct output buffers, while TP0's shared
         # inputs become the request source for its DP rank.
         registration = PleOffloadRegistration(
-            worker_id=(
-                self.dp_rank * vllm_config.parallel_config.world_size
-                + vllm_config.parallel_config.rank
-            ),
+            worker_id=self._worker_id,
             tp_rank=self.tp_rank,
             dp_rank=self.dp_rank,
             gpu_output_buffers={
@@ -226,6 +250,7 @@ def _register_with_offload_worker(
             input_ids_buf=self._input_ids_buf,
             query_start_loc_buf=self._query_start_loc_buf,
             ngram_context_buf=self._ngram_context_buf,
+            ack_addr=self._ack_addr,
         )
 
         # ForkingPickler transmits tensors through shared-memory and CUDA IPC.
@@ -250,6 +275,49 @@ def _register_with_offload_worker(
             sorted(self._layers),
         )
 
+    def _await_registration_ack(self, ipc_addr: str) -> None:
+        """Block until the offload worker acknowledges this registration.
+
+        This is the STARTUP rendezvous barrier. It is bounded WELL ABOVE a
+        legitimate cold boot: a healthy offload worker only sends the ACK once
+        it has received every GPU worker's registration and built their output
+        targets, so a slow sibling can legitimately delay it, but an infinite
+        wait means a registration was lost, the offload worker died, or it is
+        itself hung -- all of which are named here rather than left to surface
+        as an untimed ``cuStreamWaitValue32`` hang in the first warmup forward
+        (vllm-project/vllm#53960).
+        """
+        assert self._ack_socket is not None
+        timeout_s = barrier_timeout_s()
+        deadline = time.monotonic() + timeout_s
+        while True:
+            if self._ack_socket.poll(timeout=_ACK_POLL_MS):
+                ack = _PLE_OFFLOAD_ACK_DECODER.decode(self._ack_socket.recv())
+                logger.info(
+                    "PleOffload: registration ACK received "
+                    "(worker_id=%d, dp_rank=%d, tp_rank=%d, layers=%d).",
+                    ack.worker_id,
+                    self.dp_rank,
+                    self.tp_rank,
+                    ack.num_layers,
+                )
+                return
+            if time.monotonic() >= deadline:
+                raise TimeoutError(
+                    "PLE offload registration was not acknowledged within "
+                    f"{timeout_s:.0f}s. "
+                    "GPU-worker state: registration sent to the offload PULL "
+                    f"socket at {ipc_addr}, then blocked awaiting the ACK on "
+                    f"{self._ack_addr} (dp_rank={self.dp_rank}, "
+                    f"tp_rank={self.tp_rank}, worker_id={self._worker_id}). "
+                    "Offload-worker state: has NOT sent this ACK -- it is still "
+                    "collecting the expected registrations, has died during "
+                    "startup, or this registration message was lost in transit. "
+                    "This is the TP=1 PLE-offload warmup rendezvous deadlock "
+                    "(vllm-project/vllm#53960): a boot that reaches this bound "
+                    "is hung, not merely slow."
+                )
+
     def _start_request_thread(self, ipc_addr: str) -> None:
         """Start the thread that publishes batches after inputs are ready."""
         self._request_thread = threading.Thread(
@@ -469,6 +537,9 @@ def close(self) -> None:
         if self._registration_socket is not None:
             self._registration_socket.close(linger=0)
             self._registration_socket = None
+        if self._ack_socket is not None:
+            self._ack_socket.close(linger=0)
+            self._ack_socket = None
         if self._zmq_ctx is not None:
             self._zmq_ctx.term()
             self._zmq_ctx = None

vllm/v1/ple_offload/protocol.py

@@ -7,6 +7,27 @@
 import msgspec
 import torch
 
+import vllm.envs as envs
+
+# ---------------------------------------------------------------------------
+# Startup rendezvous barrier timeout
+# ---------------------------------------------------------------------------
+# The GPU worker <-> offload registration handshake (connector.py sends a
+# registration, worker.py accepts it and returns an ACK) is a STARTUP-ONLY
+# rendezvous. Bounding both ends converts the TP=1 warmup deadlock
+# (vllm-project/vllm#53960) -- where a lost registration or a hung sibling
+# leaves one side waiting forever -- into a NAMED TimeoutError. The bound is
+# deliberately WELL ABOVE a legitimate cold boot (main weight load alone is
+# ~510 s and a full boot ~13-14 min here), so it never pre-empts a slow-but-
+# healthy start; it only fires on a true hang. It is derived from the existing
+# ready-timeout knob so raising that knob raises this bound too.
+
+
+def barrier_timeout_s() -> float:
+    """Return the startup registration-rendezvous timeout in seconds."""
+    return max(2.0 * float(envs.VLLM_PLE_OFFLOAD_READY_TIMEOUT), 1200.0)
+
+
 # ---------------------------------------------------------------------------
 # IPC message dataclasses
 # ---------------------------------------------------------------------------
@@ -26,6 +47,10 @@ class PleOffloadRegistration:
     input_ids_buf: torch.Tensor
     query_start_loc_buf: torch.Tensor
     ngram_context_buf: torch.Tensor | None
+    # ZMQ endpoint the offload worker PUSHes this worker's registration ACK to.
+    # The GPU worker binds it before sending and blocks on the ACK as the
+    # startup rendezvous barrier. Empty disables the ACK path (legacy/tests).
+    ack_addr: str = ""
 
 
 @dataclass
@@ -37,4 +62,20 @@ class PleOffloadRequest:
     num_reqs: int
 
 
+@dataclass
+class PleOffloadRegistrationAck:
+    """Sent by the offload worker to one GPU worker once that worker's
+    registration has been received AND its output targets are built.
+
+    Receipt of this ACK is the startup rendezvous barrier: the GPU worker must
+    not dispatch its first (warmup) PLE forward -- which enqueues an untimed
+    ``cuStreamWaitValue32`` on the model stream -- until the offload worker is
+    known to be serving it. A missing ACK is detected (TimeoutError), never an
+    infinite GPU-stream wait (vllm-project/vllm#53960)."""
+
+    worker_id: int
+    num_layers: int
+
+
 _PLE_OFFLOAD_REQUEST_DECODER = msgspec.msgpack.Decoder(PleOffloadRequest)
+_PLE_OFFLOAD_ACK_DECODER = msgspec.msgpack.Decoder(PleOffloadRegistrationAck)