diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 7f5efcdaec3..eb4ee170819 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/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, ) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py index aab57b87a90..ab109523cf6 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py +++ b/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(): diff --git a/tests/unit/proxy/test_spend_log_cleanup.py b/tests/unit/proxy/test_spend_log_cleanup.py index 876ffe96584..3796eceda00 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/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