test(spend_log_cleanup): inject clock and sleep so the tool index budget split test stops racing the wall clock

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py

@@ -421,7 +421,12 @@ class SpendLogCleanup:
         general_settings=None,
         redis_cache: RedisCache | None = None,
         partition_manager: SpendLogsPartitionManager | None = None,
+        *,
+        clock: Callable[[], float] = time.monotonic,
+        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
     ):
+        self._clock = clock
+        self._sleep = sleep
         self.retention_seconds: int | None = None
         self.partition_manager = partition_manager or SpendLogsPartitionManager()
         from litellm.proxy.proxy_server import general_settings as default_settings
@@ -586,7 +591,7 @@ class SpendLogCleanup:
 
     async def _probe_storage(self, prisma_client: PrismaClient, deadline: float) -> StorageProbe | None:
         """Read what the spend-log tables weigh from the catalog, or None when the probe fails or has no budget."""
-        if time.monotonic() >= deadline:
+        if self._clock() >= deadline:
             return None
         timeout_ms: Final = self._timeout_ms(deadline)
         try:
@@ -624,7 +629,7 @@ class SpendLogCleanup:
             stale.allocated_bytes,
             _ANALYZE_SQL,
         )
-        if time.monotonic() >= deadline:
+        if self._clock() >= deadline:
             return None
         timeout_ms: Final = self._timeout_ms(deadline)
         try:
@@ -645,7 +650,7 @@ class SpendLogCleanup:
     async def _probe_once_settled(
         self, prisma_client: PrismaClient, deadline: float, attempts_left: int
     ) -> StorageProbe | None:
-        await asyncio.sleep(_STATS_SETTLE_SECONDS)
+        await self._sleep(_STATS_SETTLE_SECONDS)
         refreshed: Final = await self._probe_storage(prisma_client, deadline)
         if refreshed is None or not refreshed.stats_are_stale or attempts_left <= 1:
             return refreshed
@@ -653,7 +658,7 @@ class SpendLogCleanup:
 
     async def _load_size_state(self, prisma_client: PrismaClient, deadline: float) -> SizeCleanupState | None:
         """The reclaimable count saved by the last check, or None on the first check or when it cannot be read."""
-        if time.monotonic() >= deadline:
+        if self._clock() >= deadline:
             return None
         timeout_ms: Final = self._timeout_ms(deadline)
         try:
@@ -686,7 +691,7 @@ class SpendLogCleanup:
         the tool-index rows that expire with their spend logs. When fewer rows than that remain, the step is
         everything up to now.
         """
-        if time.monotonic() >= deadline:
+        if self._clock() >= deadline:
             return None
         timeout_ms: Final = self._timeout_ms(deadline)
         async with db_span("size_step_cutoff", "LiteLLM_SpendLogs"), bounded_tx(prisma_client, timeout_ms) as tx:
@@ -718,11 +723,10 @@ class SpendLogCleanup:
         Interpolating this into SQL is safe by construction: an int cannot carry
         SQL, and SET does not accept a bind parameter.
         """
-        remaining_ms: Final = int((deadline - time.monotonic()) * 1000)
+        remaining_ms: Final = int((deadline - self._clock()) * 1000)
         return max(1, min(int(self.batch_timeout_seconds * 1000), remaining_ms))
 
-    @staticmethod
-    def _group_deadline(overall_deadline: float, groups_remaining: int) -> float:
+    def _group_deadline(self, overall_deadline: float, groups_remaining: int) -> float:
         """
         Give each pending cleanup group an equal share of the time left.
 
@@ -731,7 +735,7 @@ class SpendLogCleanup:
         """
         if groups_remaining == 1:
             return overall_deadline
-        current_time: Final = time.monotonic()
+        current_time: Final = self._clock()
         if current_time >= overall_deadline:
             return overall_deadline
         return current_time + (overall_deadline - current_time) / groups_remaining
@@ -748,7 +752,7 @@ class SpendLogCleanup:
         """
 
         def remaining() -> int | None:
-            return None if time.monotonic() >= deadline else self._timeout_ms(deadline)
+            return None if self._clock() >= deadline else self._timeout_ms(deadline)
 
         return remaining
 
@@ -841,7 +845,7 @@ class SpendLogCleanup:
         run_count = 0
         consecutive_failures = 0
         while True:
-            if time.monotonic() >= deadline:
+            if self._clock() >= deadline:
                 verbose_proxy_logger.info(
                     "Run budget exhausted during %s cleanup after %d rows; the next run resumes from here",
                     table_name,
@@ -858,13 +862,13 @@ class SpendLogCleanup:
                     prisma_client, cutoff_date, table_name, time_column, total_deleted, "batch_cap_reached", deadline
                 )
             # Find rows and delete them in one go without fetching to application
-            batch_started_at = time.monotonic()
+            batch_started_at = self._clock()
             try:
                 batch_result = await self._execute_delete_batch(
                     prisma_client, delete_sql, cutoff_date, table_name, deadline
                 )
             except Exception as batch_exc:
-                if time.monotonic() >= deadline:
+                if self._clock() >= deadline:
                     # The statement timeout was clamped to the budget that was
                     # left, so this batch was cancelled by the deadline itself.
                     # That is the bound working, not a database fault, and
@@ -906,7 +910,7 @@ class SpendLogCleanup:
                     return await self._finish_table(
                         prisma_client, cutoff_date, table_name, time_column, total_deleted, "aborted", deadline
                     )
-                await asyncio.sleep(SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS)
+                await self._sleep(SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS)
                 continue
 
             if batch_result is None:
@@ -920,7 +924,7 @@ class SpendLogCleanup:
 
             consecutive_failures = 0
             deleted_count = batch_result
-            SpendLogCleanupMetrics.record_batch(table_name, deleted_count, time.monotonic() - batch_started_at)
+            SpendLogCleanupMetrics.record_batch(table_name, deleted_count, self._clock() - batch_started_at)
             verbose_proxy_logger.info("Deleted %s %s rows in this batch", deleted_count, table_name)
 
             if deleted_count == 0:
@@ -934,7 +938,7 @@ class SpendLogCleanup:
             _record_run_batch(deleted_count)
 
             # Add a small sleep to prevent overwhelming the database
-            await asyncio.sleep(0.1)
+            await self._sleep(0.1)
 
     async def _finish_table(
         self,
@@ -957,7 +961,7 @@ class SpendLogCleanup:
         operator the backlog was not drained; the gauge simply keeps its value
         from the last run that finished inside its budget.
         """
-        if time.monotonic() >= deadline:
+        if self._clock() >= deadline:
             return TableCleanupResult(table_name=table_name, rows_deleted=rows_deleted, stop_reason=stop_reason)
         remaining: Final = await self._count_remaining(prisma_client, cutoff_date, table_name, time_column, deadline)
         if remaining is not None:
@@ -1070,7 +1074,7 @@ class SpendLogCleanup:
         # behind any long-running reader for as long as that reader lives, which
         # is the one way this job could still outlast its budget without bound.
         remaining_timeout_ms: Final = self._remaining_timeout_ms(deadline)
-        if time.monotonic() >= deadline:
+        if self._clock() >= deadline:
             verbose_proxy_logger.info("Run budget already spent, skipping partition maintenance this run")
         elif self.general_settings.get(
             "use_spend_logs_partitioning", False
@@ -1210,7 +1214,7 @@ class SpendLogCleanup:
         self, prisma_client: PrismaClient, deadline: float, still_over_target: Callable[[], Awaitable[bool]]
     ) -> tuple[str, ...]:
         remaining_timeout_ms: Final = self._remaining_timeout_ms(deadline)
-        if time.monotonic() >= deadline or not self._setting("use_spend_logs_partitioning"):
+        if self._clock() >= deadline or not self._setting("use_spend_logs_partitioning"):
             return ()
         if not await self.partition_manager.is_partitioned(prisma_client, remaining_timeout_ms):
             return ()
@@ -1257,7 +1261,7 @@ class SpendLogCleanup:
         needed: Final = rows_still_needed(state.logs_deleted + state.index_deleted)
         if needed <= 0:
             return state.stopped("target_reached")
-        if time.monotonic() >= deadline:
+        if self._clock() >= deadline:
             return state.stopped("budget_exhausted")
         if state.batches_left <= 0:
             return state.stopped("batch_cap_reached")
@@ -1443,7 +1447,7 @@ class SpendLogCleanup:
 
     async def _run_passes(self, prisma_client: PrismaClient, age_rules: bool, size_rule: bool) -> None:
         lock_acquired = False
-        run_started_at: Final = time.monotonic()
+        run_started_at: Final = self._clock()
         progress: Final = _RunProgress()
         progress_token: Final = _run_progress.set(progress)
         try:
@@ -1493,7 +1497,7 @@ class SpendLogCleanup:
                     SpendLogCleanupMetrics.record_run("skipped_locked")
                     return
 
-            deadline: Final = time.monotonic() + self.run_budget_seconds
+            deadline: Final = self._clock() + self.run_budget_seconds
             size_group: Final = int(active_size_rule is not None)
             configured_group_count: Final = (
                 int(delete_spend_logs and self.retention_seconds is not None)
@@ -1561,12 +1565,12 @@ class SpendLogCleanup:
             )
             outcome: Final = self._run_outcome(results)
             SpendLogCleanupMetrics.record_run(outcome)
-            self._log_run_summary(outcome, results, time.monotonic() - run_started_at, size_report)
+            self._log_run_summary(outcome, results, self._clock() - run_started_at, size_report)
 
         except asyncio.CancelledError:
             verbose_proxy_logger.error(
                 "Spend log cleanup cancelled after %.2fs (rows_deleted=%d, batches=%d); the next run resumes from here",
-                time.monotonic() - run_started_at,
+                self._clock() - run_started_at,
                 progress.rows_deleted,
                 progress.batches,
             )

tests/unit/proxy/test_spend_log_cleanup.py

@@ -1588,36 +1588,49 @@ async def test_a_size_pass_that_fires_while_an_age_pass_is_running_in_the_same_p
 
 
 @pytest.mark.asyncio
-async def test_spend_logs_backlog_cannot_starve_tool_index_cleanup():
+async def test_spend_logs_backlog_cannot_starve_tool_index_cleanup() -> None:
     """
     Both spend-log tables share one run budget. Before the fix the spend-log
     loop ran against the whole deadline, so a backlog that outlasted the budget
     meant LiteLLM_SpendLogToolIndex never received a single delete batch, run
     after run. The index table must still get its own share of the budget.
     """
-    mock_prisma_client = MagicMock()
-    mock_db = MagicMock()
+    class FakeClock:
+        def __init__(self) -> None:
+            self.now = 0.0
+
+        def read(self) -> float:
+            return self.now
+
+        async def advance(self, seconds: float) -> None:
+            self.now += seconds
+
+    fake: Final = FakeClock()
+    mock_prisma_client: Final = MagicMock()
+    mock_db: Final = MagicMock()
     _wire_tx(mock_db)
     mock_db.execute_raw = AsyncMock(return_value=1000)
     mock_prisma_client.db = mock_db
 
-    cleaner = SpendLogCleanup(
+    cleaner: Final = SpendLogCleanup(
         general_settings={
             "maximum_spend_logs_retention_period": "7d",
             "maximum_spend_logs_cleanup_max_batches": 500,
             "maximum_spend_logs_cleanup_run_budget": "1s",
-        }
+        },
+        clock=fake.read,
+        sleep=fake.advance,
     )
     cleaner.pod_lock_manager = None
 
-    started_at = time.monotonic()
+    start: Final = fake.read()
     await cleaner.cleanup_old_spend_logs(mock_prisma_client)
-    elapsed = time.monotonic() - started_at
+    fake_now_at_end: Final = fake.read()
 
-    tables = [call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list]
+    tables: Final = [call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list]
     assert tables.count("LiteLLM_SpendLogs") > 0
     assert tables.count("LiteLLM_SpendLogToolIndex") > 0, "tool index cleanup was starved by the spend-log backlog"
-    assert elapsed < 2.5, f"splitting the budget must not extend the run: {elapsed}s"
+    assert fake_now_at_end - start <= 1.0 + 0.15
 
 
 @pytest.mark.asyncio