From b53b5473585e3217fc361eb76b5c70162ac0236b Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 6 Nov 2026 15:50:02 +0800 Subject: [PATCH] fix(alerts): bound analytics client waits and transaction cleanup --- litellm/repositories/analytics_budget.py | 33 ++++-- .../repositories/test_analytics_budget.py | 101 ++++++++++++++++++ 2 files changed, 128 insertions(+), 6 deletions(-) create mode 100644 tests/unit/repositories/test_analytics_budget.py diff --git a/litellm/repositories/analytics_budget.py b/litellm/repositories/analytics_budget.py index 042bff3fe42..d5e02e655a4 100644 --- a/litellm/repositories/analytics_budget.py +++ b/litellm/repositories/analytics_budget.py @@ -1,6 +1,6 @@ import asyncio -from collections.abc import Generator, Mapping, Sequence -from contextlib import AbstractAsyncContextManager, contextmanager +from collections.abc import AsyncGenerator, Generator, Mapping, Sequence +from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager from contextvars import ContextVar from datetime import timedelta from threading import BoundedSemaphore @@ -75,6 +75,23 @@ def _milliseconds(limit: int) -> int: return limit if remaining is None else min(limit, max(1, int(remaining * 1000))) +@asynccontextmanager +async def _bounded_transaction( + manager: AbstractAsyncContextManager[AnalyticsTransaction], +) -> AsyncGenerator[AnalyticsTransaction, None]: + transaction: Final = await manager.__aenter__() + try: + yield transaction + except BaseException as error: + suppressed: Final = await asyncio.wait_for( + manager.__aexit__(type(error), error, error.__traceback__), timeout=1 + ) + if not suppressed: + raise + else: + await asyncio.wait_for(manager.__aexit__(None, None, None), timeout=1) + + class BoundedAnalyticsDatabase: def __init__(self, database: AnalyticsDatabase, *, background: bool = False) -> None: self._database = database @@ -86,7 +103,9 @@ class BoundedAnalyticsDatabase: raise AnalyticsCapacityError("Analytics capacity is busy. Try again shortly") try: remaining: Final = _remaining_seconds() - return await asyncio.wait_for(self._query(query, args), timeout=remaining) + 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 @@ -101,9 +120,11 @@ class BoundedAnalyticsDatabase: self._admission.release() async def _query(self, query: str, args: tuple[object, ...]) -> Sequence[Mapping[str, object]]: - async with self._database.tx( - max_wait=timedelta(milliseconds=_milliseconds(250)), - timeout=timedelta(milliseconds=_milliseconds(6000)), + async with _bounded_transaction( + self._database.tx( + max_wait=timedelta(milliseconds=_milliseconds(250)), + timeout=timedelta(milliseconds=_milliseconds(6000)), + ) ) as transaction: configuration: Final = await transaction.query_raw(_CONFIGURE, self._background, _milliseconds(5000)) if not configuration or not _Admission.model_validate(configuration[0]).acquired: diff --git a/tests/unit/repositories/test_analytics_budget.py b/tests/unit/repositories/test_analytics_budget.py new file mode 100644 index 00000000000..befa645c4b6 --- /dev/null +++ b/tests/unit/repositories/test_analytics_budget.py @@ -0,0 +1,101 @@ +import asyncio +from collections import deque +from collections.abc import AsyncGenerator, Mapping, Sequence +from contextlib import asynccontextmanager, nullcontext +from datetime import timedelta +from typing import Final, Literal + +import pytest + +from litellm.repositories.analytics_budget import ( + AnalyticsCapacityError, + AnalyticsTransaction, + BoundedAnalyticsDatabase, + analytics_deadline, +) + + +class _ClockedLoop(asyncio.SelectorEventLoop): + _ready: deque[asyncio.Handle] + _scheduled: list[asyncio.TimerHandle] + + def __init__(self) -> None: + super().__init__() + self.now = 10000.0 + + def time(self) -> float: + return self.now + + def _run_once(self) -> None: + 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 + + +class _Driver: + def __init__(self, stage: Literal["begin", "query", "commit", "rollback", "healthy"]) -> None: + self.stage = stage + self.cancelled = asyncio.Event() + + async def _stall(self) -> None: + try: + await asyncio.Event().wait() + finally: + self.cancelled.set() + + @asynccontextmanager + async def tx(self, *, max_wait: timedelta, timeout: timedelta) -> AsyncGenerator[AnalyticsTransaction]: + if self.stage == "begin": + await self._stall() + try: + yield self + finally: + if self.stage in ("commit", "rollback"): + await self._stall() + + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + if query == "SELECT 1 AS result": + if self.stage in ("query", "rollback"): + await self._stall() + return ({"result": 1},) + return ({"acquired": True},) + + +@pytest.mark.parametrize("stage", ("begin", "query", "commit", "rollback")) +@pytest.mark.parametrize("deadline_seconds", (None, 0.25, 10.0)) +def test_stalled_driver_is_cancelled_without_retaining_admission_or_tasks( + stage: Literal["begin", "query", "commit", "rollback"], deadline_seconds: float | None +) -> None: + async def scenario() -> None: + loop: Final = asyncio.get_running_loop() + existing: Final = asyncio.all_tasks() + driver: Final = _Driver(stage) + start: Final = loop.time() + context: Final = nullcontext() if deadline_seconds is None else analytics_deadline(start + deadline_seconds) + with context: + task: Final = asyncio.create_task(BoundedAnalyticsDatabase(driver).query_raw("SELECT 1 AS result")) + done, _ = await asyncio.wait((task,), timeout=10) + if not done: + task.cancel() + watchdog: Final = loop.call_later(1, task.cancel) + try: + with pytest.raises(asyncio.CancelledError): + await task + finally: + watchdog.cancel() + assert done, "Analytics retained its task after the query and cleanup budget" + with pytest.raises(AnalyticsCapacityError): + await task + assert driver.cancelled.is_set() + assert loop.time() - start <= (7.25 if deadline_seconds != 0.25 else 1.25) + 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},)] + assert asyncio.all_tasks() == existing + + loop: Final = _ClockedLoop() + try: + loop.run_until_complete(scenario()) + finally: + loop.close()