From d5d516677762139e3cdbf516bb12264e210d1b49 Mon Sep 17 00:00:00 2001 From: "Jugal D. Bhatt" <55304795+jugaldb@users.noreply.github.com> Date: Fri, 16 May 2025 00:43:39 -0500 Subject: [PATCH] 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 --- .../db_transaction_queue/spend_log_cleanup.py | 117 +++++++++++++----- litellm/proxy/proxy_server.py | 18 +++ tests/litellm/proxy/test_spend_log_cleanup.py | 101 +++++++++++++++ 3 files changed, 203 insertions(+), 33 deletions(-) create mode 100644 tests/litellm/proxy/test_spend_log_cleanup.py 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 15d8c7b2e4a..51297000269 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -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 \ No newline at end of file + 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)}") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1825e360041..4f5d7edd888 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 diff --git a/tests/litellm/proxy/test_spend_log_cleanup.py b/tests/litellm/proxy/test_spend_log_cleanup.py new file mode 100644 index 00000000000..ffd506d0bf0 --- /dev/null +++ b/tests/litellm/proxy/test_spend_log_cleanup.py @@ -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() \ No newline at end of file