fix: clamp analytics queries to the remaining work deadline

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)

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