fix: pace background analytics before acquiring database capacity

litellm/repositories/analytics_budget.py

@@ -102,13 +102,17 @@ class BoundedAnalyticsDatabase:
         self._admission = _BACKGROUND if background else _FOREGROUND
 
     async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]:
+        if self._background:
+            remaining: Final = _remaining_seconds()
+            await asyncio.sleep(0.1 if remaining is None else min(0.1, remaining))
+            _remaining_seconds()
         if not self._admission.acquire(blocking=False):
             raise AnalyticsCapacityError("Analytics capacity is busy. Try again shortly")
         with defer_until_operations_finish(self._admission.release):
             try:
-                remaining: Final = _remaining_seconds()
+                available: Final = _remaining_seconds()
                 return await asyncio.wait_for(
-                    self._query(query, args), timeout=6.25 if remaining is None else min(6.25, remaining)
+                    self._query(query, args), timeout=6.25 if available is None else min(6.25, available)
                 )
             except EngineRequestError as error:
                 if error.response.status not in (503, 504):

tests/unit/repositories/test_analytics_budget.py

@@ -1,5 +1,6 @@
 import asyncio
 import heapq
+import math
 from collections import deque
 from collections.abc import AsyncGenerator, Mapping, Sequence
 from contextlib import asynccontextmanager, nullcontext
@@ -45,6 +46,7 @@ class _Driver:
         self.stage = stage
         self.error = error
         self.cancelled = asyncio.Event()
+        self.starts: tuple[float, ...] = ()
 
     async def _stall(self) -> None:
         if self.error is not None:
@@ -56,6 +58,7 @@ class _Driver:
 
     @asynccontextmanager
     async def tx(self, *, max_wait: timedelta, timeout: timedelta) -> AsyncGenerator[AnalyticsTransaction, None]:
+        self.starts += (asyncio.get_running_loop().time(),)
         if self.stage == "begin":
             await self._stall()
         try:
@@ -72,6 +75,71 @@ class _Driver:
         return ({"acquired": True},)
 
 
+def test_background_transactions_are_paced_without_delaying_foreground_work() -> None:
+    async def scenario() -> None:
+        loop: Final = asyncio.get_running_loop()
+        driver: Final = _Driver("healthy")
+        background: Final = BoundedAnalyticsDatabase(driver, background=True)
+        start: Final = loop.time()
+        for _ in range(5):
+            assert await background.query_raw("SELECT 1 AS result") == ({"result": 1},)
+        assert math.isclose(loop.time() - start, 0.5)
+        assert all(math.isclose(second - first, 0.1) for first, second in zip(driver.starts, driver.starts[1:]))
+        foreground_start: Final = loop.time()
+        assert await BoundedAnalyticsDatabase(driver).query_raw("SELECT 1 AS result") == ({"result": 1},)
+        assert loop.time() == foreground_start
+
+    loop: Final = _ClockedLoop()
+    try:
+        loop.run_until_complete(scenario())
+    finally:
+        loop.close()
+
+
+@pytest.mark.parametrize("budget", (0.0, 0.05))
+def test_background_pacing_does_not_start_a_transaction_past_its_deadline(budget: float) -> None:
+    async def scenario() -> None:
+        loop: Final = asyncio.get_running_loop()
+        existing: Final = asyncio.all_tasks()
+        driver: Final = _Driver("healthy")
+        database: Final = BoundedAnalyticsDatabase(driver, background=True)
+        start: Final = loop.time()
+        with analytics_deadline(start + budget), pytest.raises(AnalyticsCapacityError):
+            await database.query_raw("SELECT 1 AS result")
+        assert not driver.starts
+        assert math.isclose(loop.time() - start, budget)
+        assert await database.query_raw("SELECT 1 AS result") == ({"result": 1},)
+        assert asyncio.all_tasks() == existing
+
+    loop: Final = _ClockedLoop()
+    try:
+        loop.run_until_complete(scenario())
+    finally:
+        loop.close()
+
+
+def test_cancelling_background_pacing_retains_no_work_or_admission() -> None:
+    async def scenario() -> None:
+        existing: Final = asyncio.all_tasks()
+        driver: Final = _Driver("healthy")
+        database: Final = BoundedAnalyticsDatabase(driver, background=True)
+        task: Final = asyncio.create_task(database.query_raw("SELECT 1 AS result"))
+        await asyncio.sleep(0)
+        assert not driver.starts
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+        assert await database.query_raw("SELECT 1 AS result") == ({"result": 1},)
+        assert len(driver.starts) == 1
+        assert asyncio.all_tasks() == existing
+
+    loop: Final = _ClockedLoop()
+    try:
+        loop.run_until_complete(scenario())
+    finally:
+        loop.close()
+
+
 @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(