fix(proxy): kill mid-drain retired prisma engines at shutdown instead of abandoning them

Follow-up to #34749. Engine retirement tasks had no shutdown owner: a
task still waiting for its drain deadline when the proxy shut down was
cancelled at event-loop teardown before its kill ran, orphaning the
replaced query-engine subprocess and the DB connections it holds. In
containers the pod teardown reaps the orphan, but bare-metal and dev
deployments leaked it until Postgres idle timeouts fired.

stop_token_refresh_task, which the proxy shutdown hook already calls
on every wrapper (the routing wrapper forwards it to writer and
reader), now flushes pending retirements after stopping the refresh
loop: cancel each retirement task, await it, then SIGKILL the engine
pid directly. The direct kill also covers a retirement task cancelled
before it ever ran, which would otherwise skip its own kill entirely,
so retirements are tracked as (pid, task) pairs. A retirement task
cancelled anywhere else likewise SIGKILLs its engine before
propagating the cancellation

litellm/proxy/db/prisma_client.py

@@ -64,6 +64,12 @@ class _PrismaClient(Protocol):
     def _engine(self) -> _PrismaEngine: ...
 
 
+@dataclass(frozen=True, slots=True)
+class _EngineRetirement:
+    pid: int
+    task: "asyncio.Task[None]"
+
+
 class _PrismaDrainTracker:
     def __init__(self) -> None:
         self._active_operations = 0
@@ -222,7 +228,7 @@ class PrismaWrapper:
         self._reconnection_lock = asyncio.Lock()
         self._last_refresh_time: datetime | None = None
         self._active_drain_tracker = self._instrument_prisma_client(original_prisma)
-        self._retirement_tasks: frozenset[asyncio.Task[None]] = frozenset()
+        self._retirement_tasks: frozenset[_EngineRetirement] = frozenset()
 
         # Coordination for planned engine restarts (issue #29176). Every
         # `recreate_prisma_client` SIGTERMs the running query-engine on
@@ -282,30 +288,59 @@ class PrismaWrapper:
         return 0
 
     async def _retire_engine_when_drained(self, pid: int, tracker: _PrismaDrainTracker | None) -> None:
-        if tracker is not None:
-            try:
-                await asyncio.wait_for(
-                    tracker.wait_until_drained(),
-                    timeout=self.ENGINE_RETIREMENT_DRAIN_TIMEOUT_SECONDS,
-                )
-            except asyncio.TimeoutError:
-                verbose_proxy_logger.warning(
-                    "%sReplaced prisma engine PID %s did not drain within %ss; killing it with work still in flight.",
-                    self._log_prefix,
-                    pid,
-                    self.ENGINE_RETIREMENT_DRAIN_TIMEOUT_SECONDS,
-                )
-        await self._kill_engine_process(pid)
+        try:
+            if tracker is not None:
+                try:
+                    await asyncio.wait_for(
+                        tracker.wait_until_drained(),
+                        timeout=self.ENGINE_RETIREMENT_DRAIN_TIMEOUT_SECONDS,
+                    )
+                except asyncio.TimeoutError:
+                    verbose_proxy_logger.warning(
+                        "%sReplaced prisma engine PID %s did not drain within %ss; killing it with work still in flight.",
+                        self._log_prefix,
+                        pid,
+                        self.ENGINE_RETIREMENT_DRAIN_TIMEOUT_SECONDS,
+                    )
+            await self._kill_engine_process(pid)
+        except asyncio.CancelledError:
+            self._kill_engine_process_now(pid)
+            raise
+
+    def _kill_engine_process_now(self, pid: int) -> None:
+        if pid <= 0:
+            return
+        try:
+            os.kill(pid, getattr(signal, "SIGKILL", signal.SIGTERM))
+        except (ProcessLookupError, PermissionError, OSError):
+            return
+        verbose_proxy_logger.warning(
+            "%sSent SIGKILL to prisma-query-engine PID %s while flushing engine retirements.",
+            self._log_prefix,
+            pid,
+        )
+
+    async def flush_engine_retirements(self) -> None:
+        pending = self._retirement_tasks
+        if not pending:
+            return
+        for retirement in pending:
+            retirement.task.cancel()
+        await asyncio.gather(*(retirement.task for retirement in pending), return_exceptions=True)
+        for retirement in pending:
+            self._kill_engine_process_now(retirement.pid)
 
     def _schedule_engine_retirement(self, pid: int, tracker: _PrismaDrainTracker | None) -> None:
         if pid <= 0:
             return
         retirement_task = asyncio.create_task(self._retire_engine_when_drained(pid, tracker))
-        self._retirement_tasks = self._retirement_tasks.union((retirement_task,))
+        self._retirement_tasks = self._retirement_tasks.union((_EngineRetirement(pid=pid, task=retirement_task),))
         retirement_task.add_done_callback(self._retirement_finished)
 
     def _retirement_finished(self, retirement_task: asyncio.Task[None]) -> None:
-        self._retirement_tasks = self._retirement_tasks.difference((retirement_task,))
+        self._retirement_tasks = frozenset(
+            retirement for retirement in self._retirement_tasks if retirement.task is not retirement_task
+        )
 
     async def connect(self, timeout: int | timedelta | None = None) -> None:
         if timeout is None:
@@ -649,16 +684,16 @@ class PrismaWrapper:
 
         Should be called during application shutdown to clean up resources.
         """
-        if self._token_refresh_task is None:
-            return
+        if self._token_refresh_task is not None:
+            self._token_refresh_task.cancel()
+            try:
+                await self._token_refresh_task
+            except asyncio.CancelledError:
+                pass
+            self._token_refresh_task = None
+            verbose_proxy_logger.info("%sStopped RDS IAM token refresh background task", self._log_prefix)
 
-        self._token_refresh_task.cancel()
-        try:
-            await self._token_refresh_task
-        except asyncio.CancelledError:
-            pass
-        self._token_refresh_task = None
-        verbose_proxy_logger.info("%sStopped RDS IAM token refresh background task", self._log_prefix)
+        await self.flush_engine_retirements()
 
     async def _token_refresh_loop(self) -> None:
         """

tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py

@@ -71,7 +71,7 @@ def test_wrapper_instruments_generated_prisma_engine() -> None:
 
 
 async def _wait_for_retirements(wrapper: PrismaWrapper) -> None:
-    await asyncio.gather(*tuple(wrapper._retirement_tasks))
+    await asyncio.gather(*tuple(retirement.task for retirement in wrapper._retirement_tasks))
 
 
 @pytest.mark.asyncio
@@ -425,6 +425,24 @@ async def test_retirement_kills_old_engine_when_drain_never_completes(
     mock_kill.assert_any_call(111, signal.SIGTERM)
 
 
+@pytest.mark.asyncio
+async def test_stop_token_refresh_flushes_pending_engine_retirement(mock_prisma_binary):
+    """Shutdown must kill a mid-drain replaced engine immediately instead of
+    abandoning its retirement task at event-loop teardown (which would orphan
+    the engine and its DB connection pool)."""
+    wrapper = _make_wrapper(engine_pid=111, iam=True)
+    old_engine = wrapper._original_prisma._engine
+    old_engine._engine.start_transaction = AsyncMock(return_value="transaction-1")
+    await old_engine.start_transaction(content="transaction")
+    wrapper._schedule_engine_retirement(111, wrapper._active_drain_tracker)
+
+    with patch("os.kill") as mock_kill:
+        await asyncio.wait_for(wrapper.stop_token_refresh_task(), timeout=2)
+
+    assert mock_kill.call_args_list == [call(111, signal.SIGKILL)]
+    assert wrapper._retirement_tasks == frozenset()
+
+
 @pytest.mark.asyncio
 async def test_safe_refresh_cancellation_restores_token_and_cleans_replacement(
     mock_prisma_binary, monkeypatch