From e537468a18719660d91eae8819d74af508d82efd Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 28 Jul 2026 23:17:18 +0000 Subject: [PATCH 1/2] [None][fix] Make a stalled executor worker init self-report The proxy's worker-startup handshake exits only on a leader status message or on worker death, so a rank that is alive but wedged during initialization (e.g. inside an NCCL bootstrap collective) leaves the proxy spinning in worker_init_status_queue.poll(1) with nothing said. The failure surfaces only as an outer harness timeout, with a client-side stack that says nothing beyond "waiting on a queue" -- neither which rank is stuck nor whether any rank is stuck at all. This does not add a deadline. Initialization that is slow but healthy (a very large checkpoint loading from a cold mount) must not be killed, so the loop still exits only on ready, init error, or a dead rank. What changes is that the wait is no longer silent: * every worker rank arms a watchdog thread before the first collective and, while its own initialization is still pending, logs its rank/pid and dumps all of its thread stacks (reusing print_all_stacks) every TRTLLM_WORKER_INIT_STALL_WARN_SEC seconds (default 600); * the proxy logs a matching stall report naming the elapsed time and, where it can see them, how many worker tasks are still running, which distinguishes a stalled initialization from the already-covered crash. During init the proxy has an IPC channel to the rank-0 leader only, so a wedged non-leader is invisible to it: the only process that can say where rank N is stuck is rank N. Hence the per-rank self-report. The leader stays armed until it has delivered the ready signal, not merely until its constructor returns -- that is when the proxy's wait actually ends, and a leader wedged in between is precisely the case the proxy cannot see. Subordinates disarm after construction, since they then block in block_subordinates() for the life of the job. Attribution is the whole deliverable here, so the report must not claim more than it knows. RemoteMpiCommSessionClient.submit() returns [], so under trtllm-llmapi-launch mpi_futures is empty and counting it would print "0/0 worker task(s) are still running, so no rank has exited" -- a confident liveness statement made with zero visibility, the same empty-list trap pre_shutdown() already documents. With no futures the report now says plainly that it cannot tell whether a rank has exited and names check_worker_error() as the authoritative channel; it does not call it, because reading it consumes the death notice _check_remote_worker_death() acts on. For the same reason the knob is now validated instead of being coerced. float("nan") does not raise, and every comparison with nan is False, so a typo'd TRTLLM_WORKER_INIT_STALL_WARN_SEC=nan slipped past the "period <= 0" disable check and armed a watchdog whose Event.wait(nan) returns immediately (measured: ~10us, no exception) rather than sleeping -- a hot loop dumping every thread's stack at WARNING on every rank. Event.wait(inf) raises OverflowError inside the watchdog thread instead, silently killing the reporting. Non-finite and unparsable values now raise, and the proxy reads the knob before any rank is spawned: startup is exactly where information gets lost, so swallowing a misconfiguration here is the wrong default. Both are emitted at WARNING, not ERROR: a slow load is not a fault, and this fires on that run too. print_all_stacks() grows an optional log callable so the dump follows the level of the condition that triggered it; its default stays logger.error for existing callers. The knob is an environment variable rather than an LlmArgs field because the TRTLLM-prefixed environment is already forwarded to spawned MPI ranks, so a single export arms both sides, and the protected LLM API surface stays untouched. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/_utils.py | 16 +- tensorrt_llm/executor/proxy.py | 70 +- tensorrt_llm/executor/utils.py | 60 ++ tensorrt_llm/executor/worker.py | 63 +- .../executor/test_proxy_worker_startup.py | 621 ++++++++++++++++++ 5 files changed, 822 insertions(+), 8 deletions(-) create mode 100644 tests/unittest/executor/test_proxy_worker_startup.py diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 71574f3e1aa5..d005d149e893 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -700,11 +700,19 @@ def is_flashinfer_gdn_supported_arch(sm_version=None): return sm_version in (90, 100, 103) -def print_all_stacks(): - """Print stack traces for all threads""" +def print_all_stacks(log=None): + """Print stack traces for all threads + + Args: + log: logging callable used to emit the traces; defaults to + ``logger.error``. Callers dumping stacks for a condition that is + not (yet) a fault -- e.g. a slow but healthy startup -- should pass + ``logger.warning`` instead. + """ + log = logger.error if log is None else log for thread_id, frame in sys._current_frames().items(): - logger.error(f"Thread {thread_id} stack trace:\n" + - "".join(traceback.format_stack(frame))) + log(f"Thread {thread_id} stack trace:\n" + + "".join(traceback.format_stack(frame))) def is_trace_enabled(env_var: str): diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index 32afb4780ed1..3dacba53bb7d 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -19,6 +19,7 @@ import shutil import tempfile import threading +import time import weakref from queue import Empty from typing import Dict, List, Optional, Union @@ -43,11 +44,12 @@ from .result import GenerationResult, IterationResult from .rpc import RPCClient from .rpc.rpc_common import RPCError, get_unique_ipc_addr -from .utils import (EngineDeadError, ErrorResponse, RequestError, - WorkerCommIpcAddrs, create_mpi_comm_session, +from .utils import (WORKER_INIT_STALL_WARN_ENV, EngineDeadError, ErrorResponse, + RequestError, WorkerCommIpcAddrs, create_mpi_comm_session, get_spawn_proxy_process_env, is_llm_response, multi_frontend_request_addr, multi_frontend_result_addr, - namespace_client_id, print_alive_threads) + namespace_client_id, print_alive_threads, + worker_init_stall_warn_sec) from .worker import GenerationExecutorWorker, worker_main from .worker_process_monitor import WorkerProcessIdentity, WorkerProcessMonitor @@ -605,6 +607,10 @@ def _start_dispatch_threads(self): def _start_executor_workers(self, worker_kwargs): + # Read the knob before anything is spawned: a bad value raises, and it + # must do so while there are still no ranks to leave behind. + stall_warn_sec = worker_init_stall_warn_sec() + self_ref = weakref.ref(self) def mpi_done_callback(future: concurrent.futures.Future): @@ -642,6 +648,17 @@ def mpi_done_callback(future: concurrent.futures.Future): self.workers_started = True + # This loop still exits only on the leader's status (ready or init + # error) or on a rank dying. A third case exists -- every rank alive + # but nothing happening -- which neither exit covers; it is not given + # a deadline here (a slow-but-healthy startup must not be killed), but + # it is no longer silent: the stall report below, together with the + # per-rank stack dumps every worker emits on the same schedule, says + # which rank is stuck and where. + start_time = time.monotonic() + next_warn_time = (start_time + + stall_warn_sec) if stall_warn_sec > 0 else None + while True: if self.worker_init_status_queue.poll(1): status = self.worker_init_status_queue.get() @@ -654,6 +671,11 @@ def mpi_done_callback(future: concurrent.futures.Future): raise RuntimeError("Executor worker died during initialization") self._handle_background_error() + now = time.monotonic() + if next_warn_time is not None and now >= next_warn_time: + next_warn_time = now + stall_warn_sec + logger.warning(self._worker_init_stall_report(now - start_time)) + ready_signal, error_trace = status[:2] if ready_signal != GenerationExecutorProxy.READY_SIGNAL: logger.error(f"Executor worker initialization error: {error_trace}") @@ -666,6 +688,48 @@ def mpi_done_callback(future: concurrent.futures.Future): self._register_worker_processes(status) + def _worker_init_stall_report(self, elapsed: float) -> str: + """Describe a startup that has gone quiet, and where to attribute it. + + The proxy has no IPC channel to ranks other than the leader before the + ready signal arrives, so it cannot collect worker stacks itself. What + it can state is (a) whether any rank has exited -- which distinguishes + a stalled initialization from a crash, the two being indistinguishable + from the client side today -- and (b) where the per-rank stacks that + each worker dumps on the same schedule can be found. + + (a) is only answerable when the session actually hands back worker + futures. ``RemoteMpiCommSessionClient.submit()`` returns ``[]`` (the + ranks live in a separate ``mgmn_leader_node`` process under + ``trtllm-llmapi-launch``), so on that session type the proxy has *no* + liveness signal here and must say so rather than read an empty list as + "everyone is fine" -- the same empty-list trap ``pre_shutdown()`` + documents below. ``check_worker_error()`` on the session is the + authoritative channel there; it is not consulted from this report + because reading it consumes the death notice that + ``_check_remote_worker_death()`` acts on. + """ + total = len(self.mpi_futures) + if total: + running = sum(1 for fut in self.mpi_futures if not fut.done()) + liveness = ( + f"{running}/{total} worker task(s) are still running, so no " + f"rank has exited (a stalled initialization, not a worker " + f"crash)") + else: + liveness = ( + "whether a rank has exited cannot be told from here: this MPI " + "session hands back no worker futures, so the proxy has no " + "liveness signal during startup and this report can neither " + "confirm nor rule out a crashed rank. The session's " + "check_worker_error() channel is authoritative for that") + return (f"Executor worker initialization has not completed after " + f"{elapsed:.0f}s; {liveness}. Every rank dumps its own thread " + f"stacks on the same schedule: search the worker logs for 'has " + f"not finished initialization' to see which rank is stuck and " + f"where. Tune or silence this report with " + f"{WORKER_INIT_STALL_WARN_ENV}.") + def _register_worker_processes(self, status: tuple) -> None: """Register identities returned by locally spawned MPI workers. diff --git a/tensorrt_llm/executor/utils.py b/tensorrt_llm/executor/utils.py index e4a9f3333f4b..3eeabede1027 100644 --- a/tensorrt_llm/executor/utils.py +++ b/tensorrt_llm/executor/utils.py @@ -16,6 +16,7 @@ import asyncio import concurrent.futures import ctypes +import math import os import re import sys @@ -300,6 +301,65 @@ def is_llm_response(instance): return hasattr(instance, "has_error") +# --- Executor worker startup (initialization) stall reporting --------------- +# +# Both the proxy and every worker rank read this. It is an environment +# variable rather than an LlmArgs field on purpose: the ``TRTLLM``-prefixed +# environment is forwarded verbatim to spawned MPI ranks +# (``MpiPoolSession._start_mpi_pool``), so a single ``export`` sets the +# proxy-side report and the per-rank stack dumps together, and CI can set it +# without touching the protected LLM API surface. It sits next to the existing +# ``TRTLLM_WORKER_PRINT_STACKS_PERIOD`` / ``TRTLLM_WORKER_DISABLE_GC`` worker +# knobs. + +#: Seconds of silence before the proxy logs a startup-stall report and every +#: worker rank dumps its own thread stacks. Diagnostics only -- nothing is +#: killed and no deadline is imposed -- so this is safe to enable by default +#: even for a startup that is slow but healthy (e.g. a 671B checkpoint loading +#: from a cold NFS mount). Set <= 0 to silence; anything that is not a finite +#: number is a hard error rather than a silent fallback (see +#: :func:`float_from_env`). +WORKER_INIT_STALL_WARN_ENV = "TRTLLM_WORKER_INIT_STALL_WARN_SEC" +WORKER_INIT_STALL_WARN_DEFAULT_SEC = 600.0 + + +def float_from_env(name: str, default: float) -> float: + """Read a float-valued environment variable, rejecting anything unusable. + + Unset or blank takes ``default``. Anything else that is not a finite + number raises ``ValueError`` rather than falling back, because this knob + configures the reporting that a wedged startup depends on: silently + substituting the default would hide the misconfiguration behind exactly + the silence this reporting exists to remove, and the variable is only ever + set deliberately, so there is no legitimate value to preserve. + + ``nan`` in particular must not survive. Every comparison against it is + False, so it slips past the ``period <= 0`` disable check, and + ``threading.Event.wait(nan)`` then returns False immediately (measured: + ~10us, no exception) instead of sleeping -- turning the watchdog into a + hot loop that dumps every thread's stack on every rank. ``inf`` is no + better: ``Event.wait(inf)`` raises ``OverflowError`` from inside the + watchdog thread, killing the reporting silently. + """ + raw = os.getenv(name) + if raw is None or not raw.strip(): + return default + try: + value = float(raw) + except ValueError as e: + raise ValueError(f"Invalid {name}={raw!r}: expected a number of " + f"seconds (default {default}).") from e + if not math.isfinite(value): + raise ValueError(f"Invalid {name}={raw!r}: expected a *finite* number " + f"of seconds (default {default}).") + return value + + +def worker_init_stall_warn_sec() -> float: + return float_from_env(WORKER_INIT_STALL_WARN_ENV, + WORKER_INIT_STALL_WARN_DEFAULT_SEC) + + def print_alive_threads(): assert enable_llm_debug( ), "print_alive_threads must be called with enable_llm_debug() enabled" diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index f173aabc224a..53f9b44f73bc 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -26,7 +26,7 @@ from .request import CancellingRequest, GenerationRequest from .rpc_worker_mixin import RpcWorkerMixin from .utils import (ErrorResponse, IntraProcessQueue, RequestError, - WorkerCommIpcAddrs) + WorkerCommIpcAddrs, worker_init_stall_warn_sec) from .worker_process_monitor import capture_worker_process_identity __all__ = [ @@ -34,6 +34,49 @@ ] +def _worker_init_stall_watchdog(init_done: threading.Event, + period: float) -> None: + """Report where *this* rank is stuck while its initialization is pending. + + Runs in a daemon thread in every rank. During initialization the proxy + only has an IPC channel to the rank-0 leader, so a rank that is alive but + wedged -- typically inside a collective -- is invisible to it: the only + process that can say where rank N is stuck is rank N. Dumping stacks is + read-only and never interferes with a slow-but-healthy startup, which is + why this is armed by default. + + The dump relies on the wedged call having released the GIL (true for the + torch/NCCL collectives this targets). A rank stuck while holding the GIL + cannot report, and simply stays silent. + """ + rank = mpi_rank() + pid = os.getpid() + waited = 0.0 + while not init_done.wait(period): + waited += period + # WARNING, not ERROR: a very large checkpoint loading from a cold mount + # is slow, not faulty, and this fires on that run too. + logger.warning( + f"Executor worker rank {rank} (pid {pid}) has not finished " + f"initialization after {waited:.0f}s; dumping the stacks of all " + f"threads on this rank.") + print_all_stacks(log=logger.warning) + + +def _arm_worker_init_stall_watchdog( + init_done: threading.Event) -> Optional[threading.Thread]: + """Start the init stall watchdog unless it is disabled by configuration.""" + period = worker_init_stall_warn_sec() + if period <= 0: + return None + watchdog = threading.Thread(target=_worker_init_stall_watchdog, + args=(init_done, period), + name="worker_init_stall_watchdog", + daemon=True) + watchdog.start() + return watchdog + + class GenerationExecutorWorker(RpcWorkerMixin, BaseWorker): def __init__( @@ -194,6 +237,11 @@ def _print_stacks(): daemon=True) print_stacks_thread.start() + # Arm before the first collective: everything from here to the end of + # engine construction is init-phase work that can wedge a rank. + worker_init_done = threading.Event() + _arm_worker_init_stall_watchdog(worker_init_done) + mpi_comm().barrier() if llm_args is not None and llm_args.env_overrides: @@ -352,8 +400,18 @@ def notify_proxy_threads_to_quit(): error_msg = (e, traceback.format_exc()) if not worker_init_status_queue.notify_with_retry(error_msg): logger.error("Failed to deliver error message to proxy") + worker_init_done.set() return + if not is_leader: + # A subordinate's startup ends with construction: it blocks in + # block_subordinates() below and has nothing further to report. The + # leader stays armed until it has actually delivered the ready signal, + # because that -- not the end of construction -- is when the proxy's + # wait ends, and a leader wedged in between is exactly the case the + # proxy cannot see for itself. + worker_init_done.set() + # Optionally disable GC (default: not disabled) if os.getenv("TRTLLM_WORKER_DISABLE_GC", "0") == "1": gc.disable() @@ -376,6 +434,9 @@ def notify_proxy_threads_to_quit(): logger.warning( "Failed to deliver ready signal to proxy, continuing anyway" ) + # The proxy's startup wait ends here (or has been given up on); + # either way this rank has nothing more to contribute to it. + worker_init_done.set() if resource_governor_queue is not None: # Swap rank 0 to the proxy IPC queue after construction. # The resource-governor flag is already enabled on all diff --git a/tests/unittest/executor/test_proxy_worker_startup.py b/tests/unittest/executor/test_proxy_worker_startup.py new file mode 100644 index 000000000000..8cd851e0345a --- /dev/null +++ b/tests/unittest/executor/test_proxy_worker_startup.py @@ -0,0 +1,621 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. +"""Startup handshake between the executor proxy and its MPI workers. + +Every proxy test here drives the real +``GenerationExecutorProxy._start_executor_workers``; only the MPI session and +the init status queue are faked, so no GPU (and no MPI spawn) is needed. The +worker-side tests drive the real +``tensorrt_llm.executor.worker._worker_init_stall_watchdog`` / +``_arm_worker_init_stall_watchdog``. +""" + +import ast +import pathlib +import queue +import sys +import threading +import time +import types +from concurrent.futures import Future + +import pytest + +from tensorrt_llm._utils import print_all_stacks +from tensorrt_llm.executor import worker as worker_module +from tensorrt_llm.executor.proxy import GenerationExecutorProxy +from tensorrt_llm.executor.utils import WORKER_INIT_STALL_WARN_ENV, worker_init_stall_warn_sec + + +class RecordingLogger: + """Minimal stand-in for ``tensorrt_llm.logger`` that keeps the messages.""" + + def __init__(self): + self.warnings = [] + self.errors = [] + + def warning(self, message, *args, **kwargs): + self.warnings.append(str(message)) + + def error(self, message, *args, **kwargs): + self.errors.append(str(message)) + + def info(self, message, *args, **kwargs): + pass + + def debug(self, message, *args, **kwargs): + pass + + +class FakeMpiSession: + def __init__(self, futures): + self.futures = futures + self.shutdown_abort_reasons = [] + self.submitted_kwargs = None + + def submit(self, *args, **kwargs): + self.submitted_kwargs = kwargs + return self.futures + + def shutdown_abort(self, *, reason=None, grace=60): + del grace + self.shutdown_abort_reasons.append(reason) + + +class FakeWorkerInitStatusQueue: + def __init__(self, messages=None, on_poll=None): + self.messages = queue.Queue() + for message in messages or []: + self.messages.put(message) + self.on_poll = on_poll + self.acks = [] + self.poll_count = 0 + + def poll(self, timeout): + del timeout + self.poll_count += 1 + if self.on_poll is not None: + self.on_poll(self.poll_count) + return not self.messages.empty() + + def get(self): + return self.messages.get_nowait() + + def put(self, message): + self.acks.append(message) + + +def _make_proxy(monkeypatch, *, futures, init_messages=None, on_poll=None, owns_mpi_session=True): + proxy = object.__new__(GenerationExecutorProxy) + proxy._error_queue = queue.Queue() + proxy._fatal_error = None + proxy.doing_shutdown = False + proxy.worker_cls = object + proxy.workers_started = False + proxy._owns_mpi_session = owns_mpi_session + proxy.mpi_session = FakeMpiSession(futures) + proxy.worker_init_status_queue = FakeWorkerInitStatusQueue(init_messages, on_poll) + fake_modeling_auto = types.SimpleNamespace(MODEL_CLASS_MAPPING={}) + monkeypatch.setitem(sys.modules, "tensorrt_llm._torch.models.modeling_auto", fake_modeling_auto) + monkeypatch.setattr("tensorrt_llm.executor.proxy.torch.cuda.Stream", lambda: None) + monkeypatch.setattr("tensorrt_llm.executor.proxy.enable_llm_tracer", lambda: False) + # The startup knobs are environment driven; never inherit them from the + # environment the test suite happens to run in. + monkeypatch.delenv(WORKER_INIT_STALL_WARN_ENV, raising=False) + return proxy + + +def test_worker_ready_signal_exits_startup_loop(monkeypatch): + future = Future() + proxy = _make_proxy( + monkeypatch, + futures=[future], + init_messages=[(GenerationExecutorProxy.READY_SIGNAL, None)], + ) + + proxy._start_executor_workers({"tokenizer": object(), "keep": "value"}) + + assert proxy.workers_started is True + assert proxy.worker_init_status_queue.acks == ["ACK"] + assert proxy.mpi_session.shutdown_abort_reasons == [] + assert "tokenizer" not in proxy.mpi_session.submitted_kwargs + assert proxy.mpi_session.submitted_kwargs["keep"] == "value" + + +def test_worker_init_error_aborts_mpi_session(monkeypatch): + future = Future() + init_error = RuntimeError("rank 1 failed during init") + proxy = _make_proxy( + monkeypatch, + futures=[future], + init_messages=[(init_error, "rank 1 traceback")], + ) + + with pytest.raises(RuntimeError, match="Executor worker returned error"): + proxy._start_executor_workers({}) + + assert proxy.worker_init_status_queue.acks == ["ACK"] + assert proxy.mpi_session.shutdown_abort_reasons == [init_error] + + +def test_worker_future_done_before_ready_fails_fast(monkeypatch): + future = Future() + future.set_exception(RuntimeError("rank 1 exited")) + proxy = _make_proxy(monkeypatch, futures=[future]) + + with pytest.raises(RuntimeError, match="Executor worker died during initialization"): + proxy._start_executor_workers({}) + + +def test_alive_worker_without_ready_signal_keeps_waiting_by_default(monkeypatch): + """No bound is imposed unless one is asked for. + + Initialization that is slow but healthy (a very large checkpoint loading + from a cold mount) must not be killed, so the default behavior stays + "wait"; only the reporting is new. + """ + future = Future() + polled_repeatedly = threading.Event() + + def on_poll(poll_count): + if poll_count >= 5: + polled_repeatedly.set() + + proxy = _make_proxy(monkeypatch, futures=[future], on_poll=on_poll) + result = {} + + def start_workers(): + try: + proxy._start_executor_workers({}) + except BaseException as exc: + result["exception"] = exc + + startup_thread = threading.Thread(target=start_workers, daemon=True) + startup_thread.start() + + try: + assert polled_repeatedly.wait(timeout=2) + assert startup_thread.is_alive() + assert proxy.mpi_session.shutdown_abort_reasons == [] + finally: + if not future.done(): + future.set_exception(RuntimeError("rank 1 eventually exited")) + + startup_thread.join(timeout=2) + + assert not startup_thread.is_alive() + assert isinstance(result["exception"], RuntimeError) + assert str(result["exception"]) == ("Executor worker died during initialization") + + +def test_stalled_startup_is_reported_while_waiting(monkeypatch): + """A wedged-but-alive startup is reported repeatedly, not silently awaited. + + Nothing terminates the wait, so the test releases the loop the same way + reality would -- by letting a rank die -- once it has seen enough reports. + """ + records = RecordingLogger() + monkeypatch.setattr("tensorrt_llm.executor.proxy.logger", records) + future = Future() + proxy = _make_proxy(monkeypatch, futures=[future]) + monkeypatch.setenv(WORKER_INIT_STALL_WARN_ENV, "0.02") + + def start_workers(): + try: + proxy._start_executor_workers({}) + except BaseException: # noqa: BLE001 - the release path, not the subject + pass + + startup_thread = threading.Thread(target=start_workers, daemon=True) + startup_thread.start() + try: + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + if len([m for m in records.warnings if "has not completed" in m]) >= 2: + break + time.sleep(0.01) + finally: + future.set_exception(RuntimeError("rank 1 eventually exited")) + startup_thread.join(timeout=5) + assert not startup_thread.is_alive() + + stall_reports = [ + message + for message in records.warnings + if "Executor worker initialization has not completed" in message + ] + assert len(stall_reports) >= 2, records.warnings + # The report must attribute the stall: no rank died, and it must point at + # the per-rank stacks that carry the "which rank, and where" detail. + assert "1/1 worker task(s) are still running" in stall_reports[0] + assert "has not finished initialization" in stall_reports[0] + # A stall must never be dressed up as the (separately handled) crash. + assert "died during initialization" not in stall_reports[0] + + +def test_ready_signal_is_not_delayed_by_stall_reporting(monkeypatch): + """Reporting must never get between the proxy and a queued ready signal.""" + proxy = _make_proxy( + monkeypatch, + futures=[Future()], + init_messages=[(GenerationExecutorProxy.READY_SIGNAL, None)], + ) + monkeypatch.setenv(WORKER_INIT_STALL_WARN_ENV, "0.01") + + proxy._start_executor_workers({}) + + assert proxy.workers_started is True + assert proxy.mpi_session.shutdown_abort_reasons == [] + + +def test_stall_report_without_worker_futures_claims_no_liveness(monkeypatch): + """``RemoteMpiCommSessionClient.submit()`` returns ``[]``. + + An empty ``mpi_futures`` is absence of evidence, not evidence of absence: + counting an empty list yields "0/0 ... so no rank has exited", a confident + statement made with zero visibility. ``pre_shutdown()`` documents the same + empty-list trap. + """ + proxy = _make_proxy(monkeypatch, futures=[]) + proxy.mpi_futures = proxy.mpi_session.submit() + + report = proxy._worker_init_stall_report(900.0) + + assert "0/0" not in report, report + assert "no rank has exited" not in report, report + assert "worker task(s) are still running" not in report, report + # It must say the liveness question is unanswerable here, and name the + # channel that can answer it. + assert "cannot be told from here" in report, report + assert "check_worker_error()" in report, report + # The attribution the report exists for is unaffected. + assert "has not finished initialization" in report, report + + +def test_stall_report_with_worker_futures_still_states_liveness(monkeypatch): + """The visible case must keep saying what it can see.""" + alive, dead = Future(), Future() + dead.set_exception(RuntimeError("rank 1 exited")) + proxy = _make_proxy(monkeypatch, futures=[alive, dead]) + proxy.mpi_futures = proxy.mpi_session.submit() + + report = proxy._worker_init_stall_report(900.0) + + assert "1/2 worker task(s) are still running" in report, report + assert "cannot be told from here" not in report, report + + +def test_stalled_startup_without_worker_futures_is_reported(monkeypatch): + """The same, driven through the real startup loop rather than the helper.""" + records = RecordingLogger() + monkeypatch.setattr("tensorrt_llm.executor.proxy.logger", records) + + def on_poll(poll_count): + # Nothing can die here (there are no futures), so release the loop the + # only other way: let the ready signal finally arrive. + if any("has not completed" in m for m in records.warnings): + proxy.worker_init_status_queue.messages.put( + (GenerationExecutorProxy.READY_SIGNAL, None) + ) + + proxy = _make_proxy(monkeypatch, futures=[], on_poll=on_poll) + monkeypatch.setenv(WORKER_INIT_STALL_WARN_ENV, "0.02") + + proxy._start_executor_workers({}) + + stall_reports = [ + m for m in records.warnings if "Executor worker initialization has not completed" in m + ] + assert stall_reports, records.warnings + assert "0/0" not in stall_reports[0], stall_reports[0] + assert "no rank has exited" not in stall_reports[0], stall_reports[0] + assert "cannot be told from here" in stall_reports[0], stall_reports[0] + + +def test_stall_warn_knob_defaults_and_blank_values(monkeypatch): + monkeypatch.delenv(WORKER_INIT_STALL_WARN_ENV, raising=False) + assert worker_init_stall_warn_sec() > 0.0 # reporting on by default + + # Unset and blank are the only inputs that fall back to the default. + monkeypatch.setenv(WORKER_INIT_STALL_WARN_ENV, "") + assert worker_init_stall_warn_sec() > 0.0 + monkeypatch.setenv(WORKER_INIT_STALL_WARN_ENV, " ") + assert worker_init_stall_warn_sec() > 0.0 + + monkeypatch.setenv(WORKER_INIT_STALL_WARN_ENV, "30") + assert worker_init_stall_warn_sec() == 30.0 + + +@pytest.mark.parametrize("raw", ["nan", "NaN", "-nan", "inf", "-inf", "1e400"]) +def test_stall_warn_knob_rejects_non_finite_values(monkeypatch, raw): + """``float()`` accepts these; the watchdog cannot survive them. + + ``nan`` defeats the ``period <= 0`` disable check (every comparison with + it is False) and then makes ``Event.wait(nan)`` return immediately rather + than sleep, so the watchdog becomes a hot loop dumping every thread's + stack on every rank. ``Event.wait(inf)`` raises ``OverflowError`` inside + the watchdog thread instead, silently killing the reporting. + """ + monkeypatch.setenv(WORKER_INIT_STALL_WARN_ENV, raw) + + with pytest.raises(ValueError, match="finite"): + worker_init_stall_warn_sec() + + +def test_stall_warn_knob_rejects_unparsable_values(monkeypatch): + """A misconfiguration must not be swallowed at the one place it matters. + + This PR's premise is that startup is where information gets lost; falling + back to the default here would hide the typo behind exactly the silence + the reporting exists to remove. + """ + monkeypatch.setenv(WORKER_INIT_STALL_WARN_ENV, "ten minutes") + + with pytest.raises(ValueError, match=WORKER_INIT_STALL_WARN_ENV): + worker_init_stall_warn_sec() + + +def test_non_finite_knob_never_arms_a_watchdog(monkeypatch): + """The consequence, pinned on the arming path the worker ranks take.""" + dumps = [] + monkeypatch.setattr( + "tensorrt_llm.executor.worker.print_all_stacks", + lambda **kwargs: dumps.append(kwargs.get("log")), + ) + monkeypatch.setattr("tensorrt_llm.executor.worker.logger", RecordingLogger()) + monkeypatch.setenv(WORKER_INIT_STALL_WARN_ENV, "nan") + init_done = threading.Event() + armed = [] + + try: + with pytest.raises(ValueError): + armed.append(worker_module._arm_worker_init_stall_watchdog(init_done)) + finally: + # If the guard ever regresses, a watchdog is now hot-looping on + # ``Event.wait(nan)`` (which returns immediately rather than sleeping); + # retire it here so the failure is a failure and not a flooded run. + init_done.set() + for watchdog in armed: + if watchdog is not None: + watchdog.join(timeout=5) + + assert not dumps, "a nan period armed a watchdog that dumped stacks" + + +def test_non_finite_knob_never_reaches_the_startup_loop(monkeypatch): + """...and on the proxy path, before any rank is spawned. + + The future is pre-failed so that a regression here fails instead of + hanging: without the guard the loop would run with ``next_warn_time`` + permanently ``None`` (``nan > 0`` is False) and never report or exit. + """ + dead = Future() + dead.set_exception(RuntimeError("rank 1 exited")) + proxy = _make_proxy(monkeypatch, futures=[dead]) + monkeypatch.setenv(WORKER_INIT_STALL_WARN_ENV, "nan") + + with pytest.raises(ValueError, match="finite"): + proxy._start_executor_workers({}) + + # The knob is read before anything is spawned, so nothing was left behind. + assert proxy.workers_started is False + assert proxy.mpi_session.submitted_kwargs is None + + +def test_worker_init_watchdog_dumps_this_ranks_stacks(monkeypatch): + """Drives the real ``worker._worker_init_stall_watchdog``.""" + dumps = [] + records = RecordingLogger() + monkeypatch.setattr( + "tensorrt_llm.executor.worker.print_all_stacks", + lambda **kwargs: dumps.append(kwargs.get("log")), + ) + monkeypatch.setattr("tensorrt_llm.executor.worker.logger", records) + + init_done = threading.Event() + watchdog = threading.Thread( + target=worker_module._worker_init_stall_watchdog, args=(init_done, 0.02), daemon=True + ) + watchdog.start() + try: + deadline = time.monotonic() + 10 + while len(dumps) < 2 and time.monotonic() < deadline: + time.sleep(0.01) + finally: + init_done.set() + watchdog.join(timeout=5) + + assert not watchdog.is_alive() + assert len(dumps) >= 2 + # A slow-but-healthy startup is not a fault: the report and the stack dump + # it triggers must both be WARNING, never ERROR. + assert records.errors == [] + assert any("has not finished initialization" in message for message in records.warnings), ( + records.warnings + ) + assert dumps == [records.warning] * len(dumps) + + +def test_worker_init_watchdog_is_silent_once_init_completes(monkeypatch): + dumps = [] + monkeypatch.setattr( + "tensorrt_llm.executor.worker.print_all_stacks", + lambda **kwargs: dumps.append(kwargs.get("log")), + ) + monkeypatch.setattr("tensorrt_llm.executor.worker.logger", RecordingLogger()) + + init_done = threading.Event() + init_done.set() + watchdog = threading.Thread( + target=worker_module._worker_init_stall_watchdog, args=(init_done, 0.02), daemon=True + ) + watchdog.start() + watchdog.join(timeout=5) + + assert not watchdog.is_alive() + assert dumps == [] + + +def test_worker_init_watchdog_is_armed_by_default(monkeypatch): + monkeypatch.delenv(WORKER_INIT_STALL_WARN_ENV, raising=False) + init_done = threading.Event() + + watchdog = worker_module._arm_worker_init_stall_watchdog(init_done) + try: + assert watchdog is not None + assert watchdog.daemon + assert watchdog.is_alive() + finally: + init_done.set() + # Completing initialization retires the watchdog immediately, even though + # its default period is far longer than this test. + watchdog.join(timeout=5) + assert not watchdog.is_alive() + + +def test_worker_init_watchdog_can_be_disabled(monkeypatch): + monkeypatch.setenv(WORKER_INIT_STALL_WARN_ENV, "0") + + assert worker_module._arm_worker_init_stall_watchdog(threading.Event()) is None + + +def test_print_all_stacks_honours_the_log_callable(): + """Drives the real ``tensorrt_llm._utils.print_all_stacks``. + + The init watchdog relies on being able to emit the dump at WARNING; the + default must stay ERROR for the existing callers. + """ + emitted = [] + + print_all_stacks(log=emitted.append) + + assert emitted + assert all("stack trace:" in message for message in emitted) + + +# --- worker_main disarm placement ------------------------------------------ +# +# ``worker_main`` needs a live MPI world, so the *placement* of the disarm is +# pinned structurally rather than behaviourally. This is not decoration: a +# leader that disarms when its constructor returns -- instead of when it has +# delivered the ready signal -- leaves the proxy waiting on a rank that has +# stopped reporting, and every behavioural test still passes. That regression +# was found on hardware, not here; these tests are what would catch it next +# time. + + +def _worker_main_ast(): + source = pathlib.Path(worker_module.__file__).read_text() + for node in ast.parse(source).body: + if isinstance(node, ast.FunctionDef) and node.name == "worker_main": + return node + raise AssertionError("worker_main not found in tensorrt_llm.executor.worker") + + +def _disarm_linenos(scope): + return sorted( + node.lineno + for node in ast.walk(scope) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "set" + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "worker_init_done" + ) + + +def _ready_send_lineno(worker_main): + for node in ast.walk(worker_main): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "notify_with_retry" + and node.args + and isinstance(node.args[0], ast.Name) + and node.args[0].id == "ready_msg" + ): + return node.lineno + raise AssertionError("ready-signal send not found in worker_main") + + +def _subordinate_only_disarm_linenos(worker_main): + """Disarms guarded by ``if not is_leader:``.""" + linenos = [] + for node in ast.walk(worker_main): + if not isinstance(node, ast.If): + continue + test = node.test + if ( + isinstance(test, ast.UnaryOp) + and isinstance(test.op, ast.Not) + and isinstance(test.operand, ast.Name) + and test.operand.id == "is_leader" + ): + for statement in node.body: + linenos.extend(_disarm_linenos(statement)) + return sorted(linenos) + + +def test_leader_stays_armed_until_the_ready_signal_is_sent(): + worker_main = _worker_main_ast() + ready_send = _ready_send_lineno(worker_main) + + disarms = _disarm_linenos(worker_main) + assert disarms, "worker_main never disarms the init watchdog" + + after_ready = [lineno for lineno in disarms if lineno > ready_send] + assert after_ready, ( + "no worker_init_done.set() after the ready signal is sent: the leader " + "would stop reporting while the proxy is still waiting for it" + ) + + +def test_no_unguarded_disarm_before_the_ready_signal(): + """The regression guard: a bare disarm after construction re-breaks this.""" + worker_main = _worker_main_ast() + ready_send = _ready_send_lineno(worker_main) + + subordinate_only = set(_subordinate_only_disarm_linenos(worker_main)) + in_failure_path = { + lineno + for handler in ( + node for node in ast.walk(worker_main) if isinstance(node, ast.ExceptHandler) + ) + for lineno in _disarm_linenos(handler) + } + + unguarded = [ + lineno + for lineno in _disarm_linenos(worker_main) + if lineno < ready_send and lineno not in subordinate_only and lineno not in in_failure_path + ] + assert not unguarded, ( + f"worker_init_done.set() at line(s) {unguarded} runs on the leader " + "before it has delivered the ready signal; the leader must stay armed " + "until then (subordinates disarm under 'if not is_leader:', and the " + "construction-failure path disarms inside its except handler)" + ) + + +def test_subordinate_disarms_after_construction(): + worker_main = _worker_main_ast() + ready_send = _ready_send_lineno(worker_main) + + subordinate_only = _subordinate_only_disarm_linenos(worker_main) + assert subordinate_only, ( + "no 'if not is_leader:' disarm: a subordinate blocks in " + "block_subordinates() forever and would report for the life of the job" + ) + assert all(lineno < ready_send for lineno in subordinate_only) From 3106a0d4845507fb8b120b7c41019a9bd6ba8caa Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Wed, 5 Aug 2026 02:15:11 +0000 Subject: [PATCH 2/2] [TRTLLM-13409][fix] keep the init stall watchdog armed when ready delivery fails Review catches from @zhaoyangwang-nvidia and CodeRabbit. `notify_with_retry()` returning False means the ready signal never reached the proxy -- the proxy is still blocked in its startup wait, looking for a signal that will never arrive. `worker_init_done.set()` ran regardless, disarming the stall watchdog on the one path where it is most needed and leaving a silent hang: this rank healthy and serving, the proxy waiting forever, nothing reporting either fact. Disarm only on successful delivery. On failure, keep the watchdog armed so the per-rank stall reports keep coming, and log at ERROR rather than WARNING -- this outcome leaves the job hung, which is not a warning-level event. Also annotate `print_all_stacks(log)`, which gained a parameter without types. Tests read the control flow from the AST: reaching this line for real needs a constructed worker, MPI ranks and an engine, but the property under test is which branch sets the event, and that is exactly what the AST shows. One test pins that the success path disarms (otherwise every healthy startup dumps stacks forever) and one that the failure path does not. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/_utils.py | 2 +- tensorrt_llm/executor/worker.py | 30 ++++++-- .../executor/test_proxy_worker_startup.py | 70 +++++++++++++++++++ 3 files changed, 94 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index d005d149e893..52ccdeb44365 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -700,7 +700,7 @@ def is_flashinfer_gdn_supported_arch(sm_version=None): return sm_version in (90, 100, 103) -def print_all_stacks(log=None): +def print_all_stacks(log: Optional[Callable[[str], None]] = None) -> None: """Print stack traces for all threads Args: diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index 53f9b44f73bc..b3312fb3d584 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -430,13 +430,29 @@ def notify_proxy_threads_to_quit(): # Send ready signal with confirmation ready_msg = (ready_signal, None, worker_process_identities) - if not worker_init_status_queue.notify_with_retry(ready_msg): - logger.warning( - "Failed to deliver ready signal to proxy, continuing anyway" - ) - # The proxy's startup wait ends here (or has been given up on); - # either way this rank has nothing more to contribute to it. - worker_init_done.set() + ready_delivered = worker_init_status_queue.notify_with_retry( + ready_msg) + if ready_delivered: + # The proxy's startup wait has ended, so this rank has + # nothing more to contribute to it. + worker_init_done.set() + else: + # Do NOT disarm here. Delivery failing is the one case + # where the proxy is still blocked in its startup wait + # with no idea why -- it is looking for a ready signal + # that will never arrive. Disarming would remove the last + # remaining source of information about that, leaving a + # silent hang: this rank healthy and serving, the proxy + # waiting forever, and nothing reporting either fact. + # Leaving the watchdog armed keeps the per-rank stall + # reports coming, which is exactly the diagnostic this PR + # exists to provide. + logger.error( + "Failed to deliver the ready signal to the proxy. This " + "rank is initialized and will continue serving, but the " + "proxy may still be waiting for a signal it will never " + "receive; leaving the init stall watchdog armed so the " + "condition keeps being reported.") if resource_governor_queue is not None: # Swap rank 0 to the proxy IPC queue after construction. # The resource-governor flag is already enabled on all diff --git a/tests/unittest/executor/test_proxy_worker_startup.py b/tests/unittest/executor/test_proxy_worker_startup.py index 8cd851e0345a..600c43efe8eb 100644 --- a/tests/unittest/executor/test_proxy_worker_startup.py +++ b/tests/unittest/executor/test_proxy_worker_startup.py @@ -619,3 +619,73 @@ def test_subordinate_disarms_after_construction(): "block_subordinates() forever and would report for the life of the job" ) assert all(lineno < ready_send for lineno in subordinate_only) + + +# --------------------------------------------------------------------------- +# The ready signal can fail to arrive. That is the one case where the worker +# must NOT disarm its stall watchdog: the proxy is still blocked, looking for +# a signal that will never come, and the per-rank stall reports are the only +# remaining evidence of it. +# --------------------------------------------------------------------------- +def _ready_delivery_block(src: str) -> ast.If: + """The `if ready_delivered:` block inside worker_main, from source. + + Read from the AST rather than executed: reaching this line for real needs + a constructed worker, MPI ranks and an engine. The property under test is + a control-flow one -- which branch sets the event -- and that is exactly + what the AST shows. + """ + tree = ast.parse(src) + for node in ast.walk(tree): + if ( + isinstance(node, ast.If) + and isinstance(node.test, ast.Name) + and node.test.id == "ready_delivered" + ): + return node + raise AssertionError("no `if ready_delivered:` branch found in worker_main") + + +def _sets_init_done(body) -> bool: + for node in body: + for sub in ast.walk(node): + if ( + isinstance(sub, ast.Call) + and isinstance(sub.func, ast.Attribute) + and sub.func.attr == "set" + and isinstance(sub.func.value, ast.Name) + and sub.func.value.id == "worker_init_done" + ): + return True + return False + + +def test_watchdog_is_disarmed_only_when_the_ready_signal_was_delivered(): + src = pathlib.Path(worker_module.__file__).read_text() + branch = _ready_delivery_block(src) + + assert _sets_init_done(branch.body), ( + "the success path must disarm the watchdog -- otherwise every healthy " + "startup keeps dumping stacks for the life of the process" + ) + assert not _sets_init_done(branch.orelse), ( + "the failure path must NOT disarm the watchdog: the proxy is still " + "waiting for a ready signal that never arrived, and disarming removes " + "the last thing reporting that" + ) + + +def test_failed_ready_delivery_is_logged_at_error(): + """A warning is not enough: this leaves the proxy hung.""" + src = pathlib.Path(worker_module.__file__).read_text() + branch = _ready_delivery_block(src) + levels = { + sub.func.attr + for node in branch.orelse + for sub in ast.walk(node) + if isinstance(sub, ast.Call) + and isinstance(sub.func, ast.Attribute) + and isinstance(sub.func.value, ast.Name) + and sub.func.value.id == "logger" + } + assert "error" in levels, f"expected logger.error on this path, saw {levels}"