[TRTLLM-13409][feat] anti-zombie worker cleanup (PR_SET_PDEATHSIG + tree-kill)

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 <linux/prctl.h>.
+_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
 

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:

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:

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

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

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