diff --git a/tensorrt_llm/executor/__init__.py b/tensorrt_llm/executor/__init__.py index e6eb88f690b3..2bda8a05d5b6 100644 --- a/tensorrt_llm/executor/__init__.py +++ b/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", diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index ffe793bc5b48..dd795ca40027 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/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,6 +336,12 @@ 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 @@ -266,7 +349,9 @@ def _error_monitor_loop(self) -> None: 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) diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 4a16ee76587f..a3014da7072f 100644 --- a/tensorrt_llm/executor/result.py +++ b/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 diff --git a/tensorrt_llm/executor/utils.py b/tensorrt_llm/executor/utils.py index 12eaa9afefa9..a363350a1392 100644 --- a/tensorrt_llm/executor/utils.py +++ b/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. diff --git a/tensorrt_llm/llmapi/mpi_session.py b/tensorrt_llm/llmapi/mpi_session.py index 2e1eeb25a4e5..f178a34ebc16 100644 --- a/tensorrt_llm/llmapi/mpi_session.py +++ b/tensorrt_llm/llmapi/mpi_session.py @@ -302,6 +302,26 @@ class RemoteTask(NamedTuple): sync: bool = False # if True, the result will be sent back to the client +class RemoteWorkerDeath(NamedTuple): + """Worker-death notification forwarded by RemoteMpiCommSessionServer. + + Async (fire-and-forget) task submissions have no result channel, so a + crashed worker would otherwise be invisible to the client and pending + requests would block forever. The exception is carried as strings (not + the exception object) because arbitrary exceptions may not pickle. + """ + exc_type: str + message: str + + @classmethod + def from_exception(cls, e: BaseException) -> "RemoteWorkerDeath": + return cls(exc_type=type(e).__name__, message=str(e)) + + def to_exception(self) -> RuntimeError: + return RuntimeError( + f"Remote MPI worker died: {self.exc_type}: {self.message}") + + class RemoteMpiCommSessionClient(MpiSession): ''' RemoteMpiCommSessionClient is a variant of MpiCommSession that is used to connect to a remote MPI pool. @@ -344,6 +364,9 @@ def __init__(self, addr: str, hmac_key: bytes): socket_type=zmq.PAIR, use_hmac_encryption=True) self._is_shutdown = False + # Non-error messages consumed by check_worker_error() while scanning + # for RemoteWorkerDeath are buffered here for poll() (submit_sync). + self._pending_responses: list = [] self._initialized = True def submit(self, @@ -388,11 +411,33 @@ def poll(self) -> bool: ''' if self._is_shutdown: return False + if self._pending_responses: + return self._pending_responses.pop(0) response = self.queue.poll(0.1) if response: return self.queue.get() # should get a True if success return False + def check_worker_error(self) -> Optional[BaseException]: + """Non-blockingly fetch a worker-death notification, if any. + + RemoteMpiCommSessionServer forwards a RemoteWorkerDeath when an async + (fire-and-forget) worker future fails -- the only error channel in + this mode, since submit() returns no futures for the client to watch. + Non-error messages encountered while scanning are buffered for poll(). + """ + if self._is_shutdown: + return None + try: + while self.queue.poll(0): + msg = self.queue.get() + if isinstance(msg, RemoteWorkerDeath): + return msg.to_exception() + self._pending_responses.append(msg) + except Exception as e: + logger_debug(f"check_worker_error poll failed: {e}\n", "grey") + return None + def abort(self): self.shutdown() @@ -520,9 +565,36 @@ def serve(self): assert len(futures) == self.num_results == mpi_world_size() # Store futures to wait for them before the next task pending_futures = list(futures) - if message.sync: - for future in futures: + for future in futures: + if message.sync: future.add_done_callback(self.mpi_future_callback) + else: + # Fire-and-forget tasks have no result channel, but a + # crashed worker must still reach the client (the + # client-side session has no futures to watch); see + # RemoteWorkerDeath. + future.add_done_callback(self.mpi_async_error_callback) + + def mpi_async_error_callback(self, future): + """Forward a worker exception to the client for async tasks. + + Runs on the executor's callback thread, like the existing sync-path + mpi_future_callback (same pre-existing cross-thread ZMQ-put pattern). + Best-effort: the socket may already be closed at shutdown. + """ + if future.cancelled(): + return + exc = future.exception() + if exc is None: + return + print_colored( + f"RemoteMpiCommSessionServer: async MPI worker failed, forwarding " + f"to client: {type(exc).__name__}: {exc}\n", "red") + try: + self.queue.put(RemoteWorkerDeath.from_exception(exc)) + except Exception as e: + logger_debug(f"Failed to forward worker death to client: {e}\n", + "red") def mpi_future_callback(self, future): logger_debug(f"rank{global_mpi_rank()} got future: {future}\n", "red") diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 6c603e43a381..4c6c6082ef8c 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -137,6 +137,7 @@ l0_a10: - unittest/executor/test_fatal_error_health_check.py - unittest/executor/test_postprocessor_hook.py - unittest/executor/test_proxy_postproc_terminate.py + - unittest/executor/test_proxy_fast_death.py # trtllm-serve CPU-only - unittest/llmapi/apps/test_chat_utils.py - unittest/llmapi/apps/test_tool_parsers.py diff --git a/tests/unittest/executor/test_proxy_fast_death.py b/tests/unittest/executor/test_proxy_fast_death.py new file mode 100644 index 000000000000..ed0e62ae063a --- /dev/null +++ b/tests/unittest/executor/test_proxy_fast_death.py @@ -0,0 +1,324 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""EngineDeadError, the pending-result broadcast, and queue-step fast-fail (no GPU).""" + +import asyncio +import queue as _queue + +import pytest + +from tensorrt_llm.executor import EngineDeadError +from tensorrt_llm.executor.proxy import GenerationExecutorProxy +from tensorrt_llm.executor.result import GenerationResult + + +def test_engine_dead_error_is_importable_and_carries_root_cause(): + cause = RuntimeError("MPI worker exited unexpectedly") + err = EngineDeadError(cause) + assert isinstance(err, RuntimeError) + assert err.root_cause is cause + assert "Engine has died" in str(err) + assert "MPI worker exited unexpectedly" in str(err) + # No root cause is also fine. + assert "Engine has died" in str(EngineDeadError()) + + +class _FakeResult: + """Minimal stand-in for GenerationResult exposing only a `.queue`.""" + + def __init__(self): + self.queue = _queue.Queue() + + +def _bare_proxy(): + proxy = GenerationExecutorProxy.__new__(GenerationExecutorProxy) + proxy._engine_dead = False + proxy._results = {} + # Set so the __del__ -> shutdown() path is a clean no-op at GC time. + proxy.workers_started = False + return proxy + + +def test_mark_engine_dead_broadcasts_to_pending_results(): + proxy = _bare_proxy() + r1, r2 = _FakeResult(), _FakeResult() + proxy._results = {1: r1, 2: r2} + + proxy._mark_engine_dead(RuntimeError("worker died")) + + assert proxy._engine_dead is True + for r in (r1, r2): + item = r.queue.get_nowait() + assert isinstance(item, EngineDeadError) + assert isinstance(item.root_cause, RuntimeError) + + +def test_mark_engine_dead_is_idempotent(): + proxy = _bare_proxy() + r = _FakeResult() + proxy._results = {1: r} + + proxy._mark_engine_dead(RuntimeError("first")) + assert isinstance(r.queue.get_nowait(), EngineDeadError) + + # A second call must not enqueue another error. + proxy._mark_engine_dead(RuntimeError("second")) + assert r.queue.empty() + + +def test_handle_worker_death_broadcasts_event_driven(): + """Worker death propagates immediately via the future done-callback path.""" + proxy = _bare_proxy() + proxy._error_queue = _queue.Queue() + r = _FakeResult() + proxy._results = {1: r} + + cause = RuntimeError("worker segfault") + proxy._handle_worker_death(cause) + + # Broadcast happened event-driven: sticky flag set + EngineDeadError pushed. + assert proxy._engine_dead is True + item = r.queue.get_nowait() + assert isinstance(item, EngineDeadError) + assert item.root_cause is cause + # Error is also recorded for the monitor loop (which drives pre_shutdown). + assert proxy._error_queue.get_nowait() is cause + + +def test_result_step_raises_on_engine_dead(): + res = GenerationResult.__new__(GenerationResult) + res.queue = _queue.Queue() + res.queue.put(EngineDeadError(RuntimeError("boom"))) + with pytest.raises(EngineDeadError): + res._result_step() + + +def test_aresult_step_raises_on_engine_dead(): + """Async path must also unblock and raise via the _SyncQueue.put notify.""" + from tensorrt_llm.llmapi.utils import AsyncQueue + + async def run(): + res = GenerationResult.__new__(GenerationResult) + res.aqueue = AsyncQueue() + res.queue = res.aqueue.sync_q + res.queue.put(EngineDeadError(RuntimeError("boom"))) + with pytest.raises(EngineDeadError): + await res._aresult_step() + + asyncio.run(run()) + + +def test_result_stays_failed_after_engine_dead(): + """A dead engine is a sticky terminal failure surfaced on every access.""" + res = GenerationResult.__new__(GenerationResult) + res.queue = _queue.Queue() + res._done = False + res._terminal_error = None + res.queue.put(EngineDeadError(RuntimeError("boom"))) + + # First access consumes the queued error and records it as terminal. + with pytest.raises(EngineDeadError): + res._result_step() + assert res._done is True + assert res.queue.empty() + + # Repeated access must keep raising (not re-block, not return a + # successful-looking self). + with pytest.raises(EngineDeadError): + res.result() + with pytest.raises(EngineDeadError): + res.result() + # _exception() must keep surfacing the same failure, not None. + assert isinstance(res._exception(), EngineDeadError) + + +def test_aresult_stays_failed_after_engine_dead(): + from tensorrt_llm.llmapi.utils import AsyncQueue + + async def run(): + res = GenerationResult.__new__(GenerationResult) + res.aqueue = AsyncQueue() + res.queue = res.aqueue.sync_q + res._done = False + res._terminal_error = None + res.queue.put(EngineDeadError(RuntimeError("boom"))) + + with pytest.raises(EngineDeadError): + await res._aresult_step() + assert res._done is True + # Repeated await must keep raising, not return successful-looking self. + with pytest.raises(EngineDeadError): + await res.aresult() + + asyncio.run(run()) + + +def test_submit_fast_fails_when_engine_already_dead(): + """submit() must reject new work immediately once the engine is dead.""" + proxy = _bare_proxy() + proxy._fatal_error = None + proxy._engine_dead = True + # The sticky guard is the first thing submit() does, before it touches the + # request, so a placeholder request is never dereferenced. + with pytest.raises(EngineDeadError): + proxy.submit(object()) + + +class _FakeRequest: + """Minimal GenerationRequest stand-in for the submit() re-check test.""" + + disaggregated_params = None + + def set_id(self, request_id): + self.id = request_id + + +def test_submit_rechecks_engine_death_after_registration(monkeypatch): + """Close the submit-vs-death race. + + If the engine dies between the top-of-submit guard and registering the + result in _results (so _mark_engine_dead's one-shot sweep misses it), + submit() must still fail fast and not leak the dangling result. + """ + proxy = _bare_proxy() + proxy._fatal_error = None + proxy._engine_dead = False + proxy._results = {} + proxy._start_dispatch_threads = lambda: None + proxy._get_next_client_id = lambda: 42 + proxy._get_logprob_params = lambda request: None + proxy._handle_background_error = lambda *a, **k: None + + fake_result = _FakeResult() + + def fake_generation_result(*args, **kwargs): + # Simulate the error-monitor thread marking the engine dead mid-submit, + # after the top guard passed but as the result is being created. + proxy._engine_dead = True + return fake_result + + monkeypatch.setattr("tensorrt_llm.executor.proxy.GenerationResult", fake_generation_result) + + with pytest.raises(EngineDeadError): + proxy.submit(_FakeRequest()) + + # The raced result must not be left dangling in _results. + assert 42 not in proxy._results + + +# --- Remote (TLLM_SPAWN_PROXY_PROCESS / trtllm-llmapi-launch) mode coverage --- +# RemoteMpiCommSessionClient.submit() returns no futures, so worker death must +# travel server -> client as a RemoteWorkerDeath over the control socket and be +# surfaced by the proxy's _check_remote_worker_death(). + + +def test_remote_worker_death_roundtrip(): + from tensorrt_llm.llmapi.mpi_session import RemoteWorkerDeath + + death = RemoteWorkerDeath.from_exception(ValueError("rank 3 exploded")) + exc = death.to_exception() + assert isinstance(exc, RuntimeError) + assert "ValueError" in str(exc) + assert "rank 3 exploded" in str(exc) + + +def test_server_async_callback_forwards_only_failures(): + from concurrent.futures import Future + + from tensorrt_llm.llmapi.mpi_session import RemoteMpiCommSessionServer, RemoteWorkerDeath + + server = object.__new__(RemoteMpiCommSessionServer) + sent = [] + server.queue = type("Q", (), {"put": lambda self, m: sent.append(m)})() + + ok = Future() + ok.set_result(42) + server.mpi_async_error_callback(ok) + assert sent == [] + + cancelled = Future() + cancelled.cancel() + server.mpi_async_error_callback(cancelled) + assert sent == [] + + failed = Future() + failed.set_exception(RuntimeError("worker segfault")) + server.mpi_async_error_callback(failed) + assert len(sent) == 1 and isinstance(sent[0], RemoteWorkerDeath) + assert sent[0].message == "worker segfault" + + +class _FakeZmqQueue: + """poll()/get() stub fed with a fixed message sequence.""" + + def __init__(self, messages): + self._messages = list(messages) + + def poll(self, timeout): + return bool(self._messages) + + def get(self): + return self._messages.pop(0) + + +def _bare_remote_client(messages): + from tensorrt_llm.llmapi.mpi_session import RemoteMpiCommSessionClient + + client = object.__new__(RemoteMpiCommSessionClient) # bypass singleton + client._is_shutdown = False + client._pending_responses = [] + client.queue = _FakeZmqQueue(messages) + return client + + +def test_client_check_worker_error_returns_death_and_buffers_others(): + from tensorrt_llm.llmapi.mpi_session import RemoteWorkerDeath + + death = RemoteWorkerDeath.from_exception(RuntimeError("boom")) + client = _bare_remote_client([[1, 2, 3], death]) + + exc = client.check_worker_error() + assert isinstance(exc, RuntimeError) and "boom" in str(exc) + # The non-error message was buffered for poll() (submit_sync path). + assert client.poll() == [1, 2, 3] + # Nothing left. + assert client.check_worker_error() is None + + +def test_proxy_check_remote_worker_death_marks_engine_dead(): + proxy = _bare_proxy() + proxy._fatal_error = None + proxy._error_queue = _queue.Queue() + proxy.doing_shutdown = False + pre_shutdowns = [] + proxy.pre_shutdown = lambda: pre_shutdowns.append(1) + r = _FakeResult() + proxy._results = {1: r} + + death_exc = RuntimeError("Remote MPI worker died: X: boom") + proxy.mpi_session = type("S", (), {"check_worker_error": lambda self: death_exc})() + + assert proxy._check_remote_worker_death() is True + # Same fast-death behavior as the local-futures path: + assert proxy._engine_dead is True + assert isinstance(r.queue.get_nowait(), EngineDeadError) + assert proxy._fatal_error is death_exc + assert proxy._error_queue.get_nowait() is death_exc + assert pre_shutdowns == [1] + + # Sessions without the hook (MpiPoolSession) are a no-op. + proxy2 = _bare_proxy() + proxy2.mpi_session = object() + assert proxy2._check_remote_worker_death() is False