From 92763eddbe6d9908eac10b8eba60fc1b04c8612c Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Wed, 24 Jun 2026 10:00:11 +0000 Subject: [PATCH 1/3] [TRTLLM-13409][feat] anti-zombie worker cleanup (PR_SET_PDEATHSIG + tree-kill) When the proxy / MPI launcher dies abruptly (e.g. a watchdog hard-kill or a pod-kill), mpi4py worker processes can be left orphaned, keeping their CUDA context and holding GPU memory until the next run OOMs at model load and blames the wrong PR. Add two anti-zombie helpers to tensorrt_llm._utils and wire them in: - set_parent_death_signal(): prctl(PR_SET_PDEATHSIG) so the kernel signals this process when its parent dies. Called at the top of worker_main, so every spawned executor worker self-terminates if its parent goes away. Linux-only; no-op elsewhere. Not called on the in-process LLM() path. - kill_process_tree(): SIGKILL a pid and all descendants (psutil-based, covers forked grandchildren), blocking until reaped. Wired into the proxy's pre_shutdown on the fatal-error path to reap the proxy's own descendant processes (postproc / local helpers). Note: MPIPoolSession workers are children of the MPI daemon, not of the proxy, so they are covered by PR_SET_PDEATHSIG rather than the proxy's tree-kill. Adds CPU-only unit tests for both helpers. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/_utils.py | 78 ++++++++ tensorrt_llm/executor/proxy.py | 19 +- tensorrt_llm/executor/worker.py | 12 +- .../test_lists/test-db/l0_cpu_arm.yml | 1 + .../test_lists/test-db/l0_cpu_x86.yml | 1 + tests/unittest/_utils/test_anti_zombie.py | 181 ++++++++++++++++++ 6 files changed, 290 insertions(+), 2 deletions(-) create mode 100644 tests/unittest/_utils/test_anti_zombie.py diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 67f1c9d3e950..391dc0c27ddb 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -400,6 +400,84 @@ def get_free_ports(num=1) -> List[int]: return ports +# Linux prctl option; see . +_PR_SET_PDEATHSIG = 1 + + +def set_parent_death_signal(sig=None, + expected_parent_pid: Optional[int] = None) -> None: + """Ask the kernel to send ``sig`` to this process when its parent dies. + + Linux-only (``prctl(PR_SET_PDEATHSIG)``); a no-op on other platforms. Used + by spawned executor workers so that an abruptly-killed parent (proxy / MPI + launcher) cannot leave them orphaned, holding their CUDA context and GPU + memory until the job's wall-clock pod-kill. + + ``PR_SET_PDEATHSIG`` only takes effect once this syscall runs: a parent + that died *before* arming never triggers it (this process was already + reparented). Spawners that know the intended direct parent's PID can pass + ``expected_parent_pid`` to close that startup race — after arming, + ``os.getppid()`` is compared against it and ``sig`` is delivered to this + process immediately on mismatch. MPI-launched workers cannot use this + (their direct parent is the MPI/Slurm daemon, whose PID the spawner does + not know); there the window stays open but is benign — that daemon dying + means the resource manager is already tearing the whole job down. + + Caveats: PR_SET_PDEATHSIG keys off the death of the parent *thread*, and + only covers the direct parent. Deeper / forked trees are handled by + ``kill_process_tree``. + """ + import signal + if sig is None: + sig = signal.SIGKILL + if sys.platform != "linux": + return + import ctypes + libc = ctypes.CDLL("libc.so.6", use_errno=True) + if libc.prctl(_PR_SET_PDEATHSIG, sig, 0, 0, 0) != 0: + errno = ctypes.get_errno() + raise OSError(errno, f"prctl(PR_SET_PDEATHSIG, {sig}) failed") + if expected_parent_pid is not None and os.getppid() != expected_parent_pid: + # The parent died before the signal was armed (we were reparented); + # deliver the death signal ourselves rather than lingering orphaned. + os.kill(os.getpid(), sig) + + +def kill_process_tree(pid: int, + *, + include_parent: bool = True, + wait_timeout: float = 60.0) -> None: + """SIGKILL ``pid`` and all of its descendants, blocking until reaped. + + Uses ``psutil.children(recursive=True)`` so forked grandchildren are covered + (raw PR_SET_PDEATHSIG only covers the direct parent-child link). Mirrors the + anti-zombie cleanup used by SGLang / vLLM. + """ + import time + + import psutil + try: + parent = psutil.Process(pid) + except psutil.NoSuchProcess: + return + children = parent.children(recursive=True) + targets = children + ([parent] if include_parent else []) + for p in targets: + try: + p.kill() + except psutil.NoSuchProcess: + pass + deadline = time.time() + wait_timeout + for p in targets: + try: + p.wait(timeout=max(0.0, deadline - time.time())) + except psutil.TimeoutExpired: + logger.error(f"kill_process_tree: pid {p.pid} did not exit within " + f"{wait_timeout}s; abandoning.") + except psutil.NoSuchProcess: + pass + + # mpi4py only exports MPI_COMM_TYPE_SHARED, so we define OMPI_COMM_TYPE_HOST here OMPI_COMM_TYPE_HOST = 9 diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index 07d0ea0414fa..57ab5c1e1ffc 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -27,7 +27,8 @@ from tensorrt_llm.logger import logger -from .._utils import customized_gc_thresholds, mpi_rank, nvtx_range_debug +from .._utils import (customized_gc_thresholds, kill_process_tree, mpi_rank, + nvtx_range_debug) from ..llmapi.mpi_session import (MpiCommSession, MpiPoolSession, MpiSession, RemoteMpiCommSessionClient, validate_session_world_size) @@ -642,6 +643,22 @@ def pre_shutdown(self): if not self.mpi_futures or any(not f.done() for f in self.mpi_futures): self.request_queue.put_noblock(None, retry=4) + # Anti-zombie: when shutting down after a fatal error, the graceful + # sentinel above may never be drained (workers wedged / dead). Reap any + # of the proxy's own descendant processes (e.g. postproc workers, local + # helpers) so they don't orphan and leak GPU memory. include_parent is + # False so we don't kill the proxy mid-cleanup. MPI-spawned workers are + # not the proxy's children and are covered by PR_SET_PDEATHSIG instead. + if self._fatal_error is not None: + try: + kill_process_tree(os.getpid(), + include_parent=False, + wait_timeout=10.0) + except Exception as e: # noqa: BLE001 - cleanup must not raise + logger_debug( + f"kill_process_tree during pre_shutdown failed: " + f"{e}\n", "yellow") + def shutdown(self): if not self.workers_started: return diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index 342168664ef1..709fe9d0aade 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -11,7 +11,8 @@ from tensorrt_llm.logger import logger -from .._utils import mpi_comm, mpi_rank, print_all_stacks +from .._utils import (mpi_comm, mpi_rank, print_all_stacks, + set_parent_death_signal) from ..bindings import executor as tllm from ..llmapi.llm_args import BaseLlmArgs from ..llmapi.mpi_session import set_mpi_session_cpp @@ -179,6 +180,15 @@ def worker_main( hmac_key: bytes = b"", ) -> None: + # Anti-zombie: if our parent (proxy / MPI launcher) dies abruptly, have the + # kernel SIGKILL this worker so it can't orphan and leak GPU memory. + try: + set_parent_death_signal() + except OSError as e: + logger.warning( + f"PR_SET_PDEATHSIG setup failed: {e}; orphaned workers may leak " + "GPU memory if the parent dies abruptly.") + def _print_stacks(): counter = 0 while True: diff --git a/tests/integration/test_lists/test-db/l0_cpu_arm.yml b/tests/integration/test_lists/test-db/l0_cpu_arm.yml index ba693acaf756..630da1175bc6 100644 --- a/tests/integration/test_lists/test-db/l0_cpu_arm.yml +++ b/tests/integration/test_lists/test-db/l0_cpu_arm.yml @@ -14,3 +14,4 @@ l0_cpu_arm: orchestrator: mpi tests: - unittest/executor/test_rpc.py + - unittest/_utils/test_anti_zombie.py diff --git a/tests/integration/test_lists/test-db/l0_cpu_x86.yml b/tests/integration/test_lists/test-db/l0_cpu_x86.yml index 9a39993347b8..9446c691ab74 100644 --- a/tests/integration/test_lists/test-db/l0_cpu_x86.yml +++ b/tests/integration/test_lists/test-db/l0_cpu_x86.yml @@ -14,3 +14,4 @@ l0_cpu_x86: orchestrator: mpi tests: - unittest/executor/test_rpc.py + - unittest/_utils/test_anti_zombie.py diff --git a/tests/unittest/_utils/test_anti_zombie.py b/tests/unittest/_utils/test_anti_zombie.py new file mode 100644 index 000000000000..6cfb1c5e4add --- /dev/null +++ b/tests/unittest/_utils/test_anti_zombie.py @@ -0,0 +1,181 @@ +# 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. +"""Anti-zombie helpers: PR_SET_PDEATHSIG and kill_process_tree (no GPU).""" + +import signal +import subprocess +import sys +import time + +import psutil +import pytest + +pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="PR_SET_PDEATHSIG is Linux-only") + + +def _alive(pid: int) -> bool: + return psutil.pid_exists(pid) and psutil.Process(pid).status() != psutil.STATUS_ZOMBIE + + +def _wait_gone(pid: int, timeout: float = 10.0) -> bool: + deadline = time.time() + timeout + while time.time() < deadline: + if not _alive(pid): + return True + time.sleep(0.1) + return not _alive(pid) + + +# Parent spawns a child that arms PR_SET_PDEATHSIG, prints its pid, then sleeps. +_PARENT_SRC = """ +import subprocess, sys, time +child = subprocess.Popen([sys.executable, "-c", ''' +import signal, time +from tensorrt_llm._utils import set_parent_death_signal +set_parent_death_signal(signal.SIGKILL) +print("CHILD_READY", flush=True) +time.sleep(300) +'''], stdout=subprocess.PIPE, text=True) +# Relay the child's readiness + pid to our stdout. +child.stdout.readline() # wait for CHILD_READY +print(f"CHILD_PID={child.pid}", flush=True) +time.sleep(300) +""" + + +def test_prctl_kills_child_when_parent_dies() -> None: + parent = subprocess.Popen( + [sys.executable, "-c", _PARENT_SRC], stdout=subprocess.PIPE, text=True + ) + child_pid = None + try: + line = parent.stdout.readline().strip() + assert line.startswith("CHILD_PID="), f"unexpected: {line!r}" + child_pid = int(line.split("=", 1)[1]) + assert _alive(child_pid) + + # Kill the parent; the kernel should SIGKILL the child via PDEATHSIG. + parent.kill() + parent.wait(timeout=10) + + assert _wait_gone(child_pid, timeout=10.0), f"child {child_pid} survived its parent's death" + finally: + if parent.poll() is None: + parent.kill() + # Best-effort: if an assertion failed after the parent was reaped, the + # child may still be sleeping — don't leave it behind for 5 minutes. + if child_pid is not None: + try: + psutil.Process(child_pid).kill() + except psutil.Error: + pass + + +# Builds a 3-level tree (top -> child -> grandchild), all sleeping, and prints +# each pid so the test can verify kill_process_tree reaps the whole tree. +_TREE_SRC = """ +import subprocess, sys, time +gc_src = "import time; print('G', flush=True); time.sleep(300)" +ch_src = ( + "import subprocess, sys, time; " + "g = subprocess.Popen([sys.executable, '-c', %r]); " + "print('CHILD_PID=' + str(__import__('os').getpid()), flush=True); " + "print('GRANDCHILD_PID=' + str(g.pid), flush=True); " + "time.sleep(300)" +) % gc_src +child = subprocess.Popen([sys.executable, "-c", ch_src], stdout=subprocess.PIPE, text=True) +import os +print("TOP_PID=" + str(os.getpid()), flush=True) +for _ in range(2): + print(child.stdout.readline().strip(), flush=True) +time.sleep(300) +""" + + +def test_kill_process_tree_reaps_grandchildren() -> None: + from tensorrt_llm._utils import kill_process_tree + + top = subprocess.Popen([sys.executable, "-c", _TREE_SRC], stdout=subprocess.PIPE, text=True) + pids = {} + try: + for _ in range(3): + line = top.stdout.readline().strip() + key, _, val = line.partition("=") + pids[key] = int(val) + assert {"TOP_PID", "CHILD_PID", "GRANDCHILD_PID"} <= set(pids) + for pid in pids.values(): + assert _alive(pid), f"{pid} not alive at setup" + + kill_process_tree(pids["TOP_PID"], include_parent=True, wait_timeout=10.0) + + for name, pid in pids.items(): + assert _wait_gone(pid, timeout=10.0), f"{name} ({pid}) not reaped" + finally: + if top.poll() is None: + top.kill() + # Best-effort cleanup if the assertion failed mid-way. + for pid in pids.values(): + try: + psutil.Process(pid).kill() + except psutil.Error: + pass + + +def test_set_parent_death_signal_idempotent() -> None: + """Calling it must not raise. Run in a subprocess so we don't arm + PR_SET_PDEATHSIG on the pytest worker itself.""" + src = ( + "import signal\n" + "from tensorrt_llm._utils import set_parent_death_signal\n" + "set_parent_death_signal(signal.SIGTERM)\n" + "set_parent_death_signal(signal.SIGTERM)\n" + "print('OK')\n" + ) + # Generous timeout: the subprocess pays a cold `import tensorrt_llm`, which + # alone can take ~a minute on slower hosts, before the prctl calls run. + proc = subprocess.run([sys.executable, "-c", src], timeout=300, capture_output=True, text=True) + assert proc.returncode == 0, proc.stderr + assert "OK" in proc.stdout + + +def test_prearming_parent_death_detected() -> None: + """Regression for the arming race: PR_SET_PDEATHSIG only takes effect after + the prctl syscall, so a parent that died first would never trigger it. When + the spawner supplies expected_parent_pid, a reparented process must detect + the mismatch right after arming and deliver the signal to itself.""" + src = ( + "import os, signal\n" + "from tensorrt_llm._utils import set_parent_death_signal\n" + "# Simulate 'parent already died before arming': expect a parent PID\n" + "# that is guaranteed not to be our actual current parent.\n" + "set_parent_death_signal(signal.SIGKILL, expected_parent_pid=os.getppid() + 1)\n" + "print('UNREACHABLE')\n" + ) + proc = subprocess.run([sys.executable, "-c", src], timeout=300, capture_output=True, text=True) + assert proc.returncode == -signal.SIGKILL, (proc.returncode, proc.stderr) + assert "UNREACHABLE" not in proc.stdout + + # And the happy path: the expected parent matches, no self-kill. + src_ok = ( + "import os, signal\n" + "from tensorrt_llm._utils import set_parent_death_signal\n" + "set_parent_death_signal(signal.SIGKILL, expected_parent_pid=os.getppid())\n" + "print('OK')\n" + ) + proc = subprocess.run( + [sys.executable, "-c", src_ok], timeout=300, capture_output=True, text=True + ) + assert proc.returncode == 0, proc.stderr + assert "OK" in proc.stdout From 2f1bfa96100e67ffcea17e0977a4b36a18406747 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:02:03 +0000 Subject: [PATCH 2/3] [TRTLLM-13409][test] mark anti-zombie tests as cpu_only The CPU-Generic-arm CI stage invokes pytest with -m 'cpu_only and not disabled', so tests without the cpu_only marker are deselected and the stage exits with pytest code 5, which the outer runner reports as 'AssertionError: failure reported in unittests'. Add pytest.mark.cpu_only alongside the existing linux-only skipif so the four tests in tests/unittest/_utils/test_anti_zombie.py are actually selected on the CPU stages. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tests/unittest/_utils/test_anti_zombie.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_utils/test_anti_zombie.py b/tests/unittest/_utils/test_anti_zombie.py index 6cfb1c5e4add..0d3119660dd6 100644 --- a/tests/unittest/_utils/test_anti_zombie.py +++ b/tests/unittest/_utils/test_anti_zombie.py @@ -22,7 +22,10 @@ import psutil import pytest -pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="PR_SET_PDEATHSIG is Linux-only") +pytestmark = [ + pytest.mark.cpu_only, + pytest.mark.skipif(sys.platform != "linux", reason="PR_SET_PDEATHSIG is Linux-only"), +] def _alive(pid: int) -> bool: From 457a7d76e4c3bdc11f0226a50ac510c7b52b638f Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Thu, 23 Jul 2026 02:35:43 +0000 Subject: [PATCH 3/3] [TRTLLM-13409][chore] hoist anti-zombie helpers' imports to module top Addresses review nit from @brnguyen2: signal, ctypes, time (stdlib) and psutil (third-party) were lazily imported inside set_parent_death_signal and kill_process_tree. Move them to the module-level imports for consistency with the rest of _utils.py, which already imports numpy, nvtx, mpi4py unconditionally. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/_utils.py | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 391dc0c27ddb..c738d1bde9a3 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -13,17 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. import copy +import ctypes import gc import inspect import json import linecache import math import os +import signal import socket import struct import sys import tempfile import threading +import time import trace import traceback import weakref @@ -31,11 +34,11 @@ from ctypes import byref from enum import EnumMeta from functools import lru_cache, partial, wraps -from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Sequence, TypeVar, Union import numpy as np import nvtx +import psutil from mpi4py import MPI from mpi4py.util import pkl5 from typing_extensions import ParamSpec @@ -64,7 +67,7 @@ has_nvml = False # isort: on -from tensorrt_llm.bindings import DataType, GptJsonConfig, LayerType +from tensorrt_llm.bindings import DataType, LayerType from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE from tensorrt_llm.logger import logger @@ -427,12 +430,10 @@ def set_parent_death_signal(sig=None, only covers the direct parent. Deeper / forked trees are handled by ``kill_process_tree``. """ - import signal if sig is None: sig = signal.SIGKILL if sys.platform != "linux": return - import ctypes libc = ctypes.CDLL("libc.so.6", use_errno=True) if libc.prctl(_PR_SET_PDEATHSIG, sig, 0, 0, 0) != 0: errno = ctypes.get_errno() @@ -453,9 +454,6 @@ def kill_process_tree(pid: int, (raw PR_SET_PDEATHSIG only covers the direct parent-child link). Mirrors the anti-zombie cleanup used by SGLang / vLLM. """ - import time - - import psutil try: parent = psutil.Process(pid) except psutil.NoSuchProcess: @@ -853,13 +851,6 @@ def __contains__(cls, item): return True -def supports_inflight_batching(engine_dir): - config_path = Path(engine_dir) / "config.json" - json_config = GptJsonConfig.parse_file(config_path) - model_config = json_config.model_config - return model_config.supports_inflight_batching - - class QuantModeWrapper: def __init__(self, objs):