Merge 57492754e6 into aa769fbc69

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/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py

@@ -6,7 +6,8 @@ the pure placement rules on the OpenAI-format message list.
 """
 
 import time
-from itertools import chain
+from collections.abc import Iterator, Mapping
+from itertools import chain, count
 from typing import Final
 
 import pytest
@@ -327,14 +328,34 @@ EMPTY_ASSISTANT = pytest.mark.parametrize(
 )
 
 
-def test_flagged_placement_handles_a_long_run_of_unanswered_tool_search_uses_in_linear_time():
+class _ReadCountingMessage(Mapping[str, object]):
+    def __init__(self, reads: Iterator[int], fields: Mapping[str, object]) -> None:
+        self._reads: Final = reads
+        self._fields: Final = fields
+
+    def __getitem__(self, key: str) -> object:
+        next(self._reads)
+        return self._fields[key]
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._fields)
+
+    def __len__(self) -> int:
+        return len(self._fields)
+
+
+def _message_reads_placing_unanswered_tool_search_uses(search_use_count: int) -> int:
+    reads: Final = count()
     search_uses: Final = tuple(
-        {
-            "role": "assistant",
-            "content": None,
-            "tool_calls": [{**_TOOL_SEARCH_USE, "id": f"srvtoolu_{index:06d}"}],
-        }
-        for index in range(3000)
+        _ReadCountingMessage(
+            reads,
+            {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [{**_TOOL_SEARCH_USE, "id": f"srvtoolu_{index:06d}"}],
+            },
+        )
+        for index in range(search_use_count)
     )
     reminders: Final = tuple(
         chain.from_iterable(
@@ -349,12 +370,19 @@ def test_flagged_placement_handles_a_long_run_of_unanswered_tool_search_uses_in_
         {"role": "user", "content": "end"},
     ]
 
-    started: Final = time.perf_counter()
     placed: Final = place_mid_conversation_system(messages, supports_mid_conversation_system=True)
-    elapsed: Final = time.perf_counter() - started
+    total_reads: Final = next(reads)
 
-    assert set(_roles(placed)[:21]) == {"user"}
-    assert elapsed < 1.0
+    assert set(_roles(placed[:21])) == {"user"}
+    assert placed[21:-1] == search_uses
+    return total_reads
+
+
+def test_flagged_placement_handles_a_long_run_of_unanswered_tool_search_uses_in_linear_time():
+    reads_for_1500: Final = _message_reads_placing_unanswered_tool_search_uses(1500)
+    reads_for_3000: Final = _message_reads_placing_unanswered_tool_search_uses(3000)
+
+    assert reads_for_3000 < 3 * reads_for_1500
 
 
 def test_flagged_placement_merges_every_system_run_that_shares_a_user_turn():

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