mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
Feat log retention cleanup (#10872)
* Batch deletion with tasks * reset utils * change print to logger * remove print for lint * fix lint * local testing * remove constants.py var * ruff fix * type fix * - make sure job doesn't get added if maximum_spend_logs_retention_period is not set - enable user to override scheduler time - Upper bound on while true * update and fix test
This commit is contained in:
Jugal D. Bhatt
GitHub
parent
fdfef04d93
commit
d5d5166777
3 changed files with 203 additions and 33 deletions
|
|
@ -1,39 +1,90 @@
|
|||
"""
|
||||
Handles checking if spend logs should be deleted based on maximum retention period
|
||||
"""
|
||||
|
||||
from typing import Optional, Union
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from datetime import datetime, timedelta, UTC
|
||||
from typing import Optional
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
|
||||
def _should_delete_spend_logs() -> bool:
|
||||
class SpendLogCleanup:
|
||||
"""
|
||||
Checks if the Pod should delete spend logs based on maximum retention period
|
||||
|
||||
This setting enables automatic deletion of old spend logs to manage database size.
|
||||
The maximum_spend_logs_retention_period can be specified in:
|
||||
- Days (e.g., "30d")
|
||||
- Hours (e.g., "24h")
|
||||
- Minutes (e.g., "60m")
|
||||
- Seconds (e.g., "3600s" or just "3600")
|
||||
Handles cleaning up old spend logs based on maximum retention period.
|
||||
Deletes logs in batches to prevent timeouts.
|
||||
"""
|
||||
_maximum_spend_logs_retention_period: Optional[Union[int, str]] = general_settings.get(
|
||||
"maximum_spend_logs_retention_period", None
|
||||
)
|
||||
|
||||
if _maximum_spend_logs_retention_period is None:
|
||||
return False
|
||||
|
||||
try:
|
||||
if isinstance(_maximum_spend_logs_retention_period, int):
|
||||
_maximum_spend_logs_retention_period = str(_maximum_spend_logs_retention_period)
|
||||
duration_in_seconds(_maximum_spend_logs_retention_period)
|
||||
return True
|
||||
except ValueError as e:
|
||||
verbose_proxy_logger.error(
|
||||
f"Invalid maximum_spend_logs_retention_period value: {_maximum_spend_logs_retention_period}, error: {str(e)}"
|
||||
)
|
||||
return False
|
||||
def __init__(self, general_settings=None):
|
||||
self.batch_size = 1000
|
||||
self.retention_seconds: Optional[int] = None
|
||||
from litellm.proxy.proxy_server import general_settings as default_settings
|
||||
self.general_settings = general_settings or default_settings
|
||||
verbose_proxy_logger.info("SpendLogCleanup initialized with batch size: %d", self.batch_size)
|
||||
|
||||
def _should_delete_spend_logs(self) -> bool:
|
||||
"""
|
||||
Determines if logs should be deleted based on the max retention period in settings.
|
||||
"""
|
||||
retention_setting = self.general_settings.get("maximum_spend_logs_retention_period")
|
||||
verbose_proxy_logger.info("Checking retention setting: %s", retention_setting)
|
||||
|
||||
if retention_setting is None:
|
||||
verbose_proxy_logger.info("No retention setting found")
|
||||
return False
|
||||
|
||||
try:
|
||||
if isinstance(retention_setting, int):
|
||||
retention_setting = str(retention_setting)
|
||||
self.retention_seconds = duration_in_seconds(retention_setting)
|
||||
verbose_proxy_logger.info("Retention period set to %d seconds", self.retention_seconds)
|
||||
return True
|
||||
except ValueError as e:
|
||||
verbose_proxy_logger.error(
|
||||
f"Invalid maximum_spend_logs_retention_period value: {retention_setting}, error: {str(e)}"
|
||||
)
|
||||
return False
|
||||
|
||||
async def cleanup_old_spend_logs(self, prisma_client: PrismaClient) -> None:
|
||||
"""
|
||||
Main cleanup function. Deletes old spend logs in batches.
|
||||
"""
|
||||
try:
|
||||
verbose_proxy_logger.info(f"Cleanup job triggered at {datetime.now()}")
|
||||
|
||||
if not self._should_delete_spend_logs():
|
||||
verbose_proxy_logger.info("Skipping cleanup — invalid or missing retention setting.")
|
||||
return
|
||||
|
||||
if self.retention_seconds is None:
|
||||
verbose_proxy_logger.error("Retention seconds is None, cannot proceed with cleanup")
|
||||
return
|
||||
|
||||
cutoff_date = datetime.now(UTC) - timedelta(seconds=float(self.retention_seconds))
|
||||
verbose_proxy_logger.info(f"Deleting logs older than {cutoff_date.isoformat()}")
|
||||
|
||||
total_deleted = 0
|
||||
run_count = 0
|
||||
while True:
|
||||
if run_count > 100:
|
||||
verbose_proxy_logger.info("Max logs deleted - 1,00,000, rest of the logs 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:
|
||||
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)
|
||||
verbose_proxy_logger.info(f"Deleted {len(logs_to_delete)} logs in this batch")
|
||||
run_count += 1
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"Error during cleanup: {str(e)}")
|
||||
|
|
|
|||
|
|
@ -143,6 +143,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
|||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.proxy._experimental.mcp_server.server import router as mcp_router
|
||||
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
|
||||
from litellm.proxy._experimental.mcp_server.tool_registry import (
|
||||
global_mcp_tool_registry,
|
||||
)
|
||||
|
|
@ -3297,6 +3298,23 @@ class ProxyStartupEvent:
|
|||
|
||||
PrometheusLogger.initialize_budget_metrics_cron_job(scheduler=scheduler)
|
||||
|
||||
### SPEND LOG CLEANUP ###
|
||||
if general_settings.get("maximum_spend_logs_retention_period") is not None:
|
||||
spend_log_cleanup = SpendLogCleanup()
|
||||
# Get the interval from config or default to 1 day
|
||||
retention_interval = general_settings.get("maximum_spend_logs_retention_interval", "1d")
|
||||
try:
|
||||
interval_seconds = duration_in_seconds(retention_interval)
|
||||
scheduler.add_job(
|
||||
spend_log_cleanup.cleanup_old_spend_logs,
|
||||
"interval",
|
||||
seconds=interval_seconds,
|
||||
args=[prisma_client],
|
||||
)
|
||||
except ValueError:
|
||||
verbose_proxy_logger.error(f"Invalid maximum_spend_logs_retention_interval value: {retention_interval}, defaulting to 60 seconds")
|
||||
interval_seconds = 60
|
||||
|
||||
scheduler.start()
|
||||
|
||||
@classmethod
|
||||
|
|
|
|||
101
tests/litellm/proxy/test_spend_log_cleanup.py
Normal file
101
tests/litellm/proxy/test_spend_log_cleanup.py
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
"""
|
||||
Test cases for spend log cleanup functionality
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from datetime import datetime, timedelta, UTC
|
||||
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
|
||||
from unittest.mock import MagicMock, AsyncMock
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_delete_spend_logs():
|
||||
# Test case 1: No retention set
|
||||
cleaner = SpendLogCleanup(general_settings={})
|
||||
assert cleaner._should_delete_spend_logs() is False
|
||||
|
||||
# Test case 2: Valid seconds string
|
||||
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "3600s"})
|
||||
assert cleaner._should_delete_spend_logs() is True
|
||||
|
||||
# Test case 3: Valid days string
|
||||
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"})
|
||||
assert cleaner._should_delete_spend_logs() is True
|
||||
|
||||
# Test case 4: Valid hours string
|
||||
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "24h"})
|
||||
assert cleaner._should_delete_spend_logs() is True
|
||||
|
||||
# Test case 5: Invalid format
|
||||
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "invalid"})
|
||||
assert cleaner._should_delete_spend_logs() is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_old_spend_logs_batch_deletion():
|
||||
from types import SimpleNamespace
|
||||
|
||||
# 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
|
||||
]
|
||||
|
||||
# Wire up mocks
|
||||
mock_db.litellm_spendlogs = mock_spendlogs
|
||||
mock_prisma_client.db = mock_db
|
||||
|
||||
# Run cleanup
|
||||
test_settings = {"maximum_spend_logs_retention_period": "7d"}
|
||||
cleaner = SpendLogCleanup(general_settings=test_settings)
|
||||
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.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_old_spend_logs_retention_period_cutoff():
|
||||
"""
|
||||
Test that logs are filtered using correct cutoff based on retention
|
||||
"""
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_spendlogs.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma_client.db.litellm_spendlogs.delete = AsyncMock()
|
||||
|
||||
test_settings = {"maximum_spend_logs_retention_period": "24h"}
|
||||
cleaner = SpendLogCleanup(general_settings=test_settings)
|
||||
assert cleaner._should_delete_spend_logs() is True
|
||||
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
||||
|
||||
cutoff_date = mock_prisma_client.db.litellm_spendlogs.find_many.call_args[1]["where"]["startTime"]["lt"]
|
||||
expected_cutoff = datetime.now(UTC) - timedelta(hours=24)
|
||||
assert abs((cutoff_date - expected_cutoff).total_seconds()) < 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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()
|
||||
|
||||
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()
|
||||
Loading…
Add table
Reference in a new issue