Merge commit 'e48350ae23cb3e0d65194978ecb19b1f93f946c0' into litellm_alerts_pr2_final_validation

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()

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(

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(

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()

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},),
+    ]