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 1a14210dbec..34213c0d2ce 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -1,5 +1,6 @@ import asyncio import time +from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Final, Literal, TypeAlias @@ -40,6 +41,28 @@ class TableCleanupResult: stop_reason: StopReason +class _RunProgress: + """How far one cleanup run has got, reported if that run is cancelled""" + + def __init__(self) -> None: + self.rows_deleted: int = 0 + self.batches: int = 0 + + def record_batch(self, rows_deleted: int) -> None: + self.rows_deleted += rows_deleted + self.batches += 1 + + +_run_progress: ContextVar[_RunProgress] = ContextVar("spend_log_cleanup_run_progress") + + +def _record_run_batch(rows_deleted: int) -> None: + """Count a batch towards the run in progress, if a run is what issued it""" + progress: Final = _run_progress.get(None) + if progress is not None: + progress.record_batch(rows_deleted) + + class _RemainingRow(BaseModel): """One row of the capped outstanding-rows probe, validated out of prisma's untyped result.""" @@ -96,8 +119,6 @@ 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 @@ -424,8 +445,7 @@ class SpendLogCleanup: total_deleted += deleted_count run_count += 1 - self._run_rows_deleted += deleted_count - self._run_batches += 1 + _record_run_batch(deleted_count) # Add a small sleep to prevent overwhelming the database await asyncio.sleep(0.1) @@ -595,8 +615,8 @@ class SpendLogCleanup: """ lock_acquired = False run_started_at: Final = time.monotonic() - self._run_rows_deleted = 0 - self._run_batches = 0 + progress: Final = _RunProgress() + progress_token: Final = _run_progress.set(progress) try: verbose_proxy_logger.info("Cleanup job triggered at %s", datetime.now()) self._refresh_bounds() @@ -681,8 +701,8 @@ class SpendLogCleanup: 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, + progress.rows_deleted, + progress.batches, ) SpendLogCleanupMetrics.record_run("aborted") raise @@ -697,6 +717,7 @@ class SpendLogCleanup: SpendLogCleanupMetrics.record_run("aborted") return # Return after error handling finally: + _run_progress.reset(progress_token) # Only release the lock if it was actually acquired if lock_acquired and 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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c2dab669087..4d4767f892a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -664,6 +664,7 @@ from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownMan from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, cancel_in_flight_scheduler_jobs, + pause_scheduled_jobs, ) from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start from litellm.proxy.spend_tracking.spend_counter_batch import ( @@ -1376,6 +1377,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: # End of startup event yield + # Shutdown event - stop starting scheduled jobs; the ones already running keep the drain window + if scheduler is not None: + pause_scheduled_jobs(scheduler) + # Shutdown event - drain in-flight requests before tearing down dependencies # so SIGTERM (rolling update, scale-down, liveness kill) doesn't drop them. GracefulShutdownManager.start_shutdown() diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py index 3c9e791f51c..46c57e1a608 100644 --- a/litellm/proxy/shutdown/scheduled_jobs.py +++ b/litellm/proxy/shutdown/scheduled_jobs.py @@ -1,20 +1,3 @@ -""" -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 @@ -34,6 +17,8 @@ class StoppableScheduler(Protocol): @property def running(self) -> bool: ... + def pause(self) -> None: ... + def shutdown(self, wait: bool = ...) -> None: ... @@ -47,6 +32,12 @@ class AwaitableAsyncIOExecutor(AsyncIOExecutor): # pyright: ignore[reportUntype return tuple(future for future in self._pending_futures if not future.done()) +def pause_scheduled_jobs(scheduler: StoppableScheduler) -> None: + """Stop the scheduler from starting jobs that shutdown would only cancel; running jobs continue""" + if scheduler.running: + scheduler.pause() + + 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. diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py index b77d7c4ae50..3301ce34cd6 100644 --- a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -1,16 +1,8 @@ -""" -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 +from datetime import datetime, timedelta import pytest from apscheduler.schedulers.asyncio import AsyncIOScheduler @@ -19,11 +11,12 @@ import litellm.proxy.shutdown.scheduled_jobs as scheduled_jobs from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, cancel_in_flight_scheduler_jobs, + pause_scheduled_jobs, ) class _Job: - """A scheduled job that blocks until cancelled and records what it observed.""" + """A scheduled job that blocks until cancelled and records what it observed""" def __init__(self, swallow_cancellation: bool = False) -> None: self.started = asyncio.Event() @@ -45,7 +38,7 @@ class _Job: @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.""" + """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): @@ -66,10 +59,7 @@ async def _running_scheduler(*jobs: _Job) -> AsyncIterator[tuple[AsyncIOSchedule @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. - """ + """The job's own CancelledError handler records how a run ended, so shutdown must wait for it""" job = _Job() async with _running_scheduler(job) as (scheduler, executor): await cancel_in_flight_scheduler_jobs(scheduler, executor) @@ -91,10 +81,7 @@ async def test_every_in_flight_job_is_cancelled_not_only_the_first(): @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. - """ + """A job that swallows CancelledError must not hold the pod past its termination grace period""" monkeypatch.setattr(scheduled_jobs, "JOB_CANCEL_TIMEOUT_SECONDS", 0.05) job = _Job(swallow_cancellation=True) async with _running_scheduler(job) as (scheduler, executor): @@ -116,10 +103,40 @@ async def test_shutdown_with_nothing_in_flight_still_stops_the_scheduler(): @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.""" + """The proxy runs without a scheduler when it has no database""" executor = AwaitableAsyncIOExecutor() scheduler = AsyncIOScheduler(executors={"default": executor}) await cancel_in_flight_scheduler_jobs(scheduler, executor) assert scheduler.running is False + + +@pytest.mark.asyncio +async def test_pausing_stops_new_jobs_from_starting_but_leaves_running_ones_alone(): + """A job due during the shutdown drain would only be cancelled, so it must not start at all""" + running = _Job() + async with _running_scheduler(running) as (scheduler, executor): + late = _Job() + scheduler.add_job(late.run, id="late", next_run_time=datetime.now() + timedelta(seconds=0.1)) + + pause_scheduled_jobs(scheduler) + await asyncio.sleep(0.3) + + assert late.started.is_set() is False + assert running.events == [] + assert scheduler.running is True + + await cancel_in_flight_scheduler_jobs(scheduler, executor) + + assert running.events == ["cancelled", "finished"] + assert late.started.is_set() is False + + +@pytest.mark.asyncio +async def test_pausing_a_scheduler_that_never_started_is_a_no_op(): + scheduler = AsyncIOScheduler(executors={"default": AwaitableAsyncIOExecutor()}) + + pause_scheduled_jobs(scheduler) + + 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 ed35af7ee38..1691b2d174a 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -1432,12 +1432,7 @@ def _runs_recorded(outcome: str) -> float: @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. - """ + """A run cut short by shutdown must leave its outcome and how far it got behind""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module mock_logger = MagicMock() @@ -1485,11 +1480,7 @@ async def test_a_cancelled_run_records_aborted_and_logs_its_progress_before_re_r @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. - """ + """The scheduler holds one cleaner for the life of the process, so progress must not carry over""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module mock_logger = MagicMock() @@ -1509,3 +1500,43 @@ async def test_progress_reported_for_a_cancelled_run_is_that_run_only(monkeypatc (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 + + +@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""" + 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) + + first_batch_done = asyncio.Event() + second_run_done = asyncio.Event() + + async def _slow_execute_raw(sql, *args): + first_batch_done.set() + await second_run_done.wait() + return 100 + + 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]) + + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner.pod_lock_manager = None + + slow_run = asyncio.ensure_future(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) + second_run_done.set() + await asyncio.sleep(0) + slow_run.cancel() + with pytest.raises(asyncio.CancelledError): + await slow_run + + (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