fix(proxy): pause the scheduler at shutdown start and keep cleanup progress per run

Review follow-ups on #41213:

- Pause the scheduler as the first shutdown step so a job whose fire time
  falls inside the shutdown window does not start only to be cancelled.
  Jobs already running keep the whole window and are cancelled and
  awaited before the database disconnects, as before.
- Keep the cleanup run's progress in a task-scoped ContextVar rather than
  on the cleaner instance, so two runs overlapping on one cleaner
  (APSCHEDULER_MAX_INSTANCES above 1 without a Redis lock) each report
  their own rows and batches on cancellation.
- Drop the module docstrings the repository comment policy does not
  allow; the rationale lives in the PR description.

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)

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()

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.

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

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