mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(proxy): kill retirements scheduled after the shutdown flush synchronously
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
This commit is contained in:
parent
1219fde8a1
commit
a637b32528
2 changed files with 20 additions and 0 deletions
|
|
@ -229,6 +229,7 @@ class PrismaWrapper:
|
|||
self._last_refresh_time: datetime | None = None
|
||||
self._active_drain_tracker = self._instrument_prisma_client(original_prisma)
|
||||
self._retirement_tasks: frozenset[_EngineRetirement] = frozenset()
|
||||
self._retirements_flushed = False
|
||||
|
||||
# Coordination for planned engine restarts (issue #29176). Every
|
||||
# `recreate_prisma_client` SIGTERMs the running query-engine on
|
||||
|
|
@ -321,6 +322,7 @@ class PrismaWrapper:
|
|||
)
|
||||
|
||||
async def flush_engine_retirements(self) -> None:
|
||||
self._retirements_flushed = True
|
||||
pending = self._retirement_tasks
|
||||
if not pending:
|
||||
return
|
||||
|
|
@ -333,6 +335,9 @@ class PrismaWrapper:
|
|||
def _schedule_engine_retirement(self, pid: int, tracker: _PrismaDrainTracker | None) -> None:
|
||||
if pid <= 0:
|
||||
return
|
||||
if self._retirements_flushed:
|
||||
self._kill_engine_process_now(pid)
|
||||
return
|
||||
retirement_task = asyncio.create_task(self._retire_engine_when_drained(pid, tracker))
|
||||
self._retirement_tasks = self._retirement_tasks.union((_EngineRetirement(pid=pid, task=retirement_task),))
|
||||
retirement_task.add_done_callback(self._retirement_finished)
|
||||
|
|
|
|||
|
|
@ -443,6 +443,21 @@ async def test_stop_token_refresh_flushes_pending_engine_retirement(mock_prisma_
|
|||
assert wrapper._retirement_tasks == frozenset()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retirement_scheduled_after_flush_kills_engine_synchronously(mock_prisma_binary):
|
||||
"""A rotation that completes during shutdown (e.g. an in-flight __getattr__
|
||||
refresh) must not create an unowned retirement task after the flush; the
|
||||
old engine is killed inline instead."""
|
||||
wrapper = _make_wrapper(engine_pid=111, iam=True)
|
||||
await wrapper.stop_token_refresh_task()
|
||||
|
||||
with patch("os.kill") as mock_kill:
|
||||
wrapper._schedule_engine_retirement(333, wrapper._active_drain_tracker)
|
||||
|
||||
assert mock_kill.call_args_list == [call(333, 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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue