fix(proxy): kill mid-drain retired prisma engines at shutdown instead of abandoning them - #34999
Conversation
… 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
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Greptile SummaryThis PR closes the shutdown race that could leave retired Prisma query engines alive.
Confidence Score: 5/5The PR appears safe to merge. No blocking failure remains; the flush flag is set before the synchronous retirement snapshot, and every later retirement schedule observes the flag and kills its engine inline.
|
| Filename | Overview |
|---|---|
| litellm/proxy/db/prisma_client.py | Adds shutdown ownership for pending engine retirements and synchronously handles retirements scheduled after flushing begins. |
| tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py | Updates retirement-task helpers for PID/task records and verifies both pending-retirement flushing and post-flush synchronous termination. |
Reviews (3): Last reviewed commit: "fix(proxy): kill retirements scheduled a..." | Re-trigger Greptile
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
…ronously An in-flight refresh triggered by the __getattr__ stale-token fallback can complete its rotation after flush_engine_retirements has taken its snapshot, scheduling a retirement task that no longer has an owner and gets abandoned at event-loop teardown. Once the flush has run, _schedule_engine_retirement now SIGKILLs the old engine inline instead of creating a task. The event loop is single threaded and both the flag-set-plus-snapshot and the schedule are synchronous, so every schedule either lands in the flush snapshot or sees the flag; no retirement can escape both
|
Pushed a637b32 addressing the P1: once the shutdown flush has run, _schedule_engine_retirement kills the old engine inline instead of creating a task, so a rotation completing after the flush snapshot (an in-flight getattr refresh) cannot leave an unowned retirement. Single threaded event loop plus no awaits between flag set, snapshot, and schedule means every retirement either lands in the snapshot or sees the flag. Regression test added that fails without the flag @greptileai re review |
|
@greptileai re review |
1 participant
TLDR
Problem this solves:
How it solves it:
Relevant issues
Follow-up to #34749
Linear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
The bug only manifests when the proxy shuts down while a warm IAM rotation's drain window (up to 90s) is open, which needs an IAM-enabled RDS instance to reach; there is no curl-visible surface because the leak is a child process outliving the proxy. Manual verification recipe for an IAM environment: run the proxy against IAM RDS, wait for a rotation log line ("RDS IAM token refreshed successfully"), immediately SIGTERM the proxy inside the drain window, then
ps aux | grep query-engine. At7cd009caf7(before) the old engine pid survives the proxy exit until Postgres idle timeouts reap its connections; at1219fde8a1(after) no query-engine process remainsDeterministic reproduction is the added regression test: at
7cd009caf7it fails becausestop_token_refresh_taskreturns without touching the pending retirement (no kill signal is ever sent), and at1219fde8a1it passes. The full database suites (615 tests acrosstests/test_litellm/proxy/db/,tests/test_litellm/proxy/utils/prisma_and_spend/, and the engine watchdog suite) pass at1219fde8a1Type
🐛 Bug Fix
Changes
PR #34749 retires a replaced engine via a background task that waits for in-flight work to drain (bounded at 90s) before killing the old engine. Those retirement tasks had no shutdown owner: a task still mid-drain when the proxy shut down was cancelled at event-loop teardown before its kill ever ran, orphaning the replaced query-engine subprocess and the DB connections it holds. Containers reap the orphan with the pod; 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 so a mid-rotation cancellation can still schedule its cleanup before the flush runs. The flush cancels each retirement task, awaits it, and then SIGKILLs the engine pid directly. The direct kill matters because a task cancelled before its first run never executes its body, so relying on the task's own cancellation handler would silently skip the kill; retirements are therefore tracked as (pid, task) pairs. Independently, a retirement task cancelled from anywhere now SIGKILLs its engine before propagating the cancellation, covering the window between the drain deadline's SIGTERM and its SIGKILL backstopA rotation completing after the flush has taken its snapshot (an in-flight refresh spawned by the getattr stale-token fallback) cannot escape either: once the flush has run, _schedule_engine_retirement kills the old engine inline instead of creating a task. The event loop is single threaded and both the flag-set-plus-snapshot and the schedule are synchronous, so every schedule either lands in the flush snapshot or sees the flag
Final Attestation