diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 71574f3e1aa5..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 @@ -35,6 +38,7 @@ import numpy as np import nvtx +import psutil from mpi4py import MPI from mpi4py.util import pkl5 from typing_extensions import ParamSpec @@ -399,6 +403,79 @@ 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``. + """ + if sig is None: + sig = signal.SIGKILL + if sys.platform != "linux": + return + 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. + """ + 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 a7751cc6fe95..c2f52f1e270d 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -29,7 +29,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) @@ -721,6 +722,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 _get_next_client_id(self) -> int: client_id = super()._get_next_client_id() if self._num_frontends > 1: diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index f173aabc224a..dba76a3a9ba0 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 a23f8b9a8f46..38383d8f6f4d 100644 --- a/tests/integration/test_lists/test-db/l0_cpu_x86.yml +++ b/tests/integration/test_lists/test-db/l0_cpu_x86.yml @@ -14,4 +14,5 @@ l0_cpu_x86: orchestrator: mpi tests: - unittest/executor/test_rpc.py + - unittest/_utils/test_anti_zombie.py - unittest/executor/test_multi_frontend_routing.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..0d3119660dd6 --- /dev/null +++ b/tests/unittest/_utils/test_anti_zombie.py @@ -0,0 +1,184 @@ +# 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.cpu_only, + 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