fix(proxy): record aborted outcome when spend-log cleanup is cancelled at shutdown

cleanup_old_spend_logs only caught Exception, so a run cut short by
CancelledError recorded no outcome and logged nothing. Under uvicorn the
job was never cancelled at all: uvicorn re-raises the captured SIGTERM as
soon as the lifespan shutdown returns, before asyncio cancels outstanding
tasks, so an in-flight scheduler job simply died with the process.

The cleanup now handles CancelledError by logging elapsed time, rows
deleted and batch count at error level, recording outcome="aborted", and
re-raising. The lifespan shutdown stops the scheduler and awaits the jobs
it cancels while the database is still connected, so that handler runs
under uvicorn too, and the pod lock is released instead of orphaned.

Resolves LIT-6990

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.

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,

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,
+        )

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

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