diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 593f7c0503f..0cbce6339fd 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -11,9 +11,9 @@ import subprocess import time import urllib import urllib.parse -from collections.abc import Callable +from collections.abc import Awaitable, Callable, Coroutine from datetime import datetime, timedelta -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar from litellm._logging import verbose_proxy_logger from litellm.proxy.db.db_span import db_span @@ -96,10 +96,20 @@ class _PrismaDrainTracker: await self._drained.wait() +_EngineResponseT: Final = TypeVar("_EngineResponseT") + + class _TrackedPrismaEngine: + """Engine wrapper that counts in-flight work and never lets a cancelled caller drop an engine request. + + A dropped request leaves the engine holding a busy pool connection and can crash it once that statement + finishes, so every request runs to completion in its own task and only the caller is cancelled. + """ + def __init__(self, engine: _PrismaEngine, tracker: _PrismaDrainTracker) -> None: self._engine = engine self.tracker = tracker + self._detached: frozenset[asyncio.Task[object]] = frozenset() @property def process(self) -> _PrismaProcess: @@ -109,7 +119,60 @@ class _TrackedPrismaEngine: return getattr(self._engine, name) async def query(self, content: str, *, tx_id: str | None) -> object: + return await self._finish_even_if_cancelled(self._tracked_query(content, tx_id)) + + async def start_transaction(self, *, content: str) -> str: + return await self._finish_even_if_cancelled( + self._tracked_start_transaction(content), abandoned=self.rollback_transaction + ) + + async def commit_transaction(self, tx_id: str) -> None: + await self._finish_even_if_cancelled(self._tracked_commit_transaction(tx_id)) + + async def rollback_transaction(self, tx_id: str) -> None: + await self._finish_even_if_cancelled(self._tracked_rollback_transaction(tx_id)) + + async def _finish_even_if_cancelled( + self, + request: Coroutine[object, object, _EngineResponseT], + abandoned: Callable[[_EngineResponseT], Awaitable[None]] | None = None, + ) -> _EngineResponseT: self.tracker.begin_operation() + task: Final = asyncio.ensure_future(request) + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + self._detached = self._detached.union((task,)) + task.add_done_callback(lambda done: self._detached_request_finished(done, abandoned)) + raise + + def _detached_request_finished( + self, + done: "asyncio.Task[_EngineResponseT]", + abandoned: Callable[[_EngineResponseT], Awaitable[None]] | None, + ) -> None: + self._detached = self._detached.difference((done,)) + if done.cancelled(): + return + failure: Final = done.exception() + if failure is not None: + verbose_proxy_logger.debug("Prisma request abandoned by a cancelled caller failed: %s", failure) + return + if abandoned is None: + return + settle: Final = asyncio.ensure_future(abandoned(done.result())) + self._detached = self._detached.union((settle,)) + settle.add_done_callback(self._abandoned_transaction_settled) + + def _abandoned_transaction_settled(self, settle: "asyncio.Task[object]") -> None: + self._detached = self._detached.difference((settle,)) + if settle.cancelled() or settle.exception() is None: + return + verbose_proxy_logger.debug( + "Rolling back a transaction abandoned by a cancelled caller failed: %s", settle.exception() + ) + + async def _tracked_query(self, content: str, tx_id: str | None) -> object: try: if db_io_claimed(): record_db_io() @@ -121,9 +184,8 @@ class _TrackedPrismaEngine: finally: self.tracker.end_operation() - async def start_transaction(self, *, content: str) -> str: + async def _tracked_start_transaction(self, content: str) -> str: record_db_io() - self.tracker.begin_operation() try: transaction_id: Final = await self._engine.start_transaction(content=content) except (Exception, asyncio.CancelledError): @@ -132,18 +194,16 @@ class _TrackedPrismaEngine: self.tracker.transaction_started(transaction_id) return transaction_id - async def commit_transaction(self, tx_id: str) -> None: + async def _tracked_commit_transaction(self, tx_id: str) -> None: record_db_io() - self.tracker.begin_operation() try: await self._engine.commit_transaction(tx_id) finally: self.tracker.end_operation() self.tracker.transaction_finished(tx_id) - async def rollback_transaction(self, tx_id: str) -> None: + async def _tracked_rollback_transaction(self, tx_id: str) -> None: record_db_io() - self.tracker.begin_operation() try: await self._engine.rollback_transaction(tx_id) finally: diff --git a/tests/unit/proxy/db/test_prisma_client.py b/tests/unit/proxy/db/test_prisma_client.py index 7b8d000a8d5..99b500ec315 100644 --- a/tests/unit/proxy/db/test_prisma_client.py +++ b/tests/unit/proxy/db/test_prisma_client.py @@ -1,8 +1,10 @@ +import asyncio import json import os import signal import sys import urllib.parse +from typing import Any, Final, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest @@ -10,7 +12,13 @@ from fastapi.testclient import TestClient -from litellm.proxy.db.prisma_client import PrismaManager, PrismaWrapper, should_update_prisma_schema +from litellm.proxy.db.prisma_client import ( + PrismaManager, + PrismaWrapper, + _PrismaDrainTracker, + _TrackedPrismaEngine, + should_update_prisma_schema, +) @pytest.fixture(autouse=True) @@ -470,3 +478,97 @@ def test_without_proxy_extras_the_index_build_reports_failure_instead_of_raising monkeypatch.setitem(sys.modules, "litellm_proxy_extras.utils", None) assert run() is outcome + + +class _SlowEngine: + """Engine double whose requests only answer once the test lets them, recording what the engine saw.""" + + def __init__(self) -> None: + self.gates: dict[str, asyncio.Event] = {} # mutable-ok: one gate per request the test lets through + self.queries_finished = 0 + self.requests_cancelled = 0 + self.rolled_back: list[str] = [] # mutable-ok: test records + self.process = object() + + def gate(self, request: str) -> asyncio.Event: + return self.gates.setdefault(request, asyncio.Event()) + + async def _answer_once_released(self, request: str) -> None: + try: + await self.gate(request).wait() + except asyncio.CancelledError: + self.requests_cancelled += 1 + raise + + async def query(self, content: str, *, tx_id: str | None) -> object: + await self._answer_once_released(content) + self.queries_finished += 1 + return {"content": content, "tx_id": tx_id} + + async def start_transaction(self, *, content: str) -> str: + await self._answer_once_released("start") + return "tx-opened-after-caller-left" + + async def commit_transaction(self, tx_id: str) -> None: + raise AssertionError(f"unexpected commit of {tx_id}") + + async def rollback_transaction(self, tx_id: str) -> None: + self.rolled_back.append(tx_id) + + +def _tracked(engine: _SlowEngine) -> tuple[_TrackedPrismaEngine, _PrismaDrainTracker]: + tracker: Final = _PrismaDrainTracker() + return _TrackedPrismaEngine(cast(Any, engine), tracker), tracker + + +@pytest.mark.asyncio +async def test_cancelling_a_query_caller_lets_the_engine_finish_the_statement() -> None: + engine: Final = _SlowEngine() + tracked, tracker = _tracked(engine) + caller: Final = asyncio.create_task(tracked.query("x", tx_id=None)) + await asyncio.sleep(0) + caller.cancel() + with pytest.raises(asyncio.CancelledError): + await caller + drained: Final = asyncio.create_task(tracker.wait_until_drained()) + await asyncio.sleep(0) + assert engine.requests_cancelled == 0 + assert not drained.done() + engine.gate("x").set() + await asyncio.wait_for(drained, timeout=1) + assert engine.queries_finished == 1 + + +@pytest.mark.asyncio +async def test_transaction_opened_after_its_caller_was_cancelled_is_rolled_back() -> None: + engine: Final = _SlowEngine() + tracked, tracker = _tracked(engine) + caller: Final = asyncio.create_task(tracked.start_transaction(content="{}")) + await asyncio.sleep(0) + caller.cancel() + with pytest.raises(asyncio.CancelledError): + await caller + engine.gate("start").set() + await asyncio.wait_for(tracker.wait_until_drained(), timeout=1) + assert engine.requests_cancelled == 0 + assert engine.rolled_back == ["tx-opened-after-caller-left"] + + +@pytest.mark.asyncio +async def test_a_query_issued_while_the_last_one_finishes_keeps_the_engine_drain_open() -> None: + engine: Final = _SlowEngine() + tracked, tracker = _tracked(engine) + first: Final = asyncio.create_task(tracked.query("first", tx_id=None)) + await asyncio.sleep(0) + await asyncio.sleep(0) + drained: Final = asyncio.create_task(tracker.wait_until_drained()) + await asyncio.sleep(0) + second: Final = asyncio.create_task(tracked.query("second", tx_id=None)) + engine.gate("first").set() + await first + await asyncio.sleep(0) + assert not drained.done() + engine.gate("second").set() + await second + await asyncio.wait_for(drained, timeout=1) + assert engine.queries_finished == 2