From 44c100d05d2e41190111d818074adeeed661a725 Mon Sep 17 00:00:00 2001 From: "Jugal D. Bhatt" <55304795+jugaldb@users.noreply.github.com> Date: Fri, 16 May 2025 20:18:51 -0500 Subject: [PATCH] [Feat] handle multi-pod deployment for SpendLogs Retention (#10895) * handle multi-pod deployment * fix utc and ruff errors * add constants.py * add lock duration acc to interval * add lock duration on pod lock manager * update tests to use redis * update comments from review * update config_Settings.md * lint errors * remove custom ttl setting * add constants.py * add constants.py * add check for pod lock manager, and allow otherwise * remove dup try except and move to finally remove args --- docs/my-website/docs/proxy/config_settings.md | 3 +- litellm/constants.py | 4 +- .../db_transaction_queue/pod_lock_manager.py | 3 + .../db_transaction_queue/spend_log_cleanup.py | 105 ++++++++++++------ litellm/proxy/proxy_server.py | 3 +- tests/litellm/proxy/test_spend_log_cleanup.py | 48 ++++++-- 6 files changed, 120 insertions(+), 46 deletions(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 35efdb246da..3893e42633e 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -95,6 +95,7 @@ general_settings: key_management_system: google_kms # either google_kms or azure_kms master_key: string maximum_spend_logs_retention_period: 30d # The maximum time to retain spend logs before deletion. + maximum_spend_logs_retention_interval: 1d # interval in which the spend log cleanup task should run in. # Database Settings database_url: string @@ -212,7 +213,7 @@ general_settings: | forward_openai_org_id | boolean | If true, forwards the OpenAI Organization ID to the backend LLM call (if it's OpenAI). | | forward_client_headers_to_llm_api | boolean | If true, forwards the client headers (any `x-` headers) to the backend LLM call | | maximum_spend_logs_retention_period | str | Used to set the max retention time for spend logs in the db, after which they will be auto-purged | - +| maximum_spend_logs_retention_interval | str | Used to set the interval in which the spend log cleanup task should run in. | ### router_settings - Reference :::info diff --git a/litellm/constants.py b/litellm/constants.py index 4a5a00705f6..bedf54561d2 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -616,7 +616,9 @@ LITELLM_PROXY_ADMIN_NAME = "default_user_id" ########################### DB CRON JOB NAMES ########################### DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job" -PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME = "prometheus_emit_budget_metrics_job" +PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME = "prometheus_emit_budget_metrics" +SPEND_LOG_CLEANUP_JOB_NAME = "spend_log_cleanup" +SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 100)) DEFAULT_CRON_JOB_LOCK_TTL_SECONDS = int( os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60) ) # 1 minute diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index be3be645469..2dab963b02a 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -36,6 +36,9 @@ class PodLockManager: """ Attempt to acquire the lock for a specific cron job using Redis. Uses the SET command with NX and EX options to ensure atomicity. + + Args: + cronjob_id: The ID of the cron job to lock """ if self.redis_cache is None: verbose_proxy_logger.debug("redis_cache is None, skipping acquire_lock") 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 51297000269..ed76e5eddf5 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -1,29 +1,37 @@ -from datetime import datetime, timedelta, UTC +from datetime import datetime, timedelta, timezone from typing import Optional from litellm.proxy.utils import PrismaClient from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm._logging import verbose_proxy_logger +from litellm.caching import RedisCache +from litellm.constants import SPEND_LOG_CLEANUP_JOB_NAME, SPEND_LOG_RUN_LOOPS +import asyncio class SpendLogCleanup: """ Handles cleaning up old spend logs based on maximum retention period. Deletes logs in batches to prevent timeouts. + Uses PodLockManager to ensure only one pod runs cleanup in multi-pod deployments. """ - def __init__(self, general_settings=None): + def __init__(self, general_settings=None, redis_cache: Optional[RedisCache] = 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) + from litellm.proxy.proxy_server import proxy_logging_obj + + pod_lock_manager = proxy_logging_obj.db_spend_update_writer.pod_lock_manager + self.pod_lock_manager = pod_lock_manager + verbose_proxy_logger.info(f"SpendLogCleanup initialized with batch size: {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) + verbose_proxy_logger.info(f"Checking retention setting: {retention_setting}") if retention_setting is None: verbose_proxy_logger.info("No retention setting found") @@ -33,7 +41,7 @@ class SpendLogCleanup: 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) + verbose_proxy_logger.info(f"Retention period set to {self.retention_seconds} seconds") return True except ValueError as e: verbose_proxy_logger.error( @@ -41,9 +49,48 @@ class SpendLogCleanup: ) return False + async def _delete_old_logs(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int: + """ + Helper method to delete old logs in batches. + Returns the total number of logs deleted. + """ + total_deleted = 0 + run_count = 0 + while True: + 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") + 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) + run_count += 1 + + # Add a small sleep to prevent overwhelming the database + await asyncio.sleep(0.1) + + return total_deleted + async def cleanup_old_spend_logs(self, prisma_client: PrismaClient) -> None: """ Main cleanup function. Deletes old spend logs in batches. + If pod_lock_manager is available, ensures only one pod runs cleanup. + If no pod_lock_manager, runs cleanup without distributed locking. """ try: verbose_proxy_logger.info(f"Cleanup job triggered at {datetime.now()}") @@ -56,35 +103,29 @@ class SpendLogCleanup: verbose_proxy_logger.error("Retention seconds is None, cannot proceed with cleanup") return - cutoff_date = datetime.now(UTC) - timedelta(seconds=float(self.retention_seconds)) + # If we have a pod lock manager, try to acquire the lock + if self.pod_lock_manager and self.pod_lock_manager.redis_cache: + lock_acquired = await self.pod_lock_manager.acquire_lock( + cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME, + ) + verbose_proxy_logger.info(f"Lock acquisition attempt: {'successful' if lock_acquired else 'failed'} at {datetime.now()}") + + if not lock_acquired: + verbose_proxy_logger.info("Another pod is already running cleanup") + return + + cutoff_date = datetime.now(timezone.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") + # Perform the actual deletion + total_deleted = await self._delete_old_logs(prisma_client, cutoff_date) + verbose_proxy_logger.info(f"Deleted {total_deleted} logs") - 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)}") + return # Return after error handling + finally: + # Always release the lock if we have a pod lock manager + if self.pod_lock_manager and self.pod_lock_manager.redis_cache: + await self.pod_lock_manager.release_lock(cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME) + verbose_proxy_logger.info("Released cleanup lock") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4f5d7edd888..4b81dc1fab1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3312,8 +3312,7 @@ class ProxyStartupEvent: 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 + verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value") scheduler.start() diff --git a/tests/litellm/proxy/test_spend_log_cleanup.py b/tests/litellm/proxy/test_spend_log_cleanup.py index ffd506d0bf0..0b97f7c2c86 100644 --- a/tests/litellm/proxy/test_spend_log_cleanup.py +++ b/tests/litellm/proxy/test_spend_log_cleanup.py @@ -3,7 +3,7 @@ Test cases for spend log cleanup functionality """ import pytest -from datetime import datetime, timedelta, UTC +from datetime import datetime, timedelta, UTC, timezone from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup from unittest.mock import MagicMock, AsyncMock @@ -33,6 +33,7 @@ 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 MagicMock, AsyncMock, patch # Setup Prisma client mock_prisma_client = MagicMock() @@ -55,35 +56,62 @@ async def test_cleanup_old_spend_logs_batch_deletion(): mock_db.litellm_spendlogs = mock_spendlogs mock_prisma_client.db = mock_db - # Run cleanup + # Mock Redis cache and pod_lock_manager + mock_redis_cache = MagicMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = mock_redis_cache + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + # Run cleanup with mocked pod_lock_manager test_settings = {"maximum_spend_logs_retention_period": "7d"} cleaner = SpendLogCleanup(general_settings=test_settings) + cleaner.pod_lock_manager = mock_pod_lock_manager 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 - + 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)]}} + ) @pytest.mark.asyncio async def test_cleanup_old_spend_logs_retention_period_cutoff(): """ Test that logs are filtered using correct cutoff based on retention """ + # Setup Prisma client mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_spendlogs.find_many = AsyncMock(return_value=[]) - mock_prisma_client.db.litellm_spendlogs.delete = AsyncMock() + mock_db = MagicMock() + mock_spendlogs = MagicMock() + mock_spendlogs.find_many = AsyncMock(return_value=[]) + mock_spendlogs.delete_many = AsyncMock() + mock_db.litellm_spendlogs = mock_spendlogs + mock_prisma_client.db = mock_db + # Mock Redis cache and pod_lock_manager + mock_redis_cache = MagicMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = mock_redis_cache + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + # Run cleanup with mocked pod_lock_manager test_settings = {"maximum_spend_logs_retention_period": "24h"} cleaner = SpendLogCleanup(general_settings=test_settings) + cleaner.pod_lock_manager = mock_pod_lock_manager 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 - + # Verify the cutoff date is correct + cutoff_date = mock_spendlogs.find_many.call_args[1]["where"]["startTime"]["lt"] + expected_cutoff = datetime.now(timezone.utc) - timedelta(seconds=86400) + assert abs((cutoff_date - expected_cutoff).total_seconds()) < 1 # Allow 1 second difference for test execution time @pytest.mark.asyncio async def test_cleanup_old_spend_logs_no_retention_period():