diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 71574f3e1aa5..52ccdeb44365 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: Optional[Callable[[str], None]] = None) -> 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..b3312fb3d584 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() @@ -372,10 +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" - ) + 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 new file mode 100644 index 000000000000..600c43efe8eb --- /dev/null +++ b/tests/unittest/executor/test_proxy_worker_startup.py @@ -0,0 +1,691 @@ +# 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) + + +# --------------------------------------------------------------------------- +# 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}"