diff --git a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py index 4c856c741d3..32825f59860 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py @@ -59,6 +59,9 @@ class RequestLogIndex: REQUEST_LOG_INDEXES: Final = ( RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", '("api_key", "startTime")'), RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_litellm_call_id_idx", '("litellm_call_id")'), + RequestLogIndex( + "LiteLLM_SpendLogGuardrailIndex", "LiteLLM_SpendLogGuardrailIndex_start_time_idx", '("start_time")' + ), ) _IDENTIFIER_MAX_BYTES: Final = 63 diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 179b741fe7c..30ee6c35250 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1311,6 +1311,7 @@ model LiteLLM_SpendLogGuardrailIndex { @@id([request_id, guardrail_id]) @@index([guardrail_id, start_time]) @@index([policy_id, start_time]) + @@index([start_time]) } // Index for fast "last N logs for tool" from SpendLogs – see how a tool is called in production 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 170eb450339..db03f74c07b 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -77,8 +77,8 @@ class StorageProbe: """ What the catalog says about the spend-log tables, read without scanning them. - `allocated_bytes` is what the files take on disk (heap, indexes and TOAST of LiteLLM_SpendLogs and - LiteLLM_SpendLogToolIndex, partitions included). Deleting rows does not shrink it: deleted rows become dead + `allocated_bytes` is what the files take on disk (heap, indexes and TOAST of LiteLLM_SpendLogs, + LiteLLM_SpendLogToolIndex and LiteLLM_SpendLogGuardrailIndex, partitions included). Deleting rows does not shrink it: deleted rows become dead tuples, autovacuum later turns them into free space that new rows reuse, and only VACUUM FULL or pg_repack hands the bytes back to the filesystem. Neither dead tuples nor free space are live data, and Postgres tracks only the former (`n_dead_tup`, zeroed by vacuum), so the size rule keeps its own count of the @@ -253,13 +253,19 @@ class _SizeStepState: logs_deleted: int index_deleted: int + guardrail_index_deleted: int batches_left: int stop_reason: StopReason | None + @property + def rows_deleted(self) -> int: + 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, ) @@ -350,7 +356,7 @@ class _OffsetRow(LiteLLMBaseModel): _OFFSET_ROWS: Final = TypeAdapter(list[_OffsetRow]) -_SPEND_LOG_TABLES: Final = ("LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex") +_SPEND_LOG_TABLES: Final = ("LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex", "LiteLLM_SpendLogGuardrailIndex") # Partition children are counted with their parent: a partitioned parent has no storage of its own. _STORAGE_SQL: Final = """ @@ -992,6 +998,18 @@ class SpendLogCleanup: deadline=deadline, ) + async def _delete_old_guardrail_index_rows( + self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float + ) -> TableCleanupResult: + return await self._delete_old_rows_batched( + prisma_client, + cutoff_date, + table_name="LiteLLM_SpendLogGuardrailIndex", + key_columns=("request_id", "guardrail_id"), + time_column="start_time", + deadline=deadline, + ) + async def _delete_old_autorouter_session_rows( self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float ) -> TableCleanupResult: @@ -1056,7 +1074,7 @@ class SpendLogCleanup: self, prisma_client: PrismaClient, deadline: float ) -> tuple[TableCleanupResult, ...]: """ - Prune the spend logs and the tool index rows derived from them. + Prune the spend logs and the tool and guardrail index rows derived from them. When the table is range-partitioned, whole expired partitions are dropped first because that reclaims disk immediately. Expired rows can still sit in @@ -1084,13 +1102,20 @@ class SpendLogCleanup: verbose_proxy_logger.info("Dropped %d expired spend-log partitions: %s", len(dropped), dropped) logs_result: Final = await self._delete_old_logs( - prisma_client, cutoff_date, self._group_deadline(deadline, groups_remaining=2, now=self._clock()) + prisma_client, cutoff_date, self._group_deadline(deadline, groups_remaining=3, now=self._clock()) ) verbose_proxy_logger.info("Deleted %s logs", logs_result.rows_deleted) - index_result: Final = await self._delete_old_tool_index_rows(prisma_client, cutoff_date, deadline) - verbose_proxy_logger.info("Deleted %s expired tool index rows", index_result.rows_deleted) - return (logs_result, index_result) + tool_index_result: Final = await self._delete_old_tool_index_rows( + prisma_client, cutoff_date, self._group_deadline(deadline, groups_remaining=2, now=self._clock()) + ) + verbose_proxy_logger.info("Deleted %s expired tool index rows", tool_index_result.rows_deleted) + + guardrail_index_result: Final = await self._delete_old_guardrail_index_rows( + prisma_client, cutoff_date, deadline + ) + verbose_proxy_logger.info("Deleted %s expired guardrail index rows", guardrail_index_result.rows_deleted) + return (logs_result, tool_index_result, guardrail_index_result) async def _clean_spend_logs_by_size( self, prisma_client: PrismaClient, rule: SizeRule, deadline: float @@ -1104,7 +1129,7 @@ class SpendLogCleanup: once, re-measuring after each. Then rows go in steps: the rows still needed to reach the target follow from bytes per row, and each step turns that many of the oldest rows (at most SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES batches) into a cutoff and runs the same bounded DELETE the age - rule uses, for the spend logs and the tool index rows derived from them. Progress is counted from the + rule uses, for the spend logs and the tool and guardrail index rows derived from them. Progress is counted from the deletes themselves, not re-read from pg_stat_user_tables, whose counters trail a DELETE by up to a second. The run's batch cap and budget apply to the whole pass, so a pass that stops short reports how far above the target it still is and the next check continues from there. @@ -1230,7 +1255,11 @@ class SpendLogCleanup: rows_still_needed: Callable[[int], int], ) -> tuple[TableCleanupResult, ...]: state = _SizeStepState( # rebind-ok: one value threaded through the bounded step loop - logs_deleted=0, index_deleted=0, batches_left=self.max_batches, stop_reason=None + logs_deleted=0, + index_deleted=0, + guardrail_index_deleted=0, + batches_left=self.max_batches, + stop_reason=None, ) for _ in range(self.max_batches // SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES + 2): if state.stop_reason is not None: @@ -1246,6 +1275,11 @@ class SpendLogCleanup: rows_deleted=state.index_deleted, stop_reason=final_reason, ), + TableCleanupResult( + table_name="LiteLLM_SpendLogGuardrailIndex (size rule)", + rows_deleted=state.guardrail_index_deleted, + stop_reason=final_reason, + ), ) async def _run_size_step( @@ -1256,7 +1290,7 @@ class SpendLogCleanup: state: "_SizeStepState", ) -> "_SizeStepState": """One step: delete the rows still needed, at most SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES batches of the oldest.""" - needed: Final = rows_still_needed(state.logs_deleted + state.index_deleted) + needed: Final = rows_still_needed(state.rows_deleted) if needed <= 0: return state.stopped("target_reached") if self._clock() >= deadline: @@ -1287,15 +1321,29 @@ class SpendLogCleanup: deadline=deadline, batch_cap=step_cap, ) + guardrail_index_result: Final = await self._delete_old_rows_batched( + prisma_client, + cutoff, + table_name="LiteLLM_SpendLogGuardrailIndex", + key_columns=("request_id", "guardrail_id"), + time_column="start_time", + deadline=deadline, + batch_cap=state.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, - batches_left=state.batches_left - step_cap, + guardrail_index_deleted=state.guardrail_index_deleted + guardrail_index_result.rows_deleted, + batches_left=state.batches_left - max(step_cap, guardrail_batches), stop_reason=None, ) - if rows_still_needed(advanced.logs_deleted + advanced.index_deleted) <= 0: + if rows_still_needed(advanced.rows_deleted) <= 0: return advanced.stopped("target_reached") - if logs_result.stop_reason in ("aborted", "budget_exhausted") or index_result.stop_reason == "aborted": + if logs_result.stop_reason in ("aborted", "budget_exhausted") or "aborted" in ( + index_result.stop_reason, + guardrail_index_result.stop_reason, + ): return advanced.stopped(logs_result.stop_reason if logs_result.stop_reason != "exhausted" else "aborted") if logs_result.rows_deleted == 0: return advanced.stopped("exhausted") diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 179b741fe7c..30ee6c35250 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1311,6 +1311,7 @@ model LiteLLM_SpendLogGuardrailIndex { @@id([request_id, guardrail_id]) @@index([guardrail_id, start_time]) @@index([policy_id, start_time]) + @@index([start_time]) } // Index for fast "last N logs for tool" from SpendLogs – see how a tool is called in production diff --git a/schema.prisma b/schema.prisma index 179b741fe7c..30ee6c35250 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1311,6 +1311,7 @@ model LiteLLM_SpendLogGuardrailIndex { @@id([request_id, guardrail_id]) @@index([guardrail_id, start_time]) @@index([policy_id, start_time]) + @@index([start_time]) } // Index for fast "last N logs for tool" from SpendLogs – see how a tool is called in production diff --git a/tests/unit/proxy/test_spend_log_cleanup.py b/tests/unit/proxy/test_spend_log_cleanup.py index 114c970e163..899e34338a2 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/tests/unit/proxy/test_spend_log_cleanup.py @@ -270,8 +270,8 @@ async def test_cleanup_old_spend_logs_batch_deletion(): _wire_tx(mock_db) # Mock execute_raw to return deleted counts (3 spend-log batches, then the - # tool-index cleanup's first batch returning 0) - mock_db.execute_raw = AsyncMock(side_effect=[1000, 500, 0, 0]) + # tool-index and guardrail-index cleanups' first batches returning 0) + mock_db.execute_raw = AsyncMock(side_effect=[1000, 500, 0, 0, 0]) # Wire up mocks mock_prisma_client.db = mock_db @@ -291,7 +291,7 @@ async def test_cleanup_old_spend_logs_batch_deletion(): await cleaner.cleanup_old_spend_logs(mock_prisma_client) # Validate batching and deletion via raw SQL - assert mock_db.execute_raw.call_count == 4 + assert mock_db.execute_raw.call_count == 5 # Check the first call argument call_args_sql = mock_db.execute_raw.call_args_list[0][0][0] @@ -429,7 +429,7 @@ async def test_cleanup_uses_delete_when_not_partitioned(): mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) - mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0, 0]) + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0, 0, 0]) partition_manager = MagicMock() partition_manager.is_partitioned = AsyncMock(return_value=False) @@ -448,7 +448,7 @@ async def test_cleanup_uses_delete_when_not_partitioned(): await cleaner.cleanup_old_spend_logs(mock_prisma_client) partition_manager.drop_partitions_older_than.assert_not_awaited() - assert mock_prisma_client.db.execute_raw.await_count == 3 + assert mock_prisma_client.db.execute_raw.await_count == 4 delete_sql = mock_prisma_client.db.execute_raw.call_args_list[0][0][0] assert 'DELETE FROM "LiteLLM_SpendLogs"' in delete_sql @@ -614,6 +614,46 @@ async def test_delete_old_tool_index_rows_deletes_on_composite_key(): assert mock_db.execute_raw.call_args_list[0][0][1] == cutoff_date +@pytest.mark.asyncio +async def test_spend_log_retention_prunes_expired_guardrail_index_rows_on_the_same_cutoff(): + mock_prisma_client = MagicMock() + mock_db = MagicMock() + _wire_tx(mock_db) + mock_db.execute_raw = AsyncMock(return_value=0) + mock_prisma_client.db = mock_db + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner.pod_lock_manager = None + + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + deletes = { + call.args[0].split('"')[1]: call + for call in mock_db.execute_raw.call_args_list + if call.args[0].lstrip().startswith("DELETE") + } + guardrail_delete = deletes.get("LiteLLM_SpendLogGuardrailIndex") + assert guardrail_delete is not None, f"no guardrail index delete was issued: {sorted(deletes)}" + assert 'WHERE ("request_id", "guardrail_id") IN' in guardrail_delete.args[0] + assert '"start_time" <' in guardrail_delete.args[0] + assert guardrail_delete.args[1] == deletes["LiteLLM_SpendLogs"].args[1] + + +@pytest.mark.asyncio +async def test_guardrail_index_is_not_pruned_when_spend_log_retention_is_unset(): + mock_prisma_client = MagicMock() + mock_db = MagicMock() + _wire_tx(mock_db) + mock_db.execute_raw = AsyncMock(return_value=0) + mock_prisma_client.db = mock_db + cleaner = SpendLogCleanup(general_settings={"maximum_health_check_retention_period": "7d"}) + cleaner.pod_lock_manager = None + + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + statements = [call.args[0] for call in mock_db.execute_raw.call_args_list] + assert not any("LiteLLM_SpendLogGuardrailIndex" in sql for sql in statements), statements + + @pytest.mark.asyncio async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch): """A single batch failure (e.g. DB timeout) must not abort the whole run — @@ -1552,7 +1592,7 @@ async def test_progress_reported_for_a_cancelled_run_is_that_run_only(monkeypatc mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) - mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, 0, 0]) + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, 0, 0, 0]) cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cleaner.pod_lock_manager = None await cleaner.cleanup_old_spend_logs(mock_prisma_client) @@ -1594,7 +1634,7 @@ async def test_a_size_pass_that_fires_while_an_age_pass_is_running_in_the_same_p age_cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) age_cleaner.pod_lock_manager = None - size_cleaner = _size_cleaner({}) + size_cleaner = _size_cleaner({"maximum_spend_logs_cleanup_max_batches": 3}) slow_run = asyncio.ensure_future(age_cleaner.cleanup_old_spend_logs(slow_client)) await asyncio.wait_for(first_batch_done.wait(), timeout=5) @@ -1650,6 +1690,7 @@ async def test_spend_logs_backlog_cannot_starve_tool_index_cleanup(): tables = [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 tables.count("LiteLLM_SpendLogGuardrailIndex") > 0, "guardrail index cleanup was starved by the backlog" assert clock.now < 1.15, f"splitting the budget must not extend the run: {clock.now}s" @@ -1835,14 +1876,14 @@ async def test_size_pass_under_the_limit_issues_no_delete(): async def test_size_pass_deletes_oldest_first_in_steps_until_the_live_estimate_is_under_the_target(caplog): """ Over the limit, the pass turns the next step of oldest rows into a cutoff and runs the retention DELETE - for the spend logs and their tool-index rows, re-measuring after each step; it stops on the first + for the spend logs and their tool and guardrail index rows, re-measuring after each step; it stops on the first measurement under the target, not when the file size (which deletes never shrink) moves. """ db = MagicMock() _wire_tx(db) # 3 GiB over 3000 rows is 1 MiB a row, so 2560 rows take it from 3 GiB to the 512 MiB target: one step - # of 3 batches; the tool index is cleaned on the same cutoff - db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 300, 0]) + # of 3 batches; the tool and guardrail indexes are cleaned on the same cutoff + db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 300, 0, 200, 0]) step_cutoff = datetime(2026, 9, 1, tzinfo=timezone.utc) allocated = 3 * 1024**3 _wire_size_probes(db, [_storage(allocated, live=3000, dead=0)], cutoff=step_cutoff) @@ -1854,18 +1895,22 @@ async def test_size_pass_deletes_oldest_first_in_steps_until_the_live_estimate_i await cleaner.cleanup_spend_logs_by_size(client) deletes = [(call[0][0], call[0][1:]) for call in db.deletes.call_args_list] - assert len(deletes) == 5 + assert len(deletes) == 7 assert 'DELETE FROM "LiteLLM_SpendLogs"' in deletes[0][0] assert 'WHERE "startTime" < $1' in deletes[0][0] assert deletes[0][1][0] == step_cutoff assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in deletes[3][0] assert deletes[3][1][0] == step_cutoff + assert 'DELETE FROM "LiteLLM_SpendLogGuardrailIndex"' in deletes[5][0] + assert 'WHERE ("request_id", "guardrail_id") IN' in deletes[5][0] + assert deletes[5][1][0] == step_cutoff offset_query = next(call for call in db.query_raw.call_args_list if "OFFSET" in call[0][0]) assert offset_query[0][1] == 3000 summary = next(r.getMessage() for r in caplog.records if "Spend log cleanup run finished" in r.getMessage()) assert "outcome=completed" in summary assert "LiteLLM_SpendLogs (size rule): deleted=3000 stop_reason=target_reached" in summary assert "LiteLLM_SpendLogToolIndex (size rule): deleted=300 stop_reason=target_reached" in summary + assert "LiteLLM_SpendLogGuardrailIndex (size rule): deleted=200 stop_reason=target_reached" in summary assert f"allocated_before={allocated}" in summary assert "estimated_live_before=3221225472" in summary assert "estimated_live_after=0" in summary @@ -1873,12 +1918,31 @@ async def test_size_pass_deletes_oldest_first_in_steps_until_the_live_estimate_i assert "VACUUM FULL or pg_repack" in summary +@pytest.mark.asyncio +async def test_size_pass_deletes_every_guardrail_row_under_the_step_cutoff_not_one_per_spend_log(caplog): + db = MagicMock() + _wire_tx(db) + # one step of 3 SpendLogs batches, but two guardrails per request leave twice as many guardrail rows under it + db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 0, 1000, 1000, 1000, 1000, 1000, 1000, 0]) + _wire_size_probes(db, [_storage(3 * 1024**3, live=3000, dead=0)], cutoff=datetime(2026, 9, 1, tzinfo=timezone.utc)) + client = MagicMock() + client.db = db + cleaner = _size_cleaner({"maximum_spend_logs_cleanup_target_size": "512MB"}) + + 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=3000" in summary, summary + assert "LiteLLM_SpendLogGuardrailIndex (size rule): deleted=6000" 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() _wire_tx(db) # 4 GiB over 4000 rows: 3584 rows must go to reach the 512 MiB target, the cap allows 3 batches - db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 0]) + db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 0, 0]) allocated = 4 * 1024**3 saves = _wire_size_probes(db, [_storage(allocated, live=4000, dead=0)], cutoff=datetime.now(timezone.utc)) client = MagicMock() @@ -1890,7 +1954,7 @@ async def test_size_pass_stops_at_the_run_batch_cap_and_reports_how_far_above_th with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): await cleaner.cleanup_spend_logs_by_size(client) - assert db.deletes.await_count == 4 + assert db.deletes.await_count == 5 summary = next(r for r in caplog.records if "Spend log cleanup run finished" in r.getMessage()) assert summary.levelno == logging.WARNING assert "outcome=batch_cap_reached" in summary.getMessage() @@ -2052,7 +2116,7 @@ async def test_size_only_run_leaves_the_age_rules_to_the_retention_job(): await cleaner.cleanup_spend_logs_by_size(client) tables = {call[0][0].split('"')[1] for call in db.deletes.call_args_list} - assert tables == {"LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"} + assert tables == {"LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex", "LiteLLM_SpendLogGuardrailIndex"} @pytest.mark.asyncio @@ -2232,7 +2296,11 @@ async def test_probe_storage_reads_the_catalog_row_for_both_tables_and_a_failed_ 20_000, 0, ) - assert prisma_client.db.query_raw.call_args.args[1] == ["LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"] + assert prisma_client.db.query_raw.call_args.args[1] == [ + "LiteLLM_SpendLogs", + "LiteLLM_SpendLogToolIndex", + "LiteLLM_SpendLogGuardrailIndex", + ] prisma_client.db.query_raw = AsyncMock(side_effect=RuntimeError("pg_stat_user_tables unavailable")) with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):