tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 08:23:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Feat] Make batch size for maximum retention in spend logs a controllable parameter (#11459)

Browse Source 
* feat: add SPEND_LOG_CLEANUP_BATCH_SIZE

* docs update

* test: test_cleanup_batch_size_env_var
This commit is contained in:
Ishaan Jaff
2025-06-05 17:11:51 -07:00
committed by GitHub
parent d05eda0311
commit 18ea65218b
6 changed files with 81 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/proxy/config_settings.md
+2 -1
View File
@@ -650,4 +650,5 @@ router_settings:
| USE_AWS_KMS | Flag to enable AWS Key Management Service for encryption
| USE_PRISMA_MIGRATE | Flag to use prisma migrate instead of prisma db push. Recommended for production environments.
| WEBHOOK_URL | URL for receiving webhooks from external services
| SPEND_LOG_RUN_LOOPS | Constant for setting how many runs of 1000 batch deletes should spend_log_cleanup task run
| SPEND_LOG_RUN_LOOPS | Constant for setting how many runs of 1000 batch deletes should spend_log_cleanup task run |
| SPEND_LOG_CLEANUP_BATCH_SIZE | Number of logs deleted per batch during cleanup. Default is 1000 |
docs/my-website/docs/proxy/spend_logs_deletion.md
+5 -3
View File
@@ -71,18 +71,20 @@ If Redis is enabled, LiteLLM uses it to make sure only one instance runs the cle
Once cleanup starts:
- It calculates the cutoff date using the configured retention period
- Deletes logs older than the cutoff in **batches of 1000**
- Deletes logs older than the cutoff in batches (default size `1000`)
- Adds a short delay between batches to avoid overloading the database
### Default settings:
- **Batch size**: 1000 logs
- **Batch size**: 1000 logs (configurable via `SPEND_LOG_CLEANUP_BATCH_SIZE`)
- **Max batches per run**: 500
- **Max deletions per run**: 500,000 logs
You can change the number of batches using an environment variable:
You can change the cleanup parameters using environment variables:
```bash
SPEND_LOG_RUN_LOOPS=200
# optional: change batch size from the default 1000
SPEND_LOG_CLEANUP_BATCH_SIZE=2000
```
This would allow up to 200,000 logs to be deleted in one run.
docs/my-website/docs/proxy/ui_logs.md
+3 -1
View File
@@ -69,7 +69,9 @@ general_settings:
You can control how many logs are deleted per run using this environment variable:
`SPEND_LOG_RUN_LOOPS=200 # Deletes up to 200,000 logs in one run (batch size = 1000)`
`SPEND_LOG_RUN_LOOPS=200 # Deletes up to 200,000 logs in one run`
Set `SPEND_LOG_CLEANUP_BATCH_SIZE` to control how many logs are deleted per batch (default `1000`).
For detailed architecture and how it works, see [Spend Logs Deletion](../proxy/spend_logs_deletion).
litellm/constants.py
+1
View File
@@ -734,6 +734,7 @@ DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_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", 500))
SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000))
DEFAULT_CRON_JOB_LOCK_TTL_SECONDS = int(
os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)
) # 1 minute
litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
+49 -19
View File
@@ -1,11 +1,16 @@
import asyncio
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
from litellm.constants import (
SPEND_LOG_CLEANUP_BATCH_SIZE,
SPEND_LOG_CLEANUP_JOB_NAME,
SPEND_LOG_RUN_LOOPS,
)
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy.utils import PrismaClient
class SpendLogCleanup:
@@ -16,21 +21,26 @@ class SpendLogCleanup:
"""
def __init__(self, general_settings=None, redis_cache: Optional[RedisCache] = None):
self.batch_size = 1000
self.batch_size = SPEND_LOG_CLEANUP_BATCH_SIZE
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
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}")
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")
retention_setting = self.general_settings.get(
"maximum_spend_logs_retention_period"
)
verbose_proxy_logger.info(f"Checking retention setting: {retention_setting}")
if retention_setting is None:
@@ -41,7 +51,9 @@ class SpendLogCleanup:
if isinstance(retention_setting, int):
retention_setting = str(retention_setting)
self.retention_seconds = duration_in_seconds(retention_setting)
verbose_proxy_logger.info(f"Retention period set to {self.retention_seconds} seconds")
verbose_proxy_logger.info(
f"Retention period set to {self.retention_seconds} seconds"
)
return True
except ValueError as e:
verbose_proxy_logger.error(
@@ -49,7 +61,9 @@ class SpendLogCleanup:
)
return False
async def _delete_old_logs(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int:
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.
@@ -58,7 +72,9 @@ class SpendLogCleanup:
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")
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(
@@ -68,7 +84,9 @@ class SpendLogCleanup:
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}")
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]
@@ -80,7 +98,7 @@ class SpendLogCleanup:
total_deleted += len(logs_to_delete)
run_count += 1
# Add a small sleep to prevent overwhelming the database
await asyncio.sleep(0.1)
@@ -96,11 +114,15 @@ class SpendLogCleanup:
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.")
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")
verbose_proxy_logger.error(
"Retention seconds is None, cannot proceed with cleanup"
)
return
# If we have a pod lock manager, try to acquire the lock
@@ -108,14 +130,20 @@ class SpendLogCleanup:
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()}")
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()}")
cutoff_date = datetime.now(timezone.utc) - timedelta(
seconds=float(self.retention_seconds)
)
verbose_proxy_logger.info(
f"Deleting logs older than {cutoff_date.isoformat()}"
)
# Perform the actual deletion
total_deleted = await self._delete_old_logs(prisma_client, cutoff_date)
@@ -127,5 +155,7 @@ class SpendLogCleanup:
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)
await self.pod_lock_manager.release_lock(
cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME
)
verbose_proxy_logger.info("Released cleanup lock")
tests/test_litellm/proxy/test_spend_log_cleanup.py
+21
View File
@@ -142,3 +142,24 @@ async def test_cleanup_old_spend_logs_no_retention_period():
mock_prisma_client.db.litellm_spendlogs.find_many.assert_not_called()
mock_prisma_client.db.litellm_spendlogs.delete.assert_not_called()
def test_cleanup_batch_size_env_var(monkeypatch):
"""Ensure batch size is configurable via environment variable"""
import importlib
import litellm.constants as constants_module
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
# Set env var and reload modules to pick up new value
monkeypatch.setenv("SPEND_LOG_CLEANUP_BATCH_SIZE", "25")
importlib.reload(constants_module)
importlib.reload(cleanup_module)
cleaner = cleanup_module.SpendLogCleanup(general_settings={})
assert cleaner.batch_size == 25
# Remove env var and reload to restore default for other tests
monkeypatch.delenv("SPEND_LOG_CLEANUP_BATCH_SIZE", raising=False)
importlib.reload(constants_module)
importlib.reload(cleanup_module)