From cee0094d58d844cbaf1444b6934e7d911388fb0d Mon Sep 17 00:00:00 2001 From: moe-berri Date: Thu, 5 Nov 2026 16:20:01 +0800 Subject: [PATCH] fix: clamp analytics queries to the remaining work deadline --- litellm/repositories/analytics_budget.py | 58 +++++++++++++++---- .../database/test_analytics_budget.py | 23 ++++++++ 2 files changed, 69 insertions(+), 12 deletions(-) diff --git a/litellm/repositories/analytics_budget.py b/litellm/repositories/analytics_budget.py index 67524faa9d1..a4d7cc42065 100644 --- a/litellm/repositories/analytics_budget.py +++ b/litellm/repositories/analytics_budget.py @@ -1,5 +1,7 @@ -from collections.abc import Mapping, Sequence -from contextlib import AbstractAsyncContextManager +import asyncio +from collections.abc import Iterator, Mapping, Sequence +from contextlib import AbstractAsyncContextManager, contextmanager +from contextvars import ContextVar from datetime import timedelta from threading import BoundedSemaphore from typing import Final, Protocol @@ -30,12 +32,13 @@ class _PostgresError(BaseModel, frozen=True): _FOREGROUND: Final = BoundedSemaphore(2) _BACKGROUND: Final = BoundedSemaphore(1) +_DEADLINE: Final = ContextVar[float | None]("analytics_deadline", default=None) _CONFIGURE: Final = """ SELECT set_config('statement_timeout', LEAST(NULLIF( - EXTRACT(epoch FROM current_setting('statement_timeout')::interval)*1000,0),5000)::text,true), + EXTRACT(epoch FROM current_setting('statement_timeout')::interval)*1000,0),$2::int)::text,true), set_config('lock_timeout', LEAST(NULLIF( - EXTRACT(epoch FROM current_setting('lock_timeout')::interval)*1000,0),100)::text,true), + EXTRACT(epoch FROM current_setting('lock_timeout')::interval)*1000,0),100,$2::int)::text,true), set_config('work_mem', LEAST(pg_size_bytes(current_setting('work_mem')),1048576)::text || 'B',true), set_config('max_parallel_workers_per_gather','0',true), CASE WHEN $1::boolean THEN pg_try_advisory_xact_lock(1876539401,2) @@ -44,6 +47,31 @@ SELECT """ +@contextmanager +def analytics_deadline(deadline: float) -> Iterator[None]: + current: Final = _DEADLINE.get() + token: Final = _DEADLINE.set(deadline if current is None else min(current, deadline)) + try: + yield + finally: + _DEADLINE.reset(token) + + +def _remaining_seconds() -> float | None: + deadline: Final = _DEADLINE.get() + if deadline is None: + return None + remaining: Final = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + raise AnalyticsCapacityError("Analytics work budget is exhausted") + return remaining + + +def _milliseconds(limit: int) -> int: + remaining: Final = _remaining_seconds() + return limit if remaining is None else min(limit, max(1, int(remaining * 1000))) + + class BoundedAnalyticsDatabase: def __init__(self, database: AnalyticsDatabase, *, background: bool = False) -> None: self._database = database @@ -54,18 +82,24 @@ class BoundedAnalyticsDatabase: if not self._admission.acquire(blocking=False): raise AnalyticsCapacityError("Analytics capacity is busy. Try again shortly") try: - async with self._database.tx( - max_wait=timedelta(milliseconds=250), timeout=timedelta(seconds=6) - ) as transaction: - configuration: Final = await transaction.query_raw(_CONFIGURE, self._background) - if not configuration or not _Admission.model_validate(configuration[0]).acquired: - raise AnalyticsCapacityError("Analytics capacity is busy. Try again shortly") - return await transaction.query_raw(query, *args) + remaining: Final = _remaining_seconds() + return await asyncio.wait_for(self._query(query, args), timeout=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) as error: + except (TransactionExpiredError, TransactionNotStartedError, asyncio.TimeoutError) as error: raise AnalyticsCapacityError("Analytics exceeded its connection budget. Try again shortly") from error finally: 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)), + ) as transaction: + configuration: Final = await transaction.query_raw(_CONFIGURE, self._background, _milliseconds(5000)) + if not configuration or not _Admission.model_validate(configuration[0]).acquired: + raise AnalyticsCapacityError("Analytics capacity is busy. Try again shortly") + _remaining_seconds() + return await transaction.query_raw(query, *args) diff --git a/tests/integration/database/test_analytics_budget.py b/tests/integration/database/test_analytics_budget.py index 1f52df9cf43..9f2a9b5c17a 100644 --- a/tests/integration/database/test_analytics_budget.py +++ b/tests/integration/database/test_analytics_budget.py @@ -1,3 +1,4 @@ +import asyncio from datetime import timedelta from typing import Final, cast @@ -8,6 +9,7 @@ from litellm.repositories.analytics_budget import ( AnalyticsCapacityError, AnalyticsDatabase, BoundedAnalyticsDatabase, + analytics_deadline, ) _SETTINGS: Final = """SELECT current_setting('statement_timeout') AS statement_timeout, @@ -57,3 +59,24 @@ async def test_expensive_analytics_are_cancelled_and_release_capacity(db: Manage assert active[0]["running"] == 0 assert (await bounded.query_raw("SELECT 1 AS result"))[0]["result"] == 1 assert await db.prisma.writer_db.query_raw(_SETTINGS) == before + + +@pytest.mark.asyncio(loop_scope="session") +async def test_work_deadline_cancels_inflight_sql_and_restores_capacity(db: ManagementDatabase) -> None: + database: Final = cast(AnalyticsDatabase, db.prisma.writer_db) + bounded: Final = BoundedAnalyticsDatabase(database, background=True) + loop: Final = asyncio.get_running_loop() + start: Final = loop.time() + with analytics_deadline(start + 0.25), analytics_deadline(start + 10): + with pytest.raises(AnalyticsCapacityError): + await bounded.query_raw("SELECT 1 FROM pg_sleep(30) /* bounded_analytics_deadline */") + assert loop.time() - start < 2 + active: Final = await db.prisma.writer_db.query_raw( + """SELECT count(*)::int AS running FROM pg_stat_activity + WHERE pid<>pg_backend_pid() AND state='active' AND query LIKE $1""", + "%bounded_analytics_deadline%", + ) + assert active[0]["running"] == 0 + with analytics_deadline(loop.time() - 1), pytest.raises(AnalyticsCapacityError): + await bounded.query_raw("SELECT 1 AS result") + assert (await bounded.query_raw("SELECT 1 AS result"))[0]["result"] == 1