fix(proxy): let in-flight scheduled jobs finish before cancelling them at shutdown

Cancelling every in-flight job the moment shutdown reached the scheduler
dropped the rows a write job had already popped: flush_gateway_requests
drains its accumulator before committing and does not restore it on
CancelledError, and update_spend requeues its batch only after the
shutdown drain had already run.

Shutdown now waits up to JOB_FINISH_TIMEOUT_SECONDS for in-flight jobs
to finish on their own, cancels the ones still running, and does both
before the shutdown flushes so a requeued batch is still written. The
cleanup run never finishes inside the grace, so it is still cancelled
and still records outcome="aborted".

Resolves LIT-6990

litellm/proxy/proxy_server.py

@@ -682,8 +682,8 @@ 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,
     pause_scheduled_jobs,
+    stop_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 (
@@ -1457,6 +1457,13 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
 
     await _drain_spend_event_producer_on_shutdown()
 
+    # Shutdown event - finish or cancel in-flight scheduled jobs before the shutdown flushes and the DB disconnect
+    if scheduler is not None and scheduler_executor is not None:
+        try:
+            await stop_in_flight_scheduler_jobs(scheduler, scheduler_executor)
+        except Exception as e:
+            verbose_proxy_logger.error("Error stopping in-flight scheduled jobs: %s", e)
+
     await flush_spend_counters_on_shutdown()
 
     await _flush_spend_logs_queue_on_shutdown()
@@ -1465,13 +1472,6 @@ 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:

litellm/proxy/shutdown/scheduled_jobs.py

@@ -8,6 +8,7 @@ from apscheduler.executors.asyncio import AsyncIOExecutor
 
 from litellm._logging import verbose_proxy_logger
 
+JOB_FINISH_TIMEOUT_SECONDS: Final = 5.0
 JOB_CANCEL_TIMEOUT_SECONDS: Final = 5.0
 
 
@@ -38,21 +39,28 @@ def pause_scheduled_jobs(scheduler: StoppableScheduler) -> None:
         scheduler.pause()
 
 
-async def cancel_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: AwaitableAsyncIOExecutor) -> None:
+async def stop_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.
+    Let in-flight jobs finish for up to JOB_FINISH_TIMEOUT_SECONDS, then 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.
+    Must run before the database is disconnected: a write job that finishes needs its connection,
+    and a job's cancellation handler is what records the run's outcome.
     """
     if not scheduler.running:
         return
     in_flight: Final = executor.in_flight_jobs()
+    still_running: set[asyncio.Future[object]] = set()
+    if in_flight:
+        verbose_proxy_logger.info(
+            "Waiting up to %ss for %d in-flight scheduled job(s) to finish", JOB_FINISH_TIMEOUT_SECONDS, len(in_flight)
+        )
+        _done, still_running = await asyncio.wait(in_flight, timeout=JOB_FINISH_TIMEOUT_SECONDS)
     scheduler.shutdown(wait=False)
-    if not in_flight:
+    if not still_running:
         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)
+    verbose_proxy_logger.info("Cancelling %d in-flight scheduled job(s) for shutdown", len(still_running))
+    _done, pending = await asyncio.wait(still_running, 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",

tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py

@@ -10,23 +10,28 @@ 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,
+    stop_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, or for ``work_seconds``, and records what it observed"""
 
-    def __init__(self, swallow_cancellation: bool = False) -> None:
+    def __init__(self, swallow_cancellation: bool = False, work_seconds: float | None = None) -> None:
         self.started = asyncio.Event()
         self.events: list[str] = []
         self.swallow_cancellation = swallow_cancellation
+        self.work_seconds = work_seconds
 
     async def run(self) -> None:
         self.started.set()
         try:
-            await asyncio.Event().wait()
+            if self.work_seconds is None:
+                await asyncio.Event().wait()
+            else:
+                await asyncio.sleep(self.work_seconds)
+                self.events.append("committed")
         except asyncio.CancelledError:
             self.events.append("cancelled")
             if self.swallow_cancellation:
@@ -62,18 +67,32 @@ async def test_in_flight_jobs_observe_cancellation_before_shutdown_returns():
     """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)
+        await stop_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_a_job_that_is_finishing_is_allowed_to_finish_rather_than_cancelled(monkeypatch):
+    """A spend write cancelled mid-commit drops the rows it popped, so short jobs get to finish first"""
+    monkeypatch.setattr(scheduled_jobs, "JOB_FINISH_TIMEOUT_SECONDS", 2.0)
+    write = _Job(work_seconds=0.2)
+    stuck = _Job()
+    async with _running_scheduler(write, stuck) as (scheduler, executor):
+        await stop_in_flight_scheduler_jobs(scheduler, executor)
+
+        assert write.events == ["committed", "finished"]
+        assert stuck.events == ["cancelled", "finished"]
+        assert scheduler.running is False
+
+
 @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)
+        await stop_in_flight_scheduler_jobs(scheduler, executor)
 
     assert first.events == ["cancelled", "finished"]
     assert second.events == ["cancelled", "finished"]
@@ -86,7 +105,7 @@ async def test_a_job_that_ignores_cancellation_is_abandoned_after_the_timeout(mo
     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)
+            await stop_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
@@ -95,7 +114,7 @@ async def test_a_job_that_ignores_cancellation_is_abandoned_after_the_timeout(mo
 @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 stop_in_flight_scheduler_jobs(scheduler, executor)
         await asyncio.sleep(0)
 
         assert scheduler.running is False
@@ -107,7 +126,7 @@ async def test_a_scheduler_that_never_started_is_left_alone():
     executor = AwaitableAsyncIOExecutor()
     scheduler = AsyncIOScheduler(executors={"default": executor})
 
-    await cancel_in_flight_scheduler_jobs(scheduler, executor)
+    await stop_in_flight_scheduler_jobs(scheduler, executor)
 
     assert scheduler.running is False
 
@@ -127,7 +146,7 @@ async def test_pausing_stops_new_jobs_from_starting_but_leaves_running_ones_alon
         assert running.events == []
         assert scheduler.running is True
 
-        await cancel_in_flight_scheduler_jobs(scheduler, executor)
+        await stop_in_flight_scheduler_jobs(scheduler, executor)
 
     assert running.events == ["cancelled", "finished"]
     assert late.started.is_set() is False