From e34df5dfe013901c7ed0b278602d648dea599386 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 30 Jan 2026 20:13:45 +0000 Subject: [PATCH] Optimize spend log cleanup to use raw SQL DELETE with LIMIT This change addresses performance concerns with the previous implementation: - **Previous approach**: Used find_many() to load records into memory, then delete_many() with the fetched IDs. This required: - Loading records into application memory - Network transfer of record data - Memory constraints limiting batch size - **New approach**: Uses raw SQL DELETE with LIMIT in a subquery to: - Delete directly in the database without loading records - Eliminate network transfer of record data - Remove memory constraints on batch size - Significantly improve deletion speed for high-volume deployments For deployments with 10M+ requests/day, this optimization dramatically reduces the time and resources needed for spend log cleanup. The DELETE query uses a subquery pattern that PostgreSQL handles efficiently: DELETE FROM "LiteLLM_SpendLogs" WHERE request_id IN ( SELECT request_id FROM "LiteLLM_SpendLogs" WHERE "startTime" < cutoff_date LIMIT batch_size ) Tests have been updated to mock execute_raw instead of find_many/delete_many. Co-authored-by: ishaan --- .../db_transaction_queue/spend_log_cleanup.py | 47 ++++++++++------ .../proxy/test_spend_log_cleanup.py | 56 ++++++++----------- 2 files changed, 51 insertions(+), 52 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 02fa84bae30..7844b986543 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -65,38 +65,49 @@ class SpendLogCleanup: self, prisma_client: PrismaClient, cutoff_date: datetime ) -> int: """ - Helper method to delete old logs in batches. + Helper method to delete old logs in batches using raw SQL for efficiency. + + Uses DELETE with a subquery LIMIT to avoid: + - Loading records into memory + - Network transfer of record data + - Memory constraints on batch size + Returns the total number of logs deleted. """ total_deleted = 0 run_count = 0 while True: - if run_count > SPEND_LOG_RUN_LOOPS: + if run_count >= SPEND_LOG_RUN_LOOPS: verbose_proxy_logger.info( - "Max logs deleted - 1,00,000, rest of the logs will be deleted in next run" + f"Max loops reached ({SPEND_LOG_RUN_LOOPS}). Deleted {total_deleted} logs. " + "Rest will be deleted in next run." ) break - # Step 1: Find logs to delete - logs_to_delete = await prisma_client.db.litellm_spendlogs.find_many( - where={"startTime": {"lt": cutoff_date}}, - take=self.batch_size, - ) - verbose_proxy_logger.info(f"Found {len(logs_to_delete)} logs in this batch") - if not logs_to_delete: + # Use raw SQL with LIMIT for efficient batched deletion + # This avoids loading records into memory and reduces network transfer + deleted_count = await prisma_client.db.execute_raw( + ''' + DELETE FROM "LiteLLM_SpendLogs" + WHERE "request_id" IN ( + SELECT "request_id" FROM "LiteLLM_SpendLogs" + WHERE "startTime" < $1 + LIMIT $2 + ) + ''', + cutoff_date, + self.batch_size, + ) + + verbose_proxy_logger.info(f"Deleted {deleted_count} logs in this batch") + + if deleted_count == 0: verbose_proxy_logger.info( f"No more logs to delete. Total deleted: {total_deleted}" ) break - request_ids = [log.request_id for log in logs_to_delete] - - # Step 2: Delete them in one go - await prisma_client.db.litellm_spendlogs.delete_many( - where={"request_id": {"in": request_ids}} - ) - - total_deleted += len(logs_to_delete) + total_deleted += deleted_count run_count += 1 # Add a small sleep to prevent overwhelming the database diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index 1ffbb83caef..ebe77cd5de6 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -151,28 +151,17 @@ async def test_should_delete_spend_logs(): @pytest.mark.asyncio async def test_cleanup_old_spend_logs_batch_deletion(): - from types import SimpleNamespace - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import AsyncMock, MagicMock # Setup Prisma client mock_prisma_client = MagicMock() mock_db = MagicMock() - # Mock spendlogs table - mock_spendlogs = MagicMock() - mock_spendlogs.find_many = AsyncMock() - mock_spendlogs.delete_many = AsyncMock() - - # Create 1500 mocked logs with .request_id - mock_logs = [SimpleNamespace(request_id=f"req_{i}") for i in range(1500)] - mock_spendlogs.find_many.side_effect = [ - mock_logs[:1000], # Batch 1 - mock_logs[1000:], # Batch 2 - [], # Done - ] + # Mock execute_raw to simulate batched deletion + # Returns 1000 (batch 1), 500 (batch 2), 0 (done) + mock_db.execute_raw = AsyncMock(side_effect=[1000, 500, 0]) # Wire up mocks - mock_db.litellm_spendlogs = mock_spendlogs mock_prisma_client.db = mock_db # Mock Redis cache and pod_lock_manager @@ -189,15 +178,16 @@ async def test_cleanup_old_spend_logs_batch_deletion(): assert cleaner._should_delete_spend_logs() is True await cleaner.cleanup_old_spend_logs(mock_prisma_client) - # Validate batching and deletion - assert mock_spendlogs.find_many.call_count == 3 - assert mock_spendlogs.delete_many.call_count == 2 - mock_spendlogs.delete_many.assert_any_call( - where={"request_id": {"in": [f"req_{i}" for i in range(1000)]}} - ) - mock_spendlogs.delete_many.assert_any_call( - where={"request_id": {"in": [f"req_{i}" for i in range(1000, 1500)]}} - ) + # Validate execute_raw was called for batched deletion + # Should be called 3 times: 1000 deleted, 500 deleted, 0 deleted (done) + assert mock_db.execute_raw.call_count == 3 + + # Verify the SQL query structure in each call + for call in mock_db.execute_raw.call_args_list: + sql_query = call[0][0] + assert 'DELETE FROM "LiteLLM_SpendLogs"' in sql_query + assert '"startTime" < $1' in sql_query + assert "LIMIT $2" in sql_query @pytest.mark.asyncio @@ -208,10 +198,8 @@ async def test_cleanup_old_spend_logs_retention_period_cutoff(): # Setup Prisma client mock_prisma_client = MagicMock() mock_db = MagicMock() - mock_spendlogs = MagicMock() - mock_spendlogs.find_many = AsyncMock(return_value=[]) - mock_spendlogs.delete_many = AsyncMock() - mock_db.litellm_spendlogs = mock_spendlogs + # Return 0 to indicate no logs deleted (simulating empty table) + mock_db.execute_raw = AsyncMock(return_value=0) mock_prisma_client.db = mock_db # Mock Redis cache and pod_lock_manager @@ -228,8 +216,10 @@ async def test_cleanup_old_spend_logs_retention_period_cutoff(): assert cleaner._should_delete_spend_logs() is True await cleaner.cleanup_old_spend_logs(mock_prisma_client) - # Verify the cutoff date is correct - cutoff_date = mock_spendlogs.find_many.call_args[1]["where"]["startTime"]["lt"] + # Verify execute_raw was called with the correct cutoff date + assert mock_db.execute_raw.call_count == 1 + call_args = mock_db.execute_raw.call_args[0] + cutoff_date = call_args[1] # Second positional arg is the cutoff date expected_cutoff = datetime.now(timezone.utc) - timedelta(seconds=86400) assert ( abs((cutoff_date - expected_cutoff).total_seconds()) < 1 @@ -242,14 +232,12 @@ async def test_cleanup_old_spend_logs_no_retention_period(): Test that no logs are deleted when no retention period is set """ mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_spendlogs.find_many = AsyncMock() - mock_prisma_client.db.litellm_spendlogs.delete = AsyncMock() + mock_prisma_client.db.execute_raw = AsyncMock() cleaner = SpendLogCleanup(general_settings={}) # no retention await cleaner.cleanup_old_spend_logs(mock_prisma_client) - mock_prisma_client.db.litellm_spendlogs.find_many.assert_not_called() - mock_prisma_client.db.litellm_spendlogs.delete.assert_not_called() + mock_prisma_client.db.execute_raw.assert_not_called() def test_cleanup_batch_size_env_var(monkeypatch):