From 06d56eb6d7bd835b4dcae468c0bee67dedaef0f7 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 11:37:04 -0700 Subject: [PATCH] fix(proxy): one spend log cleanup pass at a time per pod, with or without Redis (#47500) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../db_transaction_queue/spend_log_cleanup.py | 14 ++++++++ tests/unit/proxy/test_spend_log_cleanup.py | 33 ++++++++++++------- 2 files changed, 35 insertions(+), 12 deletions(-) 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 726bc2b221c..e1e1d3690ff 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -354,6 +354,9 @@ _STORAGE_SQL: Final = """ ) """ +# The age job and the size job are scheduled on separate SpendLogCleanup instances, so the lock is per process. +_PASS_LOCK: Final = asyncio.Lock() + SPEND_LOG_CLEANUP_BOUND_SETTINGS: Final = ( "maximum_spend_logs_cleanup_batch_size", "maximum_spend_logs_cleanup_max_batches", @@ -1363,6 +1366,17 @@ class SpendLogCleanup: await self._run(prisma_client, age_rules=False, size_rule=True) async def _run(self, prisma_client: PrismaClient, age_rules: bool, size_rule: bool) -> None: + """ + One pass at a time per pod. The age job and the size job fire from separate schedules, and the pod + lock only exists with Redis, so this is what keeps them apart in a pod without it. The later one waits + rather than skips, because the other job applies a different rule and would not do its work for it. + """ + if _PASS_LOCK.locked(): + verbose_proxy_logger.info("Another cleanup pass is running in this pod; waiting for it to finish") + async with _PASS_LOCK: + await self._run_passes(prisma_client, age_rules, size_rule) + + async def _run_passes(self, prisma_client: PrismaClient, age_rules: bool, size_rule: bool) -> None: lock_acquired = False run_started_at: Final = time.monotonic() progress: Final = _RunProgress() diff --git a/tests/unit/proxy/test_spend_log_cleanup.py b/tests/unit/proxy/test_spend_log_cleanup.py index 8c97f829025..a0f222ed3e0 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/tests/unit/proxy/test_spend_log_cleanup.py @@ -20,13 +20,13 @@ from litellm.constants import ( SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS, ) from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( + SPEND_LOG_CLEANUP_BOUND_SETTINGS, SizeCleanupState, + SpendLogCleanup, SpendLogStorage, StorageProbe, - advance_size_state, - SPEND_LOG_CLEANUP_BOUND_SETTINGS, - SpendLogCleanup, TableCleanupResult, + advance_size_state, ) from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import ( SpendLogCleanupMetrics, @@ -1539,8 +1539,8 @@ async def test_progress_reported_for_a_cancelled_run_is_that_run_only(monkeypatc @pytest.mark.asyncio -async def test_progress_reported_by_an_overlapping_run_is_its_own(monkeypatch): - """With APSCHEDULER_MAX_INSTANCES above one, two runs share the cleaner but not their progress""" +async def test_a_size_pass_that_fires_while_an_age_pass_is_running_in_the_same_pod_waits_for_it(monkeypatch): + """Without Redis there is no pod lock, and the two jobs are scheduled on two SpendLogCleanup instances""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module mock_logger = MagicMock() @@ -1557,22 +1557,31 @@ async def test_progress_reported_by_an_overlapping_run_is_its_own(monkeypatch): slow_client = MagicMock() _wire_tx(slow_client.db) slow_client.db.execute_raw = _slow_execute_raw - fast_client = MagicMock() - _wire_tx(fast_client.db) - fast_client.db.execute_raw = AsyncMock(side_effect=[150, 150, 0, 0]) + size_db = MagicMock() + _wire_tx(size_db) + size_db.execute_raw = AsyncMock(return_value=1000) + _wire_size_probes(size_db, [_storage(allocated=10 * 1024**3, live=1000, dead=0)], cutoff=datetime.now(timezone.utc)) + size_client = MagicMock() + size_client.db = size_db - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) - cleaner.pod_lock_manager = None + age_cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + age_cleaner.pod_lock_manager = None + size_cleaner = _size_cleaner({}) - slow_run = asyncio.ensure_future(cleaner.cleanup_old_spend_logs(slow_client)) + slow_run = asyncio.ensure_future(age_cleaner.cleanup_old_spend_logs(slow_client)) await asyncio.wait_for(first_batch_done.wait(), timeout=5) - await cleaner.cleanup_old_spend_logs(fast_client) + size_run = asyncio.ensure_future(size_cleaner.cleanup_spend_logs_by_size(size_client)) + await asyncio.sleep(0.05) + assert size_db.deletes.await_count == 0 and not size_run.done(), "the size pass ran while the age pass held the pod" + second_run_done.set() await asyncio.sleep(0) slow_run.cancel() with pytest.raises(asyncio.CancelledError): await slow_run + await asyncio.wait_for(size_run, timeout=5) + assert size_db.deletes.await_count > 0, "the size pass never ran once the age pass let go" (error_call,) = mock_logger.error.call_args_list rendered = error_call[0][0] % error_call[0][1:] assert "(rows_deleted=100, batches=1)" in rendered