scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

[Feat] handle multi-pod deployment for SpendLogs Retention (#10895)

Browse source 
* 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
This commit is contained in:
Jugal D. Bhatt 2025-05-16 20:18:51 -05:00 • committed by GitHub
parent fdacf07ccb
commit 44c100d05d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 120 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
docs/my-website/docs/proxy/config_settings.md
View file

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

4
litellm/constants.py
View file

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

3
litellm/proxy/db/db_transaction_queue/pod_lock_manager.py
View file

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

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

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

3
litellm/proxy/proxy_server.py
View file

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

48
tests/litellm/proxy/test_spend_log_cleanup.py
View file

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