fix(proxy): analyze the guardrail index and give its size-rule deletes their own batch budget

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:

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)