From b1f577199a108d7ee9b3ee5616d9baab5059fa32 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 6 May 2026 11:39:15 -0700 Subject: [PATCH] fix(proxy): keep spend log cleanup running after batch failures and surface DB errors (#27303) Co-authored-by: Yassin Kortam --- litellm/constants.py | 6 + .../db_transaction_queue/spend_log_cleanup.py | 65 +++++-- .../proxy/test_spend_log_cleanup.py | 166 ++++++++++++++++++ 3 files changed, 225 insertions(+), 12 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 6918e40cad..1e96f8eac9 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1457,6 +1457,12 @@ KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job" EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME = "litellm_expired_ui_session_key_cleanup_job" 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)) +SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int( + os.getenv("SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) +) +SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS = float( + os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5) +) SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) SPEND_LOG_QUEUE_POLL_INTERVAL = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0)) SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = int( 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 bc9efb52b0..9475779cfd 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -5,8 +5,10 @@ from typing import Optional from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache from litellm.constants import ( + SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS, SPEND_LOG_CLEANUP_BATCH_SIZE, SPEND_LOG_CLEANUP_JOB_NAME, + SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES, SPEND_LOG_RUN_LOOPS, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds @@ -74,6 +76,7 @@ class SpendLogCleanup: """ total_deleted = 0 run_count = 0 + consecutive_failures = 0 while True: if run_count > SPEND_LOG_RUN_LOOPS: verbose_proxy_logger.info( @@ -82,18 +85,50 @@ class SpendLogCleanup: break # Step 1: Find logs and delete them in one go without fetching to application # Delete in batches, limited by self.batch_size - deleted_result = await prisma_client.db.execute_raw( - """ - DELETE FROM "LiteLLM_SpendLogs" - WHERE "request_id" IN ( - SELECT "request_id" FROM "LiteLLM_SpendLogs" - WHERE "startTime" < $1::timestamptz - LIMIT $2 + try: + deleted_result = await prisma_client.db.execute_raw( + """ + DELETE FROM "LiteLLM_SpendLogs" + WHERE "request_id" IN ( + SELECT "request_id" FROM "LiteLLM_SpendLogs" + WHERE "startTime" < $1::timestamptz + LIMIT $2 + ) + """, + cutoff_date, + self.batch_size, ) - """, - cutoff_date, - self.batch_size, - ) + except Exception as batch_exc: + # A single batch failure (e.g. Prisma/DB timeout) must not abort + # the whole run — subsequent batches may still succeed. + consecutive_failures += 1 + verbose_proxy_logger.exception( + "Spend log cleanup batch failed " + "(run_count=%d, consecutive_failures=%d, batch_size=%d, " + "cutoff=%s, total_deleted_so_far=%d): %s: %s", + run_count, + consecutive_failures, + self.batch_size, + cutoff_date.isoformat(), + total_deleted, + type(batch_exc).__name__, + batch_exc, + ) + if ( + consecutive_failures + >= SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES + ): + verbose_proxy_logger.error( + "Aborting spend log cleanup after %d consecutive batch " + "failures; total deleted before abort: %d", + consecutive_failures, + total_deleted, + ) + break + await asyncio.sleep(SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS) + continue + + consecutive_failures = 0 deleted_count = 0 if isinstance(deleted_result, int): @@ -168,7 +203,13 @@ class SpendLogCleanup: verbose_proxy_logger.info(f"Deleted {total_deleted} logs") except Exception as e: - verbose_proxy_logger.error(f"Error during cleanup: {str(e)}") + # .exception() captures the traceback; str(e) alone on a Prisma/DB + # timeout is often empty and gives operators no signal to diagnose. + verbose_proxy_logger.exception( + "Error during spend log cleanup: %s: %s", + type(e).__name__, + e, + ) return # Return after error handling finally: # Only release the lock if it was actually acquired diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index 4923d70a43..42bb919295 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -331,6 +331,172 @@ async def test_delete_old_logs_continues_on_valid_int_return(): assert total_deleted == 800 +@pytest.mark.asyncio +async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch): + """A single batch failure (e.g. DB timeout) must not abort the whole run — + subsequent batches should still execute and their counts accumulate.""" + import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module + + # Zero out the failure backoff so the test doesn't take ~0.5s of real sleep. + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 + ) + + mock_prisma_client = MagicMock() + mock_db = MagicMock() + # batch 1 succeeds, batch 2 raises (one-off DB timeout), batches 3-4 succeed, + # batch 5 returns 0 → loop exits naturally. + mock_db.execute_raw = AsyncMock( + side_effect=[100, TimeoutError("simulated DB timeout"), 200, 50, 0] + ) + mock_prisma_client.db = mock_db + + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) + + cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) + total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date) + + # All 5 batches should have been attempted; 100 + 200 + 50 = 350 deleted. + assert mock_db.execute_raw.call_count == 5 + assert total_deleted == 350 + + +@pytest.mark.asyncio +async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch): + """If batch failures persist for SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES + in a row (e.g. DB is down), the loop must abort instead of hot-looping.""" + import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module + + # Lower the threshold so the test is fast and deterministic. + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 + ) + + mock_prisma_client = MagicMock() + mock_db = MagicMock() + # Every batch raises — must abort after exactly 3 attempts, not loop forever. + mock_db.execute_raw = AsyncMock( + side_effect=ConnectionError("simulated persistent DB outage") + ) + mock_prisma_client.db = mock_db + + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) + + cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) + total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date) + + assert mock_db.execute_raw.call_count == 3 + assert total_deleted == 0 + + +@pytest.mark.asyncio +async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatch): + """A success between failures must reset the consecutive-failure counter so + intermittent timeouts don't trip the abort threshold.""" + import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module + + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 + ) + + mock_prisma_client = MagicMock() + mock_db = MagicMock() + # Pattern: fail, fail, success (resets counter), fail, fail, success, done. + # Without reset, three of these would trip abort; with reset, they don't. + mock_db.execute_raw = AsyncMock( + side_effect=[ + TimeoutError("t1"), + TimeoutError("t2"), + 100, + TimeoutError("t3"), + TimeoutError("t4"), + 50, + 0, + ] + ) + mock_prisma_client.db = mock_db + + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) + + cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) + total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date) + + assert mock_db.execute_raw.call_count == 7 + assert total_deleted == 150 + + +@pytest.mark.asyncio +async def test_cleanup_uses_logger_exception_for_full_traceback(monkeypatch): + """The outer error handler must call logger.exception() (not .error(str(e))) + so Prisma/DB timeouts surface a full traceback and exception type.""" + import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module + + mock_logger = MagicMock() + monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger) + + mock_prisma_client = MagicMock() + # Force the outer try/except to fire by making _should_delete_spend_logs raise. + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) + cleaner.pod_lock_manager = None + + def boom(): + raise RuntimeError("simulated prisma timeout") + + cleaner._should_delete_spend_logs = boom # type: ignore[assignment] + + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + assert mock_logger.exception.called, "expected logger.exception() to be called" + # The exception type name must appear in the formatted args so operators can + # tell *what* failed, not just "Error during cleanup:". + call_args = mock_logger.exception.call_args + formatted = call_args[0][0] % call_args[0][1:] + assert "RuntimeError" in formatted + assert "simulated prisma timeout" in formatted + + +@pytest.mark.asyncio +async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch): + """Even when batch deletion aborts due to consecutive failures, the pod lock + must still be released so the next scheduled run isn't permanently blocked.""" + import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module + + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 + ) + + mock_prisma_client = MagicMock() + mock_db = MagicMock() + mock_db.execute_raw = AsyncMock(side_effect=TimeoutError("DB down")) + mock_prisma_client.db = mock_db + + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) + cleaner.pod_lock_manager = mock_pod_lock_manager + + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + # Cleanup didn't crash; the abort-after-failures path returned cleanly. + mock_pod_lock_manager.release_lock.assert_awaited_once() + + def test_cleanup_batch_size_env_var(monkeypatch): """Ensure batch size is configurable via environment variable""" import importlib