scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Optimize spend log cleanup to use raw SQL DELETE with LIMIT

Browse source 
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 <ishaan@berri.ai>
This commit is contained in:
Cursor Agent 2026-01-30 20:13:45 +00:00
parent 974837c4e1
commit e34df5dfe0
2 changed files with 51 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
View file

@ -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

56
tests/test_litellm/proxy/test_spend_log_cleanup.py
View file

@ -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):