From dce4bbbfc223cf1716e8db41c8b99a22f63bcda4 Mon Sep 17 00:00:00 2001 From: jesus Date: Sat, 31 Oct 2026 22:38:50 +0000 Subject: [PATCH] fix(proxy): analyze the guardrail index and give its size-rule deletes their own batch budget --- .../db_transaction_queue/spend_log_cleanup.py | 21 ++++++-------- tests/unit/proxy/test_spend_log_cleanup.py | 28 ++++++++++++++++++- 2 files changed, 36 insertions(+), 13 deletions(-) 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 db03f74c07b..890b29a5721 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -4,7 +4,7 @@ import re import time from collections.abc import Awaitable, Callable from contextvars import ContextVar -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime, timedelta, timezone from typing import Final, Literal, TypeAlias @@ -255,6 +255,7 @@ class _SizeStepState: index_deleted: int guardrail_index_deleted: int batches_left: int + guardrail_batches_left: int stop_reason: StopReason | None @property @@ -262,13 +263,7 @@ class _SizeStepState: return self.logs_deleted + self.index_deleted + self.guardrail_index_deleted def stopped(self, reason: StopReason) -> "_SizeStepState": - return _SizeStepState( - logs_deleted=self.logs_deleted, - index_deleted=self.index_deleted, - guardrail_index_deleted=self.guardrail_index_deleted, - batches_left=self.batches_left, - stop_reason=reason, - ) + return replace(self, stop_reason=reason) @dataclass(frozen=True, slots=True) @@ -381,7 +376,7 @@ _STORAGE_SQL: Final = """ # The age job and the size job are scheduled on separate SpendLogCleanup instances, so the lock is per process. _PASS_LOCK: Final = asyncio.Lock() -_ANALYZE_SQL: Final = 'ANALYZE "LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"' +_ANALYZE_SQL: Final = "ANALYZE " + ", ".join(f'"{table}"' for table in _SPEND_LOG_TABLES) # pg_stat_user_tables publishes what ANALYZE counted up to a second later (the stats collector on Postgres 14). _STATS_SETTLE_SECONDS: Final = 1.0 _STATS_SETTLE_ATTEMPTS: Final = 5 @@ -1259,6 +1254,7 @@ class SpendLogCleanup: index_deleted=0, guardrail_index_deleted=0, batches_left=self.max_batches, + guardrail_batches_left=self.max_batches, stop_reason=None, ) for _ in range(self.max_batches // SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES + 2): @@ -1328,14 +1324,15 @@ class SpendLogCleanup: key_columns=("request_id", "guardrail_id"), time_column="start_time", deadline=deadline, - batch_cap=state.batches_left, + batch_cap=state.guardrail_batches_left, ) - guardrail_batches: Final = math.ceil(guardrail_index_result.rows_deleted / self.batch_size) advanced: Final = _SizeStepState( logs_deleted=state.logs_deleted + logs_result.rows_deleted, index_deleted=state.index_deleted + index_result.rows_deleted, guardrail_index_deleted=state.guardrail_index_deleted + guardrail_index_result.rows_deleted, - batches_left=state.batches_left - max(step_cap, guardrail_batches), + batches_left=state.batches_left - step_cap, + guardrail_batches_left=state.guardrail_batches_left + - math.ceil(guardrail_index_result.rows_deleted / self.batch_size), stop_reason=None, ) if rows_still_needed(advanced.rows_deleted) <= 0: diff --git a/tests/unit/proxy/test_spend_log_cleanup.py b/tests/unit/proxy/test_spend_log_cleanup.py index 899e34338a2..2a91cad60aa 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/tests/unit/proxy/test_spend_log_cleanup.py @@ -1937,6 +1937,32 @@ async def test_size_pass_deletes_every_guardrail_row_under_the_step_cutoff_not_o assert "LiteLLM_SpendLogGuardrailIndex (size rule): deleted=6000" in summary, summary +@pytest.mark.asyncio +async def test_size_pass_guardrail_backlog_does_not_spend_the_batches_spend_logs_need(caplog): + db = MagicMock() + _wire_tx(db) + + async def execute_raw(sql, *args): + return 0 if "LiteLLM_SpendLogToolIndex" in sql else 1000 + + db.execute_raw = execute_raw + _wire_size_probes( + db, [_storage(100 * 1024**3, live=100_000, dead=0)], cutoff=datetime(2026, 9, 1, tzinfo=timezone.utc) + ) + client = MagicMock() + client.db = db + cleaner = _size_cleaner( + {"maximum_spend_logs_cleanup_batch_size": 1000, "maximum_spend_logs_cleanup_max_batches": 20} + ) + + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + await cleaner.cleanup_spend_logs_by_size(client) + + summary = next(r.getMessage() for r in caplog.records if "Spend log cleanup run finished" in r.getMessage()) + assert "LiteLLM_SpendLogs (size rule): deleted=20000" in summary, summary + assert "LiteLLM_SpendLogGuardrailIndex (size rule): deleted=20000" in summary, summary + + @pytest.mark.asyncio async def test_size_pass_stops_at_the_run_batch_cap_and_reports_how_far_above_the_target_it_is(caplog): db = MagicMock() @@ -2004,7 +2030,7 @@ async def test_size_pass_analyzes_and_measures_again_when_the_row_statistics_wer await _size_cleaner({}).cleanup_spend_logs_by_size(client) issued = [call.args[0] for call in statements.await_args_list] - assert issued == ['ANALYZE "LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"'], ( + assert issued == ['ANALYZE "LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex", "LiteLLM_SpendLogGuardrailIndex"'], ( "the only statement is the ANALYZE: no DELETE, the carried count puts the table under 1GB" ) assert any("row statistics are stale" in record.getMessage() for record in caplog.records)