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..1a14210dbec 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_started_at: 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,6 +677,15 @@ class SpendLogCleanup: self._run_outcome(spend_log_results + session_results + health_check_results) ) + except asyncio.CancelledError: + verbose_proxy_logger.error( + "Spend log cleanup cancelled after %.2fs (rows_deleted=%d, batches=%d); the next run resumes from here", + time.monotonic() - run_started_at, + self._run_rows_deleted, + self._run_batches, + ) + SpendLogCleanupMetrics.record_run("aborted") + raise except Exception as e: # .exception() captures the traceback; str(e) alone on a Prisma/DB # timeout is often empty and gives operators no signal to diagnose. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 343be1c98f4..c2dab669087 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -661,6 +661,10 @@ from litellm.proxy.route_llm_request import route_request from litellm.proxy.route_priority import hot_routes_first from litellm.proxy.search_endpoints.endpoints import router as search_router from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager +from litellm.proxy.shutdown.scheduled_jobs import ( + AwaitableAsyncIOExecutor, + cancel_in_flight_scheduler_jobs, +) from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start from litellm.proxy.spend_tracking.spend_counter_batch import ( PendingSpendIncrement, @@ -1413,6 +1417,13 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: await proxy_config.stop_auth_cache_invalidation_subscriber() + # Shutdown event - cancel and await in-flight scheduled jobs while the DB is still connected + if scheduler is not None and scheduler_executor is not None: + try: + await cancel_in_flight_scheduler_jobs(scheduler, scheduler_executor) + except Exception as e: + verbose_proxy_logger.error("Error cancelling in-flight scheduled jobs: %s", e) + await proxy_shutdown_event(worker_heartbeat=worker_heartbeat) if prometheus_multiproc_dir: @@ -2418,6 +2429,7 @@ celery_app_conn: Final = None celery_fn: Final = None # Redis Queue for handling requests scheduler = None +scheduler_executor: AwaitableAsyncIOExecutor | None = None # rebind-ok: bound once the scheduler is built at startup # Global variable for anthropic beta headers reload scheduling last_anthropic_beta_headers_reload = None @@ -9726,7 +9738,7 @@ class ProxyStartupEvent: proxy_logging_obj: ProxyLogging, ) -> ProxyWorkerHeartbeat: """Initializes scheduled background jobs""" - global heuristic_v1_tuning_baselines, store_model_in_db, scheduler # rebind-ok: startup publishes the one read-only baseline snapshot + global heuristic_v1_tuning_baselines, store_model_in_db, scheduler, scheduler_executor # rebind-ok: startup publishes the one read-only baseline snapshot # MEMORY LEAK FIX: Configure scheduler with optimized settings # Memray analysis showed APScheduler's normalize() and _apply_jitter() causing @@ -9735,9 +9747,9 @@ class ProxyStartupEvent: # 1. Remove/minimize jitter to avoid normalize() memory explosion # 2. Use larger misfire_grace_time to prevent backlog calculations # 3. Set replace_existing=True to avoid duplicate jobs - from apscheduler.executors.asyncio import AsyncIOExecutor from apscheduler.jobstores.memory import MemoryJobStore + scheduler_executor = AwaitableAsyncIOExecutor() # rebind-ok: shutdown awaits the jobs this executor runs scheduler = AsyncIOScheduler( job_defaults={ "coalesce": APSCHEDULER_COALESCE, @@ -9750,7 +9762,7 @@ class ProxyStartupEvent: jobstores={"default": MemoryJobStore()}, # explicitly use memory job store # Use simple executor to minimize overhead executors={ - "default": AsyncIOExecutor(), + "default": scheduler_executor, }, # Disable timezone awareness to reduce computation timezone=None, diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py new file mode 100644 index 00000000000..3c9e791f51c --- /dev/null +++ b/litellm/proxy/shutdown/scheduled_jobs.py @@ -0,0 +1,70 @@ +""" +Cancel the proxy's in-flight scheduled jobs at shutdown so they can record how they ended. + +APScheduler's ``AsyncIOExecutor.shutdown`` cancels the job tasks it has in flight but cannot +wait for them, because it is not a coroutine, and under uvicorn nothing else ever will: uvicorn +re-raises the SIGTERM it captured as soon as the ASGI lifespan shutdown returns, so the process +dies before ``asyncio.run`` reaches its cancel-all-tasks step. A job mid-run at that point is +killed without ever observing cancellation, which is how a spend-log cleanup interrupted by a +rolling restart left no outcome metric and no log line behind. Cancelling here and awaiting the +cancelled tasks while the database is still connected is what lets a job's own +``CancelledError`` handler run. + +The wait is bounded by ``JOB_CANCEL_TIMEOUT_SECONDS``. A job that has just been cancelled has only +its own cleanup left to do, so the bound is there for a job that swallows cancellation, not one +that honours it, and it keeps shutdown well inside a Kubernetes termination grace period. +""" + +# pyright: reportMissingTypeStubs=false # apscheduler ships no type information + +import asyncio +from collections.abc import Collection +from typing import Final, Protocol + +from apscheduler.executors.asyncio import AsyncIOExecutor + +from litellm._logging import verbose_proxy_logger + +JOB_CANCEL_TIMEOUT_SECONDS: Final = 5.0 + + +class StoppableScheduler(Protocol): + """The slice of ``AsyncIOScheduler`` shutdown uses, which ships no type information""" + + @property + def running(self) -> bool: ... + + def shutdown(self, wait: bool = ...) -> None: ... + + +class AwaitableAsyncIOExecutor(AsyncIOExecutor): # pyright: ignore[reportUntypedBaseClass] # apscheduler ships no type information and is absent from the type-check env + """``AsyncIOExecutor`` whose in-flight job tasks can be awaited after ``shutdown`` cancels them""" + + _pending_futures: Collection["asyncio.Future[object]"] + + def in_flight_jobs(self) -> tuple["asyncio.Future[object]", ...]: + """The job tasks that are running right now, as a snapshot""" + return tuple(future for future in self._pending_futures if not future.done()) + + +async def cancel_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: AwaitableAsyncIOExecutor) -> None: + """ + Stop the scheduler and wait, bounded by JOB_CANCEL_TIMEOUT_SECONDS, for the jobs it cancels. + + Must run before the database is disconnected: a job's cancellation handler is what records + the run's outcome, and it needs the connection the job was using. + """ + if not scheduler.running: + return + in_flight: Final = executor.in_flight_jobs() + scheduler.shutdown(wait=False) + if not in_flight: + return + verbose_proxy_logger.info("Cancelling %d in-flight scheduled job(s) for shutdown", len(in_flight)) + _done, pending = await asyncio.wait(in_flight, timeout=JOB_CANCEL_TIMEOUT_SECONDS) + if pending: + verbose_proxy_logger.warning( + "%d scheduled job(s) did not finish within %ss of cancellation; giving up on them", + len(pending), + JOB_CANCEL_TIMEOUT_SECONDS, + ) diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py new file mode 100644 index 00000000000..b77d7c4ae50 --- /dev/null +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -0,0 +1,125 @@ +""" +Tests for cancelling in-flight scheduled jobs at proxy shutdown. + +These drive a real AsyncIOScheduler: the point of the helper is the hand-off +between APScheduler's fire-and-forget cancellation and the lifespan shutdown +that has to outlive it, and a mocked scheduler would not exercise that. +""" + +import asyncio +import logging +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from datetime import datetime + +import pytest +from apscheduler.schedulers.asyncio import AsyncIOScheduler + +import litellm.proxy.shutdown.scheduled_jobs as scheduled_jobs +from litellm.proxy.shutdown.scheduled_jobs import ( + AwaitableAsyncIOExecutor, + cancel_in_flight_scheduler_jobs, +) + + +class _Job: + """A scheduled job that blocks until cancelled and records what it observed.""" + + def __init__(self, swallow_cancellation: bool = False) -> None: + self.started = asyncio.Event() + self.events: list[str] = [] + self.swallow_cancellation = swallow_cancellation + + async def run(self) -> None: + self.started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.events.append("cancelled") + if self.swallow_cancellation: + await asyncio.Event().wait() + raise + finally: + self.events.append("finished") + + +@asynccontextmanager +async def _running_scheduler(*jobs: _Job) -> AsyncIterator[tuple[AsyncIOScheduler, AwaitableAsyncIOExecutor]]: + """A started scheduler with every job in flight; stopped on the way out whatever the test did.""" + executor = AwaitableAsyncIOExecutor() + scheduler = AsyncIOScheduler(executors={"default": executor}) + for index, job in enumerate(jobs): + scheduler.add_job(job.run, id=f"job-{index}", next_run_time=datetime.now()) + scheduler.start() + try: + for job in jobs: + await asyncio.wait_for(job.started.wait(), timeout=5) + yield scheduler, executor + finally: + if scheduler.running: + scheduler.shutdown(wait=False) + stragglers = executor.in_flight_jobs() + for straggler in stragglers: + straggler.cancel() + await asyncio.gather(*stragglers, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_in_flight_jobs_observe_cancellation_before_shutdown_returns(): + """ + The job's own CancelledError handler is what records how a run ended, so + shutdown must not return until that handler has run. + """ + job = _Job() + async with _running_scheduler(job) as (scheduler, executor): + await cancel_in_flight_scheduler_jobs(scheduler, executor) + + assert job.events == ["cancelled", "finished"] + assert scheduler.running is False + assert executor.in_flight_jobs() == () + + +@pytest.mark.asyncio +async def test_every_in_flight_job_is_cancelled_not_only_the_first(): + first, second = _Job(), _Job() + async with _running_scheduler(first, second) as (scheduler, executor): + await cancel_in_flight_scheduler_jobs(scheduler, executor) + + assert first.events == ["cancelled", "finished"] + assert second.events == ["cancelled", "finished"] + + +@pytest.mark.asyncio +async def test_a_job_that_ignores_cancellation_is_abandoned_after_the_timeout(monkeypatch, caplog): + """ + A job that swallows CancelledError must not hold the pod past its + termination grace period, so shutdown gives up on it and says so. + """ + monkeypatch.setattr(scheduled_jobs, "JOB_CANCEL_TIMEOUT_SECONDS", 0.05) + job = _Job(swallow_cancellation=True) + async with _running_scheduler(job) as (scheduler, executor): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await cancel_in_flight_scheduler_jobs(scheduler, executor) + + assert job.events == ["cancelled"] + assert "1 scheduled job(s) did not finish within 0.05s of cancellation" in caplog.text + + +@pytest.mark.asyncio +async def test_shutdown_with_nothing_in_flight_still_stops_the_scheduler(): + async with _running_scheduler() as (scheduler, executor): + await cancel_in_flight_scheduler_jobs(scheduler, executor) + await asyncio.sleep(0) + + assert scheduler.running is False + + +@pytest.mark.asyncio +async def test_a_scheduler_that_never_started_is_left_alone(): + """The proxy runs without a scheduler when it has no database; shutdown must not trip on that.""" + executor = AwaitableAsyncIOExecutor() + scheduler = AsyncIOScheduler(executors={"default": executor}) + + await cancel_in_flight_scheduler_jobs(scheduler, executor) + + assert scheduler.running is False diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index bf1538183ab..ed35af7ee38 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -7,6 +7,7 @@ import math import time from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -1417,3 +1418,94 @@ 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 + + +_OTHER_OUTCOMES: Final = ("completed", "budget_exhausted", "batch_cap_reached", "skipped_locked", "skipped_disabled") + + +def _runs_recorded(outcome: str) -> float: + """The real ``litellm_spend_log_cleanup_runs_total`` sample for one outcome, 0 when unset""" + from prometheus_client import REGISTRY + + return REGISTRY.get_sample_value("litellm_spend_log_cleanup_runs_total", {"outcome": outcome}) or 0.0 + + +@pytest.mark.asyncio +async def test_a_cancelled_run_records_aborted_and_logs_its_progress_before_re_raising(monkeypatch): + """ + Shutdown cancels a run by throwing CancelledError into whichever batch is in + flight. That is a BaseException, so the Exception handler never saw it and + an interrupted run left no outcome metric and no log line; operators could + not tell that cleanup stopped early, let alone how far it got. + """ + 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) + aborted_runs_before = _runs_recorded("aborted") + other_runs_before = {outcome: _runs_recorded(outcome) for outcome in _OTHER_OUTCOMES} + + third_batch_reached = asyncio.Event() + + async def _execute_raw(sql, *args): + if third_batch_reached.is_set(): + raise AssertionError("no batch may be issued after the cancelled one") + if _execute_raw.calls < 2: + _execute_raw.calls += 1 + return 150 + third_batch_reached.set() + await asyncio.Event().wait() + + _execute_raw.calls = 0 + mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) + mock_prisma_client.db.execute_raw = _execute_raw + + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner.pod_lock_manager = MagicMock() + cleaner.pod_lock_manager.redis_cache = MagicMock() + cleaner.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + cleaner.pod_lock_manager.release_lock = AsyncMock() + + run = asyncio.ensure_future(cleaner.cleanup_old_spend_logs(mock_prisma_client)) + await asyncio.wait_for(third_batch_reached.wait(), timeout=5) + run.cancel() + with pytest.raises(asyncio.CancelledError): + await run + + assert _runs_recorded("aborted") == aborted_runs_before + 1 + assert {outcome: _runs_recorded(outcome) for outcome in _OTHER_OUTCOMES} == other_runs_before + cleaner.pod_lock_manager.release_lock.assert_awaited_once() + mock_logger.exception.assert_not_called() + (error_call,) = mock_logger.error.call_args_list + rendered = error_call[0][0] % error_call[0][1:] + assert rendered.startswith("Spend log cleanup cancelled after ") + assert "s (rows_deleted=300, batches=2)" in rendered + + +@pytest.mark.asyncio +async def test_progress_reported_for_a_cancelled_run_is_that_run_only(monkeypatch): + """ + The scheduler holds one cleaner for the life of the process, so the + progress counters must start from zero on every run rather than carrying + an earlier run's totals into the cancellation line. + """ + 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() + _wire_tx(mock_prisma_client.db) + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, 0, 0]) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner.pod_lock_manager = None + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, asyncio.CancelledError()]) + with pytest.raises(asyncio.CancelledError): + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + (error_call,) = mock_logger.error.call_args_list + rendered = error_call[0][0] % error_call[0][1:] + assert "(rows_deleted=150, batches=1)" in rendered