From 1b3e39f7647db55f39c7440447438efdf2fa7774 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 20 Oct 2026 16:58:40 +0000 Subject: [PATCH 1/2] fix(proxy): prune expired LiteLLM_SpendLogGuardrailIndex rows with spend-log retention Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../request_log_indexes.py | 3 + .../litellm_proxy_extras/schema.prisma | 1 + .../db_transaction_queue/spend_log_cleanup.py | 35 ++++++++--- litellm/proxy/schema.prisma | 1 + schema.prisma | 1 + tests/unit/proxy/test_spend_log_cleanup.py | 61 ++++++++++++++++--- 6 files changed, 86 insertions(+), 16 deletions(-) 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 6c31e8364a9..3ced2abb417 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py @@ -53,6 +53,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 83a0e6f9e05..b2fadf07170 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1310,6 +1310,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 b57952b4992..00b4fbbfe0e 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -66,8 +66,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 @@ -277,7 +277,7 @@ _SAVE_SIZE_STATE_SQL: Final = """ """ -_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 = """ @@ -791,6 +791,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: @@ -855,7 +867,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 @@ -883,13 +895,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) + prisma_client, cutoff_date, self._group_deadline(deadline, groups_remaining=3) ) 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) + ) + 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_session_rollup( self, prisma_client: PrismaClient, retention_seconds: int, deadline: float diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 83a0e6f9e05..b2fadf07170 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1310,6 +1310,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 83a0e6f9e05..b2fadf07170 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1310,6 +1310,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 a9bc82d141d..5f15dec7d3b 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/tests/unit/proxy/test_spend_log_cleanup.py @@ -258,8 +258,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 @@ -279,7 +279,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] @@ -417,7 +417,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) @@ -436,7 +436,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 @@ -602,6 +602,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 — @@ -1524,7 +1564,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) @@ -1559,7 +1599,7 @@ async def test_progress_reported_by_an_overlapping_run_is_its_own(monkeypatch): slow_client.db.execute_raw = _slow_execute_raw fast_client = MagicMock() _wire_tx(fast_client.db) - fast_client.db.execute_raw = AsyncMock(side_effect=[150, 150, 0, 0]) + fast_client.db.execute_raw = AsyncMock(side_effect=[150, 150, 0, 0, 0]) cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cleaner.pod_lock_manager = None @@ -1608,6 +1648,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 elapsed < 2.5, f"splitting the budget must not extend the run: {elapsed}s" @@ -1800,7 +1841,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"): From e24225ecf7e27c939a3b4e2f667feaf99ad91ef7 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 20 Oct 2026 19:16:46 +0000 Subject: [PATCH 2/2] fix(proxy): size pass deletes every guardrail index row under the step cutoff Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../db_transaction_queue/spend_log_cleanup.py | 5 +++-- tests/unit/proxy/test_spend_log_cleanup.py | 21 ++++++++++++++++++- 2 files changed, 23 insertions(+), 3 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 db6a2943aa3..2fd34000466 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -1261,13 +1261,14 @@ class SpendLogCleanup: key_columns=("request_id", "guardrail_id"), time_column="start_time", deadline=deadline, - batch_cap=step_cap, + 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, guardrail_index_deleted=state.guardrail_index_deleted + guardrail_index_result.rows_deleted, - batches_left=state.batches_left - step_cap, + batches_left=state.batches_left - max(step_cap, guardrail_batches), 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 c9973c7cd7d..dfb2d1817cf 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/tests/unit/proxy/test_spend_log_cleanup.py @@ -1606,7 +1606,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) @@ -1885,6 +1885,25 @@ 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()