[TRTLLM-13409][feat] proxy fast-death detection + sticky EngineDeadError

tensorrt_llm/executor/__init__.py

@@ -3,7 +3,7 @@
 from .proxy import *
 from .request import *
 from .result import *
-from .utils import RequestError
+from .utils import EngineDeadError, RequestError
 from .worker import *
 
 __all__ = [
@@ -15,6 +15,7 @@
     "GenerationExecutorWorker",
     "GenerationExecutorProxy",
     "RequestError",
+    "EngineDeadError",
     "CompletionOutput",
     "GenerationResultBase",
     "DetokenizedGenerationResultBase",

tensorrt_llm/executor/proxy.py

@@ -40,9 +40,10 @@
 from .result import GenerationResult, IterationResult
 from .rpc import RPCClient
 from .rpc.rpc_common import RPCError, get_unique_ipc_addr
-from .utils import (ErrorResponse, RequestError, WorkerCommIpcAddrs,
-                    create_mpi_comm_session, get_spawn_proxy_process_env,
-                    is_llm_response, print_alive_threads)
+from .utils import (EngineDeadError, ErrorResponse, RequestError,
+                    WorkerCommIpcAddrs, create_mpi_comm_session,
+                    get_spawn_proxy_process_env, is_llm_response,
+                    print_alive_threads)
 from .worker import GenerationExecutorWorker, worker_main
 
 __all__ = [
@@ -121,6 +122,10 @@ def __init__(
                 "yellow")
 
         self._results: Dict[int, GenerationResult] = {}
+        # Sticky engine-dead flag: once a worker death is recorded, pending and
+        # new requests fail fast with EngineDeadError instead of blocking on a
+        # response queue whose producer is gone.
+        self._engine_dead = False
 
         self.model_world_size = model_world_size
 
@@ -242,15 +247,87 @@ def check_health(self) -> bool:
 
         return True
 
+    def _set_fatal_error(self, error: BaseException) -> None:
+        """Record the fatal error, then unblock every pending request.
+
+        Extends the base behavior so that the instant a worker death is
+        recorded (from the error monitor, a health check, or a crashed future),
+        all pending GenerationResults fail fast with EngineDeadError rather than
+        blocking forever on a response queue whose producer is gone.
+        """
+        super()._set_fatal_error(error)
+        self._mark_engine_dead(error)
+
+    def _mark_engine_dead(self, error: Optional[BaseException] = None) -> None:
+        """Set the sticky flag and push EngineDeadError onto pending results."""
+        if self._engine_dead:
+            return
+        self._engine_dead = True
+        dead_error = EngineDeadError(error)
+        # Snapshot to avoid mutation-during-iteration; best-effort per result.
+        # Use put() (not put_nowait()) so the async _SyncQueue path also wakes
+        # the awaiting event loop, not just the sync Queue path.
+        for result in list(self._results.values()):
+            try:
+                result.queue.put(dead_error)
+            except Exception:  # noqa: BLE001 - a full/closed queue must not stop the sweep
+                pass
+
+    def _handle_worker_death(self, error: BaseException) -> None:
+        """Event-driven worker-death handler.
+
+        Runs from the MPI future's done-callback thread the instant a worker
+        process exits. Enqueues the error for the monitor loop to record (which
+        drives pre_shutdown) and immediately broadcasts EngineDeadError to
+        pending requests, so they fail fast without waiting for the next poll
+        tick. Propagation is therefore no longer gated by the poll interval.
+        """
+        self._error_queue.put_nowait(error)
+        self._mark_engine_dead(error)
+
+    def _check_remote_worker_death(self) -> bool:
+        """Poll a remote (pre-spawned) MPI session for forwarded worker deaths.
+
+        Under ``TLLM_SPAWN_PROXY_PROCESS=1`` / ``trtllm-llmapi-launch`` the
+        session is a ``RemoteMpiCommSessionClient`` whose ``submit()`` returns
+        no futures, so the ``mpi_done_callback`` path never fires and
+        ``_check_mpi_futures()`` has nothing to watch. The remote server
+        forwards worker exceptions over its control socket instead
+        (``RemoteWorkerDeath``); surface them here into the same fast-death
+        path so this deployment mode gets identical EngineDeadError behavior.
+
+        Returns:
+            True if a dead worker was detected.
+        """
+        check = getattr(self.mpi_session, "check_worker_error", None)
+        if check is None:
+            return False
+        try:
+            error = check()
+        except Exception as exc:  # noqa: BLE001 - monitor must not die
+            logger.debug(f"check_worker_error failed (ignored): {exc!r}")
+            return False
+        if error is None:
+            return False
+        logger.error(f"Remote MPI worker death reported: {error!r}")
+        self._handle_worker_death(error)
+        self._set_fatal_error(error)
+        if not self.doing_shutdown:
+            self.pre_shutdown()
+        return True
+
     def _error_monitor_loop(self) -> None:
-        """Background thread that polls for fatal errors every ~5 seconds.
+        """Background thread that reaps a dead engine and drives pre_shutdown.
 
-        Checks MPI worker futures and drains the error queue using
-        the shared ``_check_mpi_futures()`` and ``_drain_error_queue()``
-        helpers.
+        Checks MPI worker futures, remote-session worker-death notifications,
+        and the error queue using the shared ``_check_mpi_futures()``,
+        ``_check_remote_worker_death()`` and ``_drain_error_queue()`` helpers.
 
-        Uses ``_shutdown_event`` for clean wakeup instead of a sleep loop,
-        so shutdown is immediate rather than waiting up to 5 seconds.
+        Propagation to pending requests is event-driven via
+        ``_handle_worker_death`` (the MPI future done-callback) where futures
+        exist; for remote sessions it is bounded by this loop's poll interval.
+        Uses ``_shutdown_event`` for clean wakeup instead of a sleep loop, so
+        shutdown is immediate.
         """
         while not self.doing_shutdown and self._fatal_error is None:
             try:
@@ -259,14 +336,22 @@ def _error_monitor_loop(self) -> None:
                                  "shutting down")
                     return
 
+                if self._check_remote_worker_death():
+                    logger.error(
+                        "Error monitor: remote MPI worker death detected, "
+                        "shutting down")
+                    return
+
                 self._drain_error_queue()
                 if self._fatal_error is not None:
                     return
             except Exception as exc:
                 logger.debug(f"Error monitor: unexpected exception (ignored): "
                              f"{exc!r}")
 
-            # Wait up to 5s, but wake immediately if _shutdown_event is set
+            # Backstop poll only; the latency-sensitive propagation path is
+            # event-driven (see _handle_worker_death), so a coarse interval is
+            # fine. Wakes immediately if _shutdown_event is set.
             self._shutdown_event.wait(timeout=5.0)
 
     def _setup_queues(self) -> WorkerCommIpcAddrs:
@@ -396,7 +481,7 @@ def mpi_done_callback(future: concurrent.futures.Future):
             # will not block.
             if future.exception() is not None:
                 if self_ := self_ref():
-                    self_._error_queue.put_nowait(future.exception())
+                    self_._handle_worker_death(future.exception())
 
         tracer_init_kwargs = get_tracer().init_kwargs if enable_llm_tracer(
         ) else None
@@ -542,6 +627,10 @@ def submit(self, request: GenerationRequest) -> GenerationResult:
             Forwards the request to the workers through the request queue.
         """
 
+        # Sticky fast-fail: don't accept new work once the engine is known dead.
+        if self._engine_dead or self._fatal_error is not None:
+            raise EngineDeadError(self._fatal_error)
+
         self._start_dispatch_threads()
 
         request.set_id(self._get_next_client_id())
@@ -555,6 +644,15 @@ def submit(self, request: GenerationRequest) -> GenerationResult:
             logprob_params=logprob_params)
         self._results[request.id] = result
 
+        # Close the submit-vs-death race: _mark_engine_dead() runs on the
+        # error-monitor thread and its one-shot sweep snapshots _results, so a
+        # result registered just after that snapshot would never be unblocked
+        # and would hang forever on a dead worker. Re-check after registering
+        # and fail fast if the engine died in that window.
+        if self._engine_dead or self._fatal_error is not None:
+            self._results.pop(request.id, None)
+            raise EngineDeadError(self._fatal_error)
+
         with nvtx_range_debug("request_queue.put"):
             self.request_queue.put(request)
 

tensorrt_llm/executor/result.py

@@ -31,7 +31,8 @@
     process_req_perf_metrics as _process_req_perf_metrics
 from ..sampling_params import LogprobParams, SamplingParams
 from .postprocessor_hook import PostProcessorHook, apply_post_processor_hook
-from .utils import ErrorResponse, has_event_loop, is_llm_response
+from .utils import (EngineDeadError, ErrorResponse, has_event_loop,
+                    is_llm_response)
 
 if TYPE_CHECKING:
     from .executor import GenerationExecutor
@@ -191,6 +192,10 @@ def __init__(self,
         # None indicates not yet available (e.g., before first step/stream).
         self.avg_decoded_tokens_per_iter: Optional[float] = None
         self._done = False
+        # Sticky terminal exception (e.g. EngineDeadError). Once set, the result
+        # is permanently failed: result()/aresult()/_exception() re-raise it on
+        # every subsequent call instead of looking successful.
+        self._terminal_error: Optional[BaseException] = None
         self._aborted = False
         self.metrics_dict = {}
         self.candidate_metrics: list[dict] = []
@@ -991,12 +996,26 @@ def _handle_ray_response(self, response: Any):
 
     def _result_step(self, timeout: Optional[float] = None):
         response = self.queue.get()
+        # Fast-fail: when a worker dies, the proxy enqueues EngineDeadError onto
+        # every pending result so this get() unblocks instead of hanging forever
+        # on a queue whose producer is gone. Record it as the sticky terminal
+        # error and mark the result done before raising, so subsequent
+        # result()/aresult()/_exception() calls keep surfacing the failure
+        # instead of re-blocking on an empty queue or looking successful.
+        if isinstance(response, EngineDeadError):
+            self._terminal_error = response
+            self._done = True
+            raise response
         self._handle_response(response)
 
     async def _aresult_step(self):
         assert self.aqueue is not None, "The asyncio event loop was not present during initialization, so async operations are not available."
         response = await self.aqueue.get()
         global_tracer().log_instant("result_step.get")
+        if isinstance(response, EngineDeadError):
+            self._terminal_error = response
+            self._done = True
+            raise response
         self._handle_response(response)
 
     def result(self, timeout: Optional[float] = None) -> "GenerationResult":
@@ -1008,6 +1027,8 @@ def result(self, timeout: Optional[float] = None) -> "GenerationResult":
         Returns:
             tensorrt_llm.executor.result.GenerationResult: generation result.
         """
+        if self._terminal_error is not None:
+            raise self._terminal_error
         while not self._done:
             self._result_step(timeout)
         return self
@@ -1018,6 +1039,8 @@ async def aresult(self) -> "GenerationResult":
         Returns:
             tensorrt_llm.executor.result.GenerationResult: generation result.
         """
+        if self._terminal_error is not None:
+            raise self._terminal_error
         while not self._done:
             await self._aresult_step()
         return self
@@ -1029,6 +1052,8 @@ def __iter__(self):
         return self
 
     def __next__(self):
+        if self._terminal_error is not None:
+            raise self._terminal_error
         if self._done:
             raise StopIteration
 
@@ -1039,6 +1064,8 @@ def __aiter__(self):
         return self
 
     async def __anext__(self):
+        if self._terminal_error is not None:
+            raise self._terminal_error
         if self._done:
             raise StopAsyncIteration
 

tensorrt_llm/executor/utils.py

@@ -80,6 +80,23 @@ class RequestError(RuntimeError):
     ''' The error raised when the request is failed. '''
 
 
+class EngineDeadError(RuntimeError):
+    """Raised by pending and new requests once the engine is known dead.
+
+    Sticky and engine-level (unlike the per-request ``RequestError``): when a
+    worker process dies, every queued ``GenerationResult`` is unblocked with this
+    error and every subsequent ``submit()`` raises it immediately, instead of
+    blocking forever on a response queue whose producer is gone.
+    """
+
+    def __init__(self, root_cause: Optional[BaseException] = None):
+        msg = "Engine has died"
+        if root_cause is not None:
+            msg += f": {type(root_cause).__name__}: {root_cause}"
+        super().__init__(msg)
+        self.root_cause = root_cause
+
+
 class ProcessPoolExecutorSession(MpiSession):
     # This process pool is introduced for better recoverable exceptions handling.
     # It replaces MpiPoolExecutor for single-gpu case.