scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(proxy_cli): reap orphaned prisma query-engine processes when a worker dies (#33424)

Browse source 
* fix(proxy_cli): reap orphaned prisma query-engine processes when a worker dies

When the proxy runs multi-worker (uvicorn multiprocess supervisor or the
gunicorn arbiter), a worker that crashes or is force-killed never runs its
in-process atexit cleanup, so its prisma query-engine subprocess reparents
to PID 1 and keeps its database connection pool established forever while
the replacement worker opens a fresh pool. Active DB connections then grow
past database_connection_pool_limit with every worker death.

Run a reaper thread in the supervisor process that marks itself a child
subreaper on Linux, scans for adopted query-engine children whose worker
is gone, and terminates them with SIGTERM escalating to SIGKILL after a
bounded grace period. Engines owned by live workers are children of those
workers, never of the supervisor, so they are structurally out of reach.

Resolves LIT-4449
Fixes https://github.com/BerriAI/litellm/issues/33023

* fix(proxy_cli): address review findings on the query-engine reaper

Make start_query_engine_reaper idempotent, reap simultaneous orphans
under one shared grace period instead of serially, and log when a PID
survives SIGKILL. Also regenerate schema.d.ts for the update_team
docstring line that documents the existing mcp_rpm_limit param (fixes
the walk-order-dependent documentation CI failure) and avoid a cast in
the prctl wrapper

* test(proxy): fix reaper idempotency-test isolation and widen coverage

The daemon-thread startup test now stubs threading.enumerate so a
reaper thread left running by an earlier test in the same xdist worker
cannot satisfy the idempotency guard and skip the code under test. Add
coverage for stat-file truncation, non-numeric ppid, non-child reap,
signal-to-dead-pid, subreaper capability, and reaper-loop resilience
This commit is contained in:
Yassin Kortam 2026-07-16 13:16:41 -07:00 • committed by GitHub
parent fba7ac4428
commit a8ae515bee
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 539 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

212
litellm/proxy/db/query_engine_reaper.py Normal file
View file

@ -0,0 +1,212 @@
"""Supervisor-side reaper for orphaned Prisma query-engine processes.
Each proxy worker owns a Prisma query-engine subprocess whose only cleanup
hook is an in-process ``atexit`` handler. When a multi-worker supervisor
(uvicorn's multiprocess manager, the gunicorn arbiter) force-kills a hung or
crashed worker, that handler never runs: the engine reparents to the nearest
subreaper (PID 1 in a container, which is the supervisor itself under the
standard docker entrypoint) and keeps its database connection pool
established forever, while the replacement worker opens a fresh pool. Over
repeated worker deaths the active database connections grow without bound.
The reaper runs only in the supervisor process, where a query-engine process
can never be a legitimate direct child: workers own their engines, and the
supervisor never starts one. Any direct child whose command name begins with
``query-engine`` is therefore an adopted orphan and is terminated
(SIGTERM, bounded grace, SIGKILL) and reaped. On Linux the supervisor also
marks itself a child subreaper so orphans reparent to it even when it is not
PID 1.
Linux-only by construction (``/proc`` scan, ``prctl``); a no-op elsewhere.
"""
import ctypes
import os
import signal
import sys
import threading
import time
from typing import Optional
from litellm._logging import verbose_proxy_logger
QUERY_ENGINE_COMM_PREFIX = "query-engine"
REAPER_SCAN_INTERVAL_SECONDS = 5.0
SIGTERM_GRACE_SECONDS = 10.0
PR_SET_CHILD_SUBREAPER = 36
def set_child_subreaper() -> bool:
"""Mark this process as a child subreaper so orphaned descendants
reparent to it instead of PID 1. Best-effort: when it fails (or on
non-Linux) the reaper still covers the containerized case where the
supervisor already is PID 1."""
if not sys.platform.startswith("linux"):
return False
try:
libc = ctypes.CDLL(None, use_errno=True)
result: int = libc.prctl( # pyright: ignore[reportAny] # ctypes types foreign calls as Any; default restype is c_int
PR_SET_CHILD_SUBREAPER, 1, 0, 0, 0
)
return result == 0
except (OSError, AttributeError):
return False
def _read_comm_and_ppid(pid: int, proc_root: str) -> Optional[tuple[str, int]]:
try:
with open(f"{proc_root}/{pid}/stat", encoding="ascii", errors="replace") as stat_file:
data = stat_file.read()
except (FileNotFoundError, ProcessLookupError, PermissionError, OSError):
return None
lparen = data.find("(")
rparen = data.rfind(")")
if lparen == -1 or rparen == -1 or rparen < lparen:
return None
comm = data[lparen + 1 : rparen]
fields = data[rparen + 2 :].split()
if len(fields) < 2:
return None
try:
ppid = int(fields[1])
except ValueError:
return None
return comm, ppid
def list_orphaned_engine_pids(parent_pid: int, proc_root: str = "/proc") -> tuple[int, ...]:
"""PIDs of direct children of ``parent_pid`` whose command name marks
them as Prisma query engines. In the supervisor these are always
adopted orphans: live engines are children of workers, not of the
supervisor."""
try:
entries = os.listdir(proc_root)
except (FileNotFoundError, OSError):
return ()
candidate_pids = (int(entry) for entry in entries if entry.isdigit())
return tuple(
pid
for pid in candidate_pids
if (info := _read_comm_and_ppid(pid, proc_root)) is not None
and info[1] == parent_pid
and info[0].startswith(QUERY_ENGINE_COMM_PREFIX)
)
def _try_reap(pid: int) -> bool:
try:
reaped_pid, _ = os.waitpid(pid, os.WNOHANG)
except ChildProcessError:
return True
except OSError:
return True
return reaped_pid == pid
def _send_signal(pid: int, signum: int) -> None:
try:
os.kill(pid, signum)
except (ProcessLookupError, PermissionError, OSError):
pass
def _await_reaped(pids: tuple[int, ...], timeout_seconds: float) -> tuple[int, ...]:
"""Poll until every PID is reaped or the shared deadline passes.
Returns the PIDs still alive at the deadline."""
deadline = time.monotonic() + timeout_seconds
remaining = pids
while remaining and time.monotonic() < deadline:
remaining = tuple(pid for pid in remaining if not _try_reap(pid))
if remaining:
time.sleep(0.2)
return remaining
def terminate_and_reap(pid: int, grace_seconds: float = SIGTERM_GRACE_SECONDS) -> None:
"""SIGTERM the orphaned engine, escalate to SIGKILL after the grace
period, and reap it so it does not linger as a zombie."""
terminate_and_reap_all((pid,), grace_seconds=grace_seconds)
def terminate_and_reap_all(
pids: tuple[int, ...],
grace_seconds: float = SIGTERM_GRACE_SECONDS,
) -> None:
"""Terminate a batch of orphaned engines concurrently: SIGTERM all of
them, share one grace period, SIGKILL the stragglers, and reap. The
shared deadline keeps cleanup time bounded when several workers die
at once instead of paying the grace period once per orphan."""
for pid in pids:
verbose_proxy_logger.warning(
"Reaping orphaned prisma query-engine PID %s (its worker process exited without cleanup).",
pid,
)
_send_signal(pid, signal.SIGTERM)
survivors = _await_reaped(pids, grace_seconds)
if not survivors:
return
for pid in survivors:
verbose_proxy_logger.warning(
"Orphaned prisma query-engine PID %s did not exit within %.1fs of SIGTERM; sending SIGKILL.",
pid,
grace_seconds,
)
_send_signal(pid, signal.SIGKILL)
unkillable = _await_reaped(survivors, 5.0)
for pid in unkillable:
verbose_proxy_logger.error(
"Orphaned prisma query-engine PID %s survived SIGKILL; will retry on the next scan.",
pid,
)
def reap_orphaned_engines(parent_pid: int, proc_root: str = "/proc") -> tuple[int, ...]:
"""One scan-and-reap pass. Returns the PIDs it acted on."""
orphaned_pids = list_orphaned_engine_pids(parent_pid, proc_root=proc_root)
if orphaned_pids:
terminate_and_reap_all(orphaned_pids)
return orphaned_pids
def _reaper_loop(parent_pid: int) -> None:
while True:
try:
reap_orphaned_engines(parent_pid)
except Exception as scan_error: # noqa: BLE001 # reaper thread must survive any scan failure
verbose_proxy_logger.debug("Orphaned query-engine scan failed: %s", scan_error)
time.sleep(REAPER_SCAN_INTERVAL_SECONDS)
REAPER_THREAD_NAME = "litellm-orphan-query-engine-reaper"
def start_query_engine_reaper() -> Optional[threading.Thread]:
"""Start the reaper daemon thread in the supervisor process.
Must only be called from a process that never hosts the proxy app
itself (uvicorn with ``workers > 1``, the gunicorn arbiter): with a
single in-process uvicorn worker the query engine is a legitimate
direct child and must not be touched. Idempotent: a reaper already
running in this process is returned instead of starting a second one.
"""
if not sys.platform.startswith("linux"):
return None
existing = next(
(thread for thread in threading.enumerate() if thread.name == REAPER_THREAD_NAME),
None,
)
if existing is not None:
return existing
set_child_subreaper()
reaper_thread = threading.Thread(
target=_reaper_loop,
args=(os.getpid(),),
daemon=True,
name=REAPER_THREAD_NAME,
)
reaper_thread.start()
verbose_proxy_logger.info(
"Started orphaned prisma query-engine reaper in supervisor process %s.",
os.getpid(),
)
return reaper_thread

4
litellm/proxy/proxy_cli.py
View file

@ -15,6 +15,7 @@ from dotenv import load_dotenv
import litellm
from litellm.constants import DEFAULT_NUM_WORKERS_LITELLM_PROXY
from litellm.proxy.db.query_engine_reaper import start_query_engine_reaper
from litellm.secret_managers.main import get_secret_bool
if TYPE_CHECKING:
@ -495,6 +496,7 @@ class ProxyInitializationHelpers:
gunicorn_options["certfile"] = ssl_certfile_path
gunicorn_options["keyfile"] = ssl_keyfile_path
start_query_engine_reaper()
StandaloneApplication(app=app, options=gunicorn_options).run() # Run gunicorn
@staticmethod
@ -1261,6 +1263,8 @@ def run_server(
if reload:
ProxyInitializationHelpers._configure_dev_reload(uvicorn_args, config)
if num_workers > 1:
start_query_engine_reaper()
uvicorn.run(
**uvicorn_args,
workers=num_workers,

244
tests/test_litellm/proxy/db/test_query_engine_reaper.py Normal file
View file

@ -0,0 +1,244 @@
import os
import signal
import subprocess
import sys
import time
from unittest.mock import MagicMock, patch
import pytest
from litellm.proxy.db.query_engine_reaper import (
REAPER_THREAD_NAME,
_read_comm_and_ppid,
_reaper_loop,
_send_signal,
_try_reap,
list_orphaned_engine_pids,
reap_orphaned_engines,
set_child_subreaper,
start_query_engine_reaper,
terminate_and_reap,
terminate_and_reap_all,
)
def _write_stat(proc_root, pid, comm, ppid):
pid_dir = proc_root / str(pid)
pid_dir.mkdir()
(pid_dir / "stat").write_text(f"{pid} ({comm}) S {ppid} {pid} {pid} 0 -1 4194304 100 0 0 0")
class TestReadCommAndPpid:
def test_parses_comm_and_ppid(self, tmp_path):
_write_stat(tmp_path, 137, "query-engine-de", 81)
assert _read_comm_and_ppid(137, str(tmp_path)) == ("query-engine-de", 81)
def test_comm_containing_parens_and_spaces(self, tmp_path):
_write_stat(tmp_path, 42, "weird) (name", 1)
assert _read_comm_and_ppid(42, str(tmp_path)) == ("weird) (name", 1)
def test_missing_pid_returns_none(self, tmp_path):
assert _read_comm_and_ppid(999, str(tmp_path)) is None
def test_malformed_stat_returns_none(self, tmp_path):
pid_dir = tmp_path / "55"
pid_dir.mkdir()
(pid_dir / "stat").write_text("garbage with no parens")
assert _read_comm_and_ppid(55, str(tmp_path)) is None
def test_truncated_fields_after_comm_returns_none(self, tmp_path):
pid_dir = tmp_path / "56"
pid_dir.mkdir()
(pid_dir / "stat").write_text("56 (proc) S")
assert _read_comm_and_ppid(56, str(tmp_path)) is None
def test_non_numeric_ppid_returns_none(self, tmp_path):
pid_dir = tmp_path / "57"
pid_dir.mkdir()
(pid_dir / "stat").write_text("57 (proc) S notanint 57")
assert _read_comm_and_ppid(57, str(tmp_path)) is None
class TestListOrphanedEnginePids:
def test_finds_only_engine_children_of_parent(self, tmp_path):
_write_stat(tmp_path, 137, "query-engine-de", 1)
_write_stat(tmp_path, 138, "query-engine-de", 1)
_write_stat(tmp_path, 260, "python", 1)
_write_stat(tmp_path, 285, "query-engine-de", 260)
(tmp_path / "not-a-pid").mkdir()
assert sorted(list_orphaned_engine_pids(1, proc_root=str(tmp_path))) == [137, 138]
def test_no_matches_returns_empty(self, tmp_path):
_write_stat(tmp_path, 260, "python", 1)
assert list_orphaned_engine_pids(1, proc_root=str(tmp_path)) == ()
def test_missing_proc_root_returns_empty(self, tmp_path):
assert list_orphaned_engine_pids(1, proc_root=str(tmp_path / "absent")) == ()
class TestSetChildSubreaper:
def test_matches_platform_capability(self):
result = set_child_subreaper()
if sys.platform.startswith("linux"):
assert result is True
else:
assert result is False
@pytest.mark.skipif(sys.platform == "win32", reason="POSIX signals and waitpid")
class TestSignalHelpers:
def test_try_reap_true_for_non_child_pid(self):
assert _try_reap(1) is True
def test_send_signal_swallows_missing_pid(self):
child = subprocess.Popen([sys.executable, "-c", "pass"])
child.wait()
_send_signal(child.pid, signal.SIGTERM)
class TestReaperLoop:
def test_survives_scan_failure_and_continues(self):
calls = []
def flaky_scan(parent_pid, proc_root="/proc"):
calls.append(parent_pid)
if len(calls) == 1:
raise RuntimeError("scan blew up")
raise KeyboardInterrupt
with (
patch(
"litellm.proxy.db.query_engine_reaper.reap_orphaned_engines",
side_effect=flaky_scan,
),
patch("litellm.proxy.db.query_engine_reaper.time.sleep"),
pytest.raises(KeyboardInterrupt),
):
_reaper_loop(1234)
assert calls == [1234, 1234]
@pytest.mark.skipif(sys.platform == "win32", reason="POSIX signals and waitpid")
class TestTerminateAndReap:
def test_sigterm_terminates_and_reaps_child(self):
child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(300)"])
terminate_and_reap(child.pid, grace_seconds=10.0)
with pytest.raises((ChildProcessError, OSError)):
os.waitpid(child.pid, os.WNOHANG)
child.returncode = -signal.SIGTERM
def test_escalates_to_sigkill_when_sigterm_ignored(self):
child = subprocess.Popen(
[
sys.executable,
"-c",
"import signal, time; signal.signal(signal.SIGTERM, signal.SIG_IGN); time.sleep(300)",
]
)
deadline = time.monotonic() + 5
while time.monotonic() < deadline:
probe = subprocess.run(
[sys.executable, "-c", f"import os, signal; os.kill({child.pid}, 0)"],
capture_output=True,
)
if probe.returncode == 0:
break
time.sleep(0.05)
time.sleep(0.3)
terminate_and_reap(child.pid, grace_seconds=0.5)
with pytest.raises((ChildProcessError, OSError)):
os.waitpid(child.pid, os.WNOHANG)
child.returncode = -signal.SIGKILL
class TestReapOrphanedEngines:
def test_terminates_each_orphan(self, tmp_path):
_write_stat(tmp_path, 137, "query-engine-de", 1)
_write_stat(tmp_path, 138, "query-engine-de", 1)
_write_stat(tmp_path, 285, "query-engine-de", 260)
with patch("litellm.proxy.db.query_engine_reaper.terminate_and_reap_all") as mock_terminate:
acted_on = reap_orphaned_engines(1, proc_root=str(tmp_path))
assert sorted(acted_on) == [137, 138]
assert sorted(mock_terminate.call_args.args[0]) == [137, 138]
def test_no_orphans_no_kills(self, tmp_path):
_write_stat(tmp_path, 285, "query-engine-de", 260)
with patch("litellm.proxy.db.query_engine_reaper.terminate_and_reap_all") as mock_terminate:
assert reap_orphaned_engines(1, proc_root=str(tmp_path)) == ()
mock_terminate.assert_not_called()
@pytest.mark.skipif(sys.platform == "win32", reason="POSIX signals and waitpid")
class TestTerminateAndReapAll:
def test_batch_shares_one_grace_period(self):
children = [
subprocess.Popen(
[
sys.executable,
"-c",
"import signal, time; signal.signal(signal.SIGTERM, signal.SIG_IGN); time.sleep(300)",
]
)
for _ in range(3)
]
time.sleep(0.5)
start = time.monotonic()
terminate_and_reap_all(tuple(child.pid for child in children), grace_seconds=1.0)
elapsed = time.monotonic() - start
assert elapsed < 3.0
for child in children:
with pytest.raises((ChildProcessError, OSError)):
os.waitpid(child.pid, os.WNOHANG)
child.returncode = -signal.SIGKILL
class TestStartQueryEngineReaper:
def test_noop_on_non_linux(self):
with patch("litellm.proxy.db.query_engine_reaper.sys.platform", "darwin"):
assert start_query_engine_reaper() is None
def test_starts_daemon_thread_on_linux(self):
with (
patch("litellm.proxy.db.query_engine_reaper.sys.platform", "linux"),
patch(
"litellm.proxy.db.query_engine_reaper.threading.enumerate",
return_value=[],
),
patch("litellm.proxy.db.query_engine_reaper.set_child_subreaper") as mock_subreaper,
patch("litellm.proxy.db.query_engine_reaper.threading.Thread") as mock_thread_cls,
):
thread = start_query_engine_reaper()
mock_subreaper.assert_called_once()
mock_thread_cls.assert_called_once()
assert mock_thread_cls.call_args.kwargs["daemon"] is True
assert mock_thread_cls.call_args.kwargs["args"] == (os.getpid(),)
mock_thread_cls.return_value.start.assert_called_once()
assert thread is mock_thread_cls.return_value
def test_second_call_returns_existing_thread(self):
existing = MagicMock()
existing.name = REAPER_THREAD_NAME
with (
patch("litellm.proxy.db.query_engine_reaper.sys.platform", "linux"),
patch(
"litellm.proxy.db.query_engine_reaper.threading.enumerate",
return_value=[existing],
),
patch("litellm.proxy.db.query_engine_reaper.threading.Thread") as mock_thread_cls,
):
thread = start_query_engine_reaper()
assert thread is existing
mock_thread_cls.assert_not_called()

79
tests/test_litellm/proxy/test_proxy_cli.py
View file

@ -1557,6 +1557,85 @@ class TestProxyInitializationHelpers:
mock_uvicorn_run.assert_called_once()
class TestQueryEngineReaperWiring:
def _invoke_run_server(self, args):
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
clean_env = {
k: v
for k, v in os.environ.items()
if k not in ("DATABASE_URL", "DIRECT_URL")
}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
},
),
patch("uvicorn.run") as mock_uvicorn_run,
patch(
"litellm.proxy.proxy_cli.start_query_engine_reaper"
) as mock_start_reaper,
patch(
"litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args"
) as mock_get_args,
):
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(run_server, args)
return result, mock_uvicorn_run, mock_start_reaper
def test_multi_worker_uvicorn_starts_reaper(self):
result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server(
["--local", "--num_workers", "2"]
)
assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}"
mock_uvicorn_run.assert_called_once()
mock_start_reaper.assert_called_once()
def test_single_worker_uvicorn_does_not_start_reaper(self):
result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server(
["--local", "--num_workers", "1"]
)
assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}"
mock_uvicorn_run.assert_called_once()
mock_start_reaper.assert_not_called()
@pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows")
def test_gunicorn_arbiter_starts_reaper(self):
pytest.importorskip("gunicorn")
with (
patch("gunicorn.app.base.BaseApplication.run"),
patch(
"litellm.proxy.proxy_cli.start_query_engine_reaper"
) as mock_start_reaper,
):
ProxyInitializationHelpers._run_gunicorn_server(
host="127.0.0.1",
port=4010,
app=MagicMock(),
num_workers=1,
ssl_certfile_path=None,
ssl_keyfile_path=None,
)
mock_start_reaper.assert_called_once()
class TestRunServerDbSetup:
"""Tests for run_server's prisma setup_database behavior."""