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 e97e9f6e683..3068abbc1f4 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -96,6 +96,8 @@ class SpendLogCleanup: self.general_settings = general_settings or default_settings self._refresh_bounds() + self._run_rows_deleted: int = 0 + self._run_batches: int = 0 from litellm.proxy.proxy_server import proxy_logging_obj pod_lock_manager: Final = proxy_logging_obj.db_spend_update_writer.pod_lock_manager @@ -422,6 +424,8 @@ class SpendLogCleanup: total_deleted += deleted_count run_count += 1 + self._run_rows_deleted += deleted_count + self._run_batches += 1 # Add a small sleep to prevent overwhelming the database await asyncio.sleep(0.1) @@ -590,6 +594,9 @@ class SpendLogCleanup: If no pod_lock_manager, runs cleanup without distributed locking. """ lock_acquired = False + run_start: Final = time.monotonic() + self._run_rows_deleted = 0 + self._run_batches = 0 try: verbose_proxy_logger.info("Cleanup job triggered at %s", datetime.now()) self._refresh_bounds() @@ -670,7 +677,16 @@ class SpendLogCleanup: self._run_outcome(spend_log_results + session_results + health_check_results) ) - except Exception as e: + except asyncio.CancelledError: + verbose_proxy_logger.exception( + "Spend log cleanup cancelled (elapsed=%.2fs, rows_deleted=%d, batches=%d)", + time.monotonic() - run_start, + self._run_rows_deleted, + self._run_batches, + ) + SpendLogCleanupMetrics.record_run("aborted") + raise + except Exception as e: # noqa: BLE001 - top-level cleanup job exception handler # .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( diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index bf1538183ab..9175fc6a914 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -1417,3 +1417,88 @@ def test_the_reported_run_outcome_is_the_most_significant_reason_in_any_order(st """ results = tuple(TableCleanupResult(rows_deleted=0, stop_reason=reason) for reason in stop_reasons) assert SpendLogCleanup._run_outcome(results) == expected + + +@pytest.mark.asyncio +async def test_cancelled_cleanup_records_and_logs_aborted_run(monkeypatch): + """ + When cleanup_old_spend_logs is cancelled (e.g. during pod shutdown), + asyncio.CancelledError must record outcome='aborted', emit an error-level + diagnostic containing progress metrics (including rows deleted and batches + completed prior to cancellation), and re-raise CancelledError. + """ + import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module + + mock_db = MagicMock() + _wire_tx(mock_db) + # First batch deletes 150 rows. Second batch simulates cancellation mid-run. + mock_db.execute_raw = AsyncMock( + side_effect=[150, asyncio.CancelledError("simulated shutdown mid-run")] + ) + mock_prisma_client = MagicMock() + mock_prisma_client.db = mock_db + + mock_lock_manager = MagicMock() + mock_lock_manager.redis_cache = MagicMock() + mock_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_lock_manager.release_lock = AsyncMock() + + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) + cleaner.pod_lock_manager = mock_lock_manager + + logger = MagicMock() + recorded_runs: list[str] = [] + + monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", logger) + monkeypatch.setattr( + cleanup_module.SpendLogCleanupMetrics, "record_run", lambda outcome: recorded_runs.append(outcome) + ) + + with pytest.raises(asyncio.CancelledError): + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + assert recorded_runs == ["aborted"] + assert cleaner._run_rows_deleted == 150 + assert cleaner._run_batches == 1 + mock_lock_manager.release_lock.assert_awaited_once() + + logger.exception.assert_called_once() + log_call = logger.exception.call_args[0] + rendered_log = log_call[0] % log_call[1:] + assert "cancelled" in rendered_log.lower() + assert "elapsed=" in rendered_log + assert "rows_deleted=150" in rendered_log + assert "batches=1" in rendered_log + + +@pytest.mark.asyncio +async def test_cancelled_cleanup_releases_lock(): + """ + When a cleanup run holds a distributed pod lock and is cancelled, + the finally block must release the lock before propagating CancelledError. + """ + mock_db = MagicMock() + _wire_tx(mock_db) + mock_db.execute_raw = AsyncMock( + side_effect=asyncio.CancelledError("simulated immediate shutdown") + ) + mock_prisma_client = MagicMock() + mock_prisma_client.db = mock_db + + mock_lock_manager = MagicMock() + mock_lock_manager.redis_cache = MagicMock() + mock_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_lock_manager.release_lock = AsyncMock() + + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) + cleaner.pod_lock_manager = mock_lock_manager + + with pytest.raises(asyncio.CancelledError): + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + mock_lock_manager.release_lock.assert_awaited_once() +