From e48350ae23cb3e0d65194978ecb19b1f93f946c0 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 6 Nov 2026 16:01:41 +0800 Subject: [PATCH] fix(alerts): retain admission while cancelled database work settles --- .../litellm_core_utils/pending_operations.py | 57 ++++++++ litellm/proxy/db/prisma_client.py | 25 +++- litellm/repositories/analytics_budget.py | 35 ++--- .../test_pending_operations.py | 32 +++++ .../repositories/test_analytics_budget.py | 126 ++++++++++++++++++ 5 files changed, 254 insertions(+), 21 deletions(-) create mode 100644 litellm/litellm_core_utils/pending_operations.py create mode 100644 tests/unit/litellm_core_utils/test_pending_operations.py diff --git a/litellm/litellm_core_utils/pending_operations.py b/litellm/litellm_core_utils/pending_operations.py new file mode 100644 index 00000000000..5c8a242c849 --- /dev/null +++ b/litellm/litellm_core_utils/pending_operations.py @@ -0,0 +1,57 @@ +from collections.abc import Callable, Generator +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Final + + +class PendingOperations: + def __init__(self, completed: Callable[[], None]) -> None: + self._completed = completed + self._pending = 0 + self._closed = False + self._released = False + + def started(self) -> bool: + if self._released: + return False + self._pending += 1 + return True + + def finished(self) -> None: + self._pending -= 1 + if self._closed and self._pending == 0: + self._release() + + def close(self) -> None: + self._closed = True + if self._pending == 0: + self._release() + + def _release(self) -> None: + if not self._released: + self._released = True + self._completed() + + +_PENDING: Final[ContextVar[PendingOperations | None]] = ContextVar("litellm_pending_operations", default=None) + + +@contextmanager +def defer_until_operations_finish(completed: Callable[[], None]) -> Generator[None, None, None]: + operations: Final = PendingOperations(completed) + token: Final = _PENDING.set(operations) + try: + yield + finally: + _PENDING.reset(token) + operations.close() + + +def register_pending_operation() -> PendingOperations | None: + operations: Final = _PENDING.get() + return operations if operations is not None and operations.started() else None + + +def finish_pending_operation(operations: PendingOperations | None) -> None: + if operations is not None: + operations.finished() diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index be5de040fdf..6bfb5c7680f 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -19,6 +19,11 @@ from typing import TYPE_CHECKING, Final, Protocol, TypeVar from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.pending_operations import ( + PendingOperations, + finish_pending_operation, + register_pending_operation, +) from litellm.proxy.db.db_span import db_span from litellm.proxy.db.db_url_settings import add_missing_query_params, token_refresh_params_from_url from litellm.proxy.db.log_db_metrics import db_io_claimed, record_db_io @@ -148,33 +153,45 @@ class _TrackedPrismaEngine: ) -> _EngineResponseT: self.tracker.begin_operation() task: Final = asyncio.ensure_future(request) + pending: Final = register_pending_operation() try: - return await asyncio.shield(task) + response: Final = 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)) + task.add_done_callback(lambda done: self._detached_request_finished(done, abandoned, pending)) raise + except BaseException: + finish_pending_operation(pending) + raise + else: + finish_pending_operation(pending) + return response def _detached_request_finished( self, done: "asyncio.Task[_EngineResponseT]", abandoned: Callable[[_EngineResponseT], Awaitable[None]] | None, + pending: PendingOperations | None, ) -> None: self._detached = self._detached.difference((done,)) if done.cancelled(): + finish_pending_operation(pending) return failure: Final = done.exception() if failure is not None: verbose_proxy_logger.debug("Prisma request abandoned by a cancelled caller failed: %s", failure) + finish_pending_operation(pending) return if abandoned is None: + finish_pending_operation(pending) return settle: Final = asyncio.ensure_future(abandoned(done.result())) self._detached = self._detached.union((settle,)) - settle.add_done_callback(self._abandoned_transaction_settled) + settle.add_done_callback(lambda done: self._abandoned_transaction_settled(done, pending)) - def _abandoned_transaction_settled(self, settle: "asyncio.Task[object]") -> None: + def _abandoned_transaction_settled(self, settle: "asyncio.Task[object]", pending: PendingOperations | None) -> None: self._detached = self._detached.difference((settle,)) + finish_pending_operation(pending) if settle.cancelled() or settle.exception() is None: return verbose_proxy_logger.debug( diff --git a/litellm/repositories/analytics_budget.py b/litellm/repositories/analytics_budget.py index d5e02e655a4..11ac67d970a 100644 --- a/litellm/repositories/analytics_budget.py +++ b/litellm/repositories/analytics_budget.py @@ -9,6 +9,8 @@ from typing import Final, Protocol from prisma.errors import RawQueryError, TransactionExpiredError, TransactionNotStartedError from pydantic import BaseModel +from litellm.litellm_core_utils.pending_operations import defer_until_operations_finish + class AnalyticsTransaction(Protocol): async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ... @@ -101,23 +103,22 @@ class BoundedAnalyticsDatabase: async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: if not self._admission.acquire(blocking=False): raise AnalyticsCapacityError("Analytics capacity is busy. Try again shortly") - try: - remaining: Final = _remaining_seconds() - return await asyncio.wait_for( - self._query(query, args), timeout=6.25 if remaining is None else min(6.25, remaining) - ) - except RawQueryError as error: - if _PostgresError.model_validate(error.meta or {}).code not in ("57014", "55P03"): - raise - raise AnalyticsCapacityError("Analytics exceeded its query budget. Try again shortly") from error - except ( - TransactionExpiredError, - TransactionNotStartedError, - asyncio.TimeoutError, - ) as error: - raise AnalyticsCapacityError("Analytics exceeded its connection budget. Try again shortly") from error - finally: - self._admission.release() + with defer_until_operations_finish(self._admission.release): + try: + remaining: Final = _remaining_seconds() + return await asyncio.wait_for( + self._query(query, args), timeout=6.25 if remaining is None else min(6.25, remaining) + ) + except RawQueryError as error: + if _PostgresError.model_validate(error.meta or {}).code not in ("57014", "55P03"): + raise + raise AnalyticsCapacityError("Analytics exceeded its query budget. Try again shortly") from error + except ( + TransactionExpiredError, + TransactionNotStartedError, + asyncio.TimeoutError, + ) as error: + raise AnalyticsCapacityError("Analytics exceeded its connection budget. Try again shortly") from error async def _query(self, query: str, args: tuple[object, ...]) -> Sequence[Mapping[str, object]]: async with _bounded_transaction( diff --git a/tests/unit/litellm_core_utils/test_pending_operations.py b/tests/unit/litellm_core_utils/test_pending_operations.py new file mode 100644 index 00000000000..965334bd397 --- /dev/null +++ b/tests/unit/litellm_core_utils/test_pending_operations.py @@ -0,0 +1,32 @@ +import asyncio +from threading import BoundedSemaphore +from typing import Final + +import pytest + +from litellm.litellm_core_utils.pending_operations import ( + defer_until_operations_finish, + finish_pending_operation, + register_pending_operation, +) + + +@pytest.mark.asyncio +async def test_late_child_cannot_release_a_drained_parent_scope_twice() -> None: + admission: Final = BoundedSemaphore(1) + assert admission.acquire(blocking=False) + proceed: Final = asyncio.Event() + + async def late_child() -> None: + await proceed.wait() + operation: Final = register_pending_operation() + assert operation is None + finish_pending_operation(operation) + + with defer_until_operations_finish(admission.release): + child: Final = asyncio.create_task(late_child()) + assert admission.acquire(blocking=False) + proceed.set() + await child + assert not admission.acquire(blocking=False) + admission.release() diff --git a/tests/unit/repositories/test_analytics_budget.py b/tests/unit/repositories/test_analytics_budget.py index befa645c4b6..90f8ab83f3b 100644 --- a/tests/unit/repositories/test_analytics_budget.py +++ b/tests/unit/repositories/test_analytics_budget.py @@ -1,4 +1,5 @@ import asyncio +import heapq from collections import deque from collections.abc import AsyncGenerator, Mapping, Sequence from contextlib import asynccontextmanager, nullcontext @@ -27,6 +28,8 @@ class _ClockedLoop(asyncio.SelectorEventLoop): return self.now def _run_once(self) -> None: + while self._scheduled and self._scheduled[0].cancelled(): + heapq.heappop(self._scheduled) if not self._ready and self._scheduled: self.now = max(self.now, self._scheduled[0].when()) super()._run_once() # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue] # CPython loop hook lacks stubs @@ -99,3 +102,126 @@ def test_stalled_driver_is_cancelled_without_retaining_admission_or_tasks( loop.run_until_complete(scenario()) finally: loop.close() + + +class _Process: + pid: object = None + + +class _HeldEngine: + def __init__(self, stage: Literal["begin", "query", "commit", "rollback"], *, failure: bool = False) -> None: + self.stage = stage + self.failure = failure + self.entered = asyncio.Event() + self.released = asyncio.Event() + self.rollback_entered = asyncio.Event() + self.rollback_released = asyncio.Event() + self.process = _Process() + self.calls = 0 + self.starts = 0 + + async def _request(self, stage: str) -> None: + self.calls += 1 + if stage == self.stage: + self.entered.set() + await self.released.wait() + if self.failure: + raise RuntimeError("synthetic engine error") + + async def query(self, content: str, *, tx_id: str | None) -> object: + if content == "SELECT 1 AS result": + await self._request("query") + return {} + + async def start_transaction(self, *, content: str) -> str: + self.starts += 1 + identifier: Final = f"synthetic-transaction-{self.starts}" + await self._request("begin") + return identifier + + async def commit_transaction(self, tx_id: str) -> None: + await self._request("commit") + + async def rollback_transaction(self, tx_id: str) -> None: + await self._request("rollback") + if self.stage == "begin": + self.rollback_entered.set() + await self.rollback_released.wait() + + +class _TrackedDriver: + def __init__(self, engine: _HeldEngine) -> None: + from litellm.proxy.db.prisma_client import ( + _PrismaDrainTracker, # pyright: ignore[reportPrivateUsage] # exercise real cancellation ownership + _TrackedPrismaEngine, # pyright: ignore[reportPrivateUsage] # exercise real cancellation ownership + ) + + self.stage = engine.stage + self.tracker = _PrismaDrainTracker() + self.engine = _TrackedPrismaEngine(engine, self.tracker) + + @asynccontextmanager + async def tx(self, *, max_wait: timedelta, timeout: timedelta) -> AsyncGenerator[AnalyticsTransaction]: + identifier: Final = await self.engine.start_transaction(content="{}") + try: + yield self + if self.stage == "rollback": + raise RuntimeError("synthetic query failure") + except BaseException: + await self.engine.rollback_transaction(identifier) + raise + else: + await self.engine.commit_transaction(identifier) + + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + from litellm.proxy.db.log_db_metrics import claim_db_io + + with claim_db_io(): + await self.engine.query(query, tx_id="synthetic-transaction") + return ({"result": 1},) if query == "SELECT 1 AS result" else ({"acquired": True},) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stage", ("begin", "query", "commit", "rollback")) +@pytest.mark.parametrize("failure", (False, True)) +async def test_cancelled_analytics_keeps_admission_until_detached_engine_work_and_rollback_finish( + stage: Literal["begin", "query", "commit", "rollback"], + failure: bool, +) -> None: + engines: Final = (_HeldEngine(stage, failure=failure), _HeldEngine(stage, failure=failure)) + drivers: Final = tuple(_TrackedDriver(engine) for engine in engines) + tasks: Final = tuple( + asyncio.create_task(BoundedAnalyticsDatabase(driver).query_raw("SELECT 1 AS result")) for driver in drivers + ) + try: + await asyncio.gather(*(engine.entered.wait() for engine in engines)) + for task in tasks: + task.cancel() + for task in tasks: + with pytest.raises(asyncio.CancelledError): + await task + calls: Final = tuple(engine.calls for engine in engines) + for _ in range(64): + with pytest.raises(AnalyticsCapacityError): + await BoundedAnalyticsDatabase(drivers[0]).query_raw("SELECT 1 AS result") + assert tuple(engine.calls for engine in engines) == calls + assert tuple(len(driver.engine._detached) for driver in drivers) == (1, 1) # pyright: ignore[reportPrivateUsage] # bounded retained engine tasks are the behavior under test + for engine in engines: + engine.released.set() + if stage == "begin" and not failure: + await asyncio.gather(*(engine.rollback_entered.wait() for engine in engines)) + with pytest.raises(AnalyticsCapacityError): + await BoundedAnalyticsDatabase(drivers[0]).query_raw("SELECT 1 AS result") + finally: + for engine in engines: + engine.released.set() + engine.rollback_released.set() + await asyncio.gather(*(driver.tracker.wait_until_drained() for driver in drivers)) + await asyncio.sleep(0) + await asyncio.sleep(0) + assert not any(driver.engine._detached for driver in drivers) # pyright: ignore[reportPrivateUsage] # verify detached engine tasks drained + healthy: Final = BoundedAnalyticsDatabase(_Driver("healthy")) + assert await asyncio.gather(healthy.query_raw("SELECT 1 AS result"), healthy.query_raw("SELECT 1 AS result")) == [ + ({"result": 1},), + ({"result": 1},), + ]