From c1c566db875f31e28e07f662be39207dca9d8344 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Tue, 15 Sep 2026 01:18:18 -0700 Subject: [PATCH 01/10] 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 --- .../db_transaction_queue/spend_log_cleanup.py | 16 +++ litellm/proxy/proxy_server.py | 18 ++- litellm/proxy/shutdown/scheduled_jobs.py | 70 ++++++++++ .../proxy/shutdown/test_scheduled_jobs.py | 125 ++++++++++++++++++ .../proxy/test_spend_log_cleanup.py | 92 +++++++++++++ 5 files changed, 318 insertions(+), 3 deletions(-) create mode 100644 litellm/proxy/shutdown/scheduled_jobs.py create mode 100644 tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py 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 3b5f4236d22..cab0f4d0733 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -680,6 +680,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, @@ -1456,6 +1460,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: @@ -2451,6 +2462,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 @@ -9763,7 +9775,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 @@ -9772,9 +9784,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, @@ -9787,7 +9799,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 From 8ce8dd9b3b184a8e8e8884cde29f07e07063f6f7 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Tue, 15 Sep 2026 02:25:36 -0700 Subject: [PATCH 02/10] 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. --- .../db_transaction_queue/spend_log_cleanup.py | 37 +++++++++--- litellm/proxy/proxy_server.py | 5 ++ litellm/proxy/shutdown/scheduled_jobs.py | 25 +++----- .../proxy/shutdown/test_scheduled_jobs.py | 57 ++++++++++++------- .../proxy/test_spend_log_cleanup.py | 53 +++++++++++++---- 5 files changed, 121 insertions(+), 56 deletions(-) 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 cab0f4d0733..1a7b3a6ccfb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -683,6 +683,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 ( @@ -1419,6 +1420,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: if model_info_scheduler is not scheduler: model_info_scheduler.shutdown(wait=False) + # 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 From 39a199d9a2285a0647fd0d70b3f2a7e2d72120d1 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Tue, 15 Sep 2026 03:13:35 -0700 Subject: [PATCH 03/10] 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 | 16 ++++---- litellm/proxy/shutdown/scheduled_jobs.py | 22 +++++++---- .../proxy/shutdown/test_scheduled_jobs.py | 39 ++++++++++++++----- 3 files changed, 52 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1a7b3a6ccfb..7505714b418 100644 --- a/litellm/proxy/proxy_server.py +++ b/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: diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py index 46c57e1a608..e7625a73b47 100644 --- a/litellm/proxy/shutdown/scheduled_jobs.py +++ b/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", diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py index 3301ce34cd6..7defd6cef6c 100644 --- a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py +++ b/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 From 39a14f39e5961605ba70deacc0475892f50d3cb7 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 20 Sep 2026 08:20:47 +0000 Subject: [PATCH 04/10] style(proxy): format cleanup shutdown tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/shutdown/test_scheduled_jobs.py | 4 +- .../proxy/test_spend_log_cleanup.py | 127 +++++------------- 2 files changed, 37 insertions(+), 94 deletions(-) diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py index 7defd6cef6c..87adc464608 100644 --- a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -7,11 +7,11 @@ from datetime import datetime, timedelta import pytest from apscheduler.schedulers.asyncio import AsyncIOScheduler -import litellm.proxy.shutdown.scheduled_jobs as scheduled_jobs +from litellm.proxy.shutdown import scheduled_jobs from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, - stop_in_flight_scheduler_jobs, pause_scheduled_jobs, + stop_in_flight_scheduler_jobs, ) diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index 1691b2d174a..b8f28d0c780 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -83,10 +83,10 @@ def test_spend_log_cleanup_cron_scheduling(): assert trigger_weekly is not None # Invalid cron expression should raise ValueError - with pytest.raises(ValueError, match='Wrong number of fields; got'): + with pytest.raises(ValueError, match="Wrong number of fields; got"): CronTrigger.from_crontab("invalid cron") - with pytest.raises(ValueError, match='is higher than the maximum value'): + with pytest.raises(ValueError, match="is higher than the maximum value"): CronTrigger.from_crontab("60 25 * * *") # Invalid minute and hour @@ -99,6 +99,7 @@ def test_spend_log_cleanup_cron_scheduler_integration(): a real database connection. """ from unittest.mock import MagicMock + from apscheduler.triggers.cron import CronTrigger # Mock scheduler @@ -145,15 +146,11 @@ def test_spend_log_cleanup_cron_scheduler_integration(): # No cron, so it should fall back to interval } - cleanup_cron_fallback = general_settings_interval.get( - "maximum_spend_logs_cleanup_cron" - ) + cleanup_cron_fallback = general_settings_interval.get("maximum_spend_logs_cleanup_cron") assert cleanup_cron_fallback is None # No cron configured # Simulate interval-based scheduling fallback - retention_interval = general_settings_interval.get( - "maximum_spend_logs_retention_interval", "1d" - ) + retention_interval = general_settings_interval.get("maximum_spend_logs_retention_interval", "1d") from litellm.litellm_core_utils.duration_parser import duration_in_seconds interval_seconds = duration_in_seconds(retention_interval) @@ -181,27 +178,19 @@ async def test_should_delete_spend_logs(): assert cleaner._should_delete_spend_logs() is False # Test case 2: Valid seconds string - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "3600s"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "3600s"}) assert cleaner._should_delete_spend_logs() is True # Test case 3: Valid days string - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "30d"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"}) assert cleaner._should_delete_spend_logs() is True # Test case 4: Valid hours string - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "24h"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "24h"}) assert cleaner._should_delete_spend_logs() is True # Test case 5: Invalid format - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "invalid"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "invalid"}) assert cleaner._should_delete_spend_logs() is False @@ -288,9 +277,7 @@ async def test_cleanup_old_spend_logs_retention_period_cutoff(): # Verify the cutoff date is correct cutoff_date = mock_db.execute_raw.call_args[0][1] expected_cutoff = datetime.now(timezone.utc) - timedelta(seconds=86400) - assert ( - abs((cutoff_date - expected_cutoff).total_seconds()) < 1 - ) # Allow 1 second difference for test execution time + assert abs((cutoff_date - expected_cutoff).total_seconds()) < 1 # Allow 1 second difference for test execution time @pytest.mark.asyncio @@ -310,9 +297,7 @@ async def test_cleanup_drops_partitions_when_enabled_and_partitioned(): partition_manager = MagicMock() partition_manager.is_partitioned = AsyncMock(return_value=True) partition_manager.ensure_partitions = AsyncMock(return_value=["p1"]) - partition_manager.drop_partitions_older_than = AsyncMock( - return_value=["LiteLLM_SpendLogs_p20260601"] - ) + partition_manager.drop_partitions_older_than = AsyncMock(return_value=["LiteLLM_SpendLogs_p20260601"]) cleaner = SpendLogCleanup( general_settings={ @@ -450,9 +435,7 @@ async def test_integer_retention_treated_as_days(): An integer value for maximum_spend_logs_retention_period should be treated as days (e.g., 3 → '3d' → 259200 seconds). """ - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": 3} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": 3}) result = cleaner._should_delete_spend_logs() assert result is True assert cleaner.retention_seconds == 3 * 86400 # 3 days in seconds @@ -469,13 +452,11 @@ def test_string_retention_still_works(): ("2w", 2 * 604800), ] for setting, expected_seconds in cases: - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": setting} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": setting}) assert cleaner._should_delete_spend_logs() is True, f"Failed for {setting}" - assert ( - cleaner.retention_seconds == expected_seconds - ), f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}" + assert cleaner.retention_seconds == expected_seconds, ( + f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}" + ) @pytest.mark.asyncio @@ -489,9 +470,7 @@ async def test_delete_old_logs_aborts_on_non_int_execute_raw_return(): mock_db.execute_raw = AsyncMock(return_value=None) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -510,9 +489,7 @@ async def test_delete_old_logs_continues_on_valid_int_return(): mock_db.execute_raw = AsyncMock(side_effect=[500, 300, 0]) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -559,9 +536,7 @@ async def test_delete_old_tool_index_rows_deletes_on_composite_key(): mock_db.execute_raw = AsyncMock(side_effect=[5, 0]) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_tool_index_rows(mock_prisma_client, cutoff_date, _far_deadline()) @@ -581,9 +556,7 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch) import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module # Zero out the failure backoff so the test doesn't take ~0.5s of real sleep. - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 - ) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -591,14 +564,10 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch) _wire_tx(mock_db) # batch 1 succeeds, batch 2 raises (one-off DB timeout), batches 3-4 succeed, # batch 5 returns 0 → loop exits naturally. - mock_db.execute_raw = AsyncMock( - side_effect=[100, TimeoutError("simulated DB timeout"), 200, 50, 0] - ) + mock_db.execute_raw = AsyncMock(side_effect=[100, TimeoutError("simulated DB timeout"), 200, 50, 0]) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -615,26 +584,18 @@ async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch): import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module # Lower the threshold so the test is fast and deterministic. - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3 - ) - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 - ) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) mock_db = MagicMock() _wire_tx(mock_db) # Every batch raises — must abort after exactly 3 attempts, not loop forever. - mock_db.execute_raw = AsyncMock( - side_effect=ConnectionError("simulated persistent DB outage") - ) + mock_db.execute_raw = AsyncMock(side_effect=ConnectionError("simulated persistent DB outage")) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -649,12 +610,8 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc intermittent timeouts don't trip the abort threshold.""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3 - ) - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 - ) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -675,9 +632,7 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc ) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -698,9 +653,7 @@ async def test_cleanup_uses_logger_exception_for_full_traceback(monkeypatch): mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) # Force the outer try/except to fire by making _should_delete_spend_logs raise. - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cleaner.pod_lock_manager = None def boom(): @@ -725,12 +678,8 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch must still be released so the next scheduled run isn't permanently blocked.""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2 - ) - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 - ) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -744,9 +693,7 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) mock_pod_lock_manager.release_lock = AsyncMock() - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cleaner.pod_lock_manager = mock_pod_lock_manager await cleaner.cleanup_old_spend_logs(mock_prisma_client) @@ -996,9 +943,7 @@ async def test_each_batch_carries_a_statement_and_lock_timeout(): } ) - await cleaner._delete_old_logs( - mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline() - ) + await cleaner._delete_old_logs(mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline()) assert "SET LOCAL statement_timeout = 12000" in recorded assert "SET LOCAL lock_timeout = 12000" in recorded @@ -1134,9 +1079,7 @@ async def test_remaining_rows_probe_is_capped_so_it_cannot_scan_the_table(): cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) - await cleaner._delete_old_logs( - mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline() - ) + await cleaner._delete_old_logs(mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline()) count_sql = mock_db.query_raw.call_args[0][0] assert "count(*)" in count_sql From 0b2d52edc29a7098e95314fbf2c45e58b508efc5 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 20 Sep 2026 08:21:35 +0000 Subject: [PATCH 05/10] Revert "style(proxy): format cleanup shutdown tests" This reverts commit 39a14f39e5961605ba70deacc0475892f50d3cb7. --- .../proxy/shutdown/test_scheduled_jobs.py | 4 +- .../proxy/test_spend_log_cleanup.py | 127 +++++++++++++----- 2 files changed, 94 insertions(+), 37 deletions(-) diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py index 87adc464608..7defd6cef6c 100644 --- a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -7,11 +7,11 @@ from datetime import datetime, timedelta import pytest from apscheduler.schedulers.asyncio import AsyncIOScheduler -from litellm.proxy.shutdown import scheduled_jobs +import litellm.proxy.shutdown.scheduled_jobs as scheduled_jobs from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, - pause_scheduled_jobs, stop_in_flight_scheduler_jobs, + pause_scheduled_jobs, ) diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index b8f28d0c780..1691b2d174a 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -83,10 +83,10 @@ def test_spend_log_cleanup_cron_scheduling(): assert trigger_weekly is not None # Invalid cron expression should raise ValueError - with pytest.raises(ValueError, match="Wrong number of fields; got"): + with pytest.raises(ValueError, match='Wrong number of fields; got'): CronTrigger.from_crontab("invalid cron") - with pytest.raises(ValueError, match="is higher than the maximum value"): + with pytest.raises(ValueError, match='is higher than the maximum value'): CronTrigger.from_crontab("60 25 * * *") # Invalid minute and hour @@ -99,7 +99,6 @@ def test_spend_log_cleanup_cron_scheduler_integration(): a real database connection. """ from unittest.mock import MagicMock - from apscheduler.triggers.cron import CronTrigger # Mock scheduler @@ -146,11 +145,15 @@ def test_spend_log_cleanup_cron_scheduler_integration(): # No cron, so it should fall back to interval } - cleanup_cron_fallback = general_settings_interval.get("maximum_spend_logs_cleanup_cron") + cleanup_cron_fallback = general_settings_interval.get( + "maximum_spend_logs_cleanup_cron" + ) assert cleanup_cron_fallback is None # No cron configured # Simulate interval-based scheduling fallback - retention_interval = general_settings_interval.get("maximum_spend_logs_retention_interval", "1d") + retention_interval = general_settings_interval.get( + "maximum_spend_logs_retention_interval", "1d" + ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds interval_seconds = duration_in_seconds(retention_interval) @@ -178,19 +181,27 @@ async def test_should_delete_spend_logs(): assert cleaner._should_delete_spend_logs() is False # Test case 2: Valid seconds string - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "3600s"}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "3600s"} + ) assert cleaner._should_delete_spend_logs() is True # Test case 3: Valid days string - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "30d"} + ) assert cleaner._should_delete_spend_logs() is True # Test case 4: Valid hours string - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "24h"}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "24h"} + ) assert cleaner._should_delete_spend_logs() is True # Test case 5: Invalid format - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "invalid"}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "invalid"} + ) assert cleaner._should_delete_spend_logs() is False @@ -277,7 +288,9 @@ async def test_cleanup_old_spend_logs_retention_period_cutoff(): # Verify the cutoff date is correct cutoff_date = mock_db.execute_raw.call_args[0][1] expected_cutoff = datetime.now(timezone.utc) - timedelta(seconds=86400) - assert abs((cutoff_date - expected_cutoff).total_seconds()) < 1 # Allow 1 second difference for test execution time + assert ( + abs((cutoff_date - expected_cutoff).total_seconds()) < 1 + ) # Allow 1 second difference for test execution time @pytest.mark.asyncio @@ -297,7 +310,9 @@ async def test_cleanup_drops_partitions_when_enabled_and_partitioned(): partition_manager = MagicMock() partition_manager.is_partitioned = AsyncMock(return_value=True) partition_manager.ensure_partitions = AsyncMock(return_value=["p1"]) - partition_manager.drop_partitions_older_than = AsyncMock(return_value=["LiteLLM_SpendLogs_p20260601"]) + partition_manager.drop_partitions_older_than = AsyncMock( + return_value=["LiteLLM_SpendLogs_p20260601"] + ) cleaner = SpendLogCleanup( general_settings={ @@ -435,7 +450,9 @@ async def test_integer_retention_treated_as_days(): An integer value for maximum_spend_logs_retention_period should be treated as days (e.g., 3 → '3d' → 259200 seconds). """ - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": 3}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": 3} + ) result = cleaner._should_delete_spend_logs() assert result is True assert cleaner.retention_seconds == 3 * 86400 # 3 days in seconds @@ -452,11 +469,13 @@ def test_string_retention_still_works(): ("2w", 2 * 604800), ] for setting, expected_seconds in cases: - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": setting}) - assert cleaner._should_delete_spend_logs() is True, f"Failed for {setting}" - assert cleaner.retention_seconds == expected_seconds, ( - f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}" + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": setting} ) + assert cleaner._should_delete_spend_logs() is True, f"Failed for {setting}" + assert ( + cleaner.retention_seconds == expected_seconds + ), f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}" @pytest.mark.asyncio @@ -470,7 +489,9 @@ async def test_delete_old_logs_aborts_on_non_int_execute_raw_return(): mock_db.execute_raw = AsyncMock(return_value=None) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -489,7 +510,9 @@ async def test_delete_old_logs_continues_on_valid_int_return(): mock_db.execute_raw = AsyncMock(side_effect=[500, 300, 0]) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -536,7 +559,9 @@ async def test_delete_old_tool_index_rows_deletes_on_composite_key(): mock_db.execute_raw = AsyncMock(side_effect=[5, 0]) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_tool_index_rows(mock_prisma_client, cutoff_date, _far_deadline()) @@ -556,7 +581,9 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch) import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module # Zero out the failure backoff so the test doesn't take ~0.5s of real sleep. - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 + ) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -564,10 +591,14 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch) _wire_tx(mock_db) # batch 1 succeeds, batch 2 raises (one-off DB timeout), batches 3-4 succeed, # batch 5 returns 0 → loop exits naturally. - mock_db.execute_raw = AsyncMock(side_effect=[100, TimeoutError("simulated DB timeout"), 200, 50, 0]) + mock_db.execute_raw = AsyncMock( + side_effect=[100, TimeoutError("simulated DB timeout"), 200, 50, 0] + ) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -584,18 +615,26 @@ async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch): import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module # Lower the threshold so the test is fast and deterministic. - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3 + ) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 + ) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) mock_db = MagicMock() _wire_tx(mock_db) # Every batch raises — must abort after exactly 3 attempts, not loop forever. - mock_db.execute_raw = AsyncMock(side_effect=ConnectionError("simulated persistent DB outage")) + mock_db.execute_raw = AsyncMock( + side_effect=ConnectionError("simulated persistent DB outage") + ) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -610,8 +649,12 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc intermittent timeouts don't trip the abort threshold.""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3 + ) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 + ) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -632,7 +675,9 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc ) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -653,7 +698,9 @@ async def test_cleanup_uses_logger_exception_for_full_traceback(monkeypatch): mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) # Force the outer try/except to fire by making _should_delete_spend_logs raise. - cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cleaner.pod_lock_manager = None def boom(): @@ -678,8 +725,12 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch must still be released so the next scheduled run isn't permanently blocked.""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2) - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2 + ) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 + ) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -693,7 +744,9 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) mock_pod_lock_manager.release_lock = AsyncMock() - cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner = cleanup_module.SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"} + ) cleaner.pod_lock_manager = mock_pod_lock_manager await cleaner.cleanup_old_spend_logs(mock_prisma_client) @@ -943,7 +996,9 @@ async def test_each_batch_carries_a_statement_and_lock_timeout(): } ) - await cleaner._delete_old_logs(mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline()) + await cleaner._delete_old_logs( + mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline() + ) assert "SET LOCAL statement_timeout = 12000" in recorded assert "SET LOCAL lock_timeout = 12000" in recorded @@ -1079,7 +1134,9 @@ async def test_remaining_rows_probe_is_capped_so_it_cannot_scan_the_table(): cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) - await cleaner._delete_old_logs(mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline()) + await cleaner._delete_old_logs( + mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline() + ) count_sql = mock_db.query_raw.call_args[0][0] assert "count(*)" in count_sql From 29bd2ceb2b5c8f2a172d6b08e6c54cb20c41daaa Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 20 Sep 2026 08:27:06 +0000 Subject: [PATCH 06/10] fix(proxy): avoid mutable shutdown wait set Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/shutdown/scheduled_jobs.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py index e7625a73b47..5345d380112 100644 --- a/litellm/proxy/shutdown/scheduled_jobs.py +++ b/litellm/proxy/shutdown/scheduled_jobs.py @@ -50,12 +50,13 @@ async def stop_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: 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) + still_running: Final = ( + (await asyncio.wait(in_flight, timeout=JOB_FINISH_TIMEOUT_SECONDS))[1] if in_flight else frozenset() + ) scheduler.shutdown(wait=False) if not still_running: return From 9c411dd6f2e569ea7ac54c08f03847d5c70cddba Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 19:35:29 +0000 Subject: [PATCH 07/10] refactor(proxy): make scheduled job shutdown timeouts configurable via env Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 ++ litellm/proxy/shutdown/scheduled_jobs.py | 23 +++++++++++-------- .../proxy/shutdown/test_scheduled_jobs.py | 15 ++++++------ 3 files changed, 24 insertions(+), 16 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index bbeb4846e27..842adf62f6b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1742,6 +1742,8 @@ SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = float( SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", "300")) SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", "30")) SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = int(os.getenv("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", "100000")) +SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS: Final = float(os.getenv("SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS", "5")) +SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS: Final = float(os.getenv("SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS", "5")) TOOL_SPEND_TOP_TOOLS: Final = 100 SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day") SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)) diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py index 5345d380112..cf4937780b8 100644 --- a/litellm/proxy/shutdown/scheduled_jobs.py +++ b/litellm/proxy/shutdown/scheduled_jobs.py @@ -7,9 +7,10 @@ from typing import Final, Protocol 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 +from litellm.constants import ( + SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS, + SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS, +) class StoppableScheduler(Protocol): @@ -41,8 +42,8 @@ def pause_scheduled_jobs(scheduler: StoppableScheduler) -> None: async def stop_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: AwaitableAsyncIOExecutor) -> None: """ - 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. + Let in-flight jobs finish for up to SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS, then stop the scheduler and + wait, bounded by SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS, for the jobs it cancels. 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. @@ -52,19 +53,23 @@ async def stop_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: in_flight: Final = executor.in_flight_jobs() 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) + "Waiting up to %ss for %d in-flight scheduled job(s) to finish", + SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS, + len(in_flight), ) still_running: Final = ( - (await asyncio.wait(in_flight, timeout=JOB_FINISH_TIMEOUT_SECONDS))[1] if in_flight else frozenset() + (await asyncio.wait(in_flight, timeout=SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS))[1] + if in_flight + else frozenset() ) scheduler.shutdown(wait=False) if not still_running: return 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) + _done, pending = await asyncio.wait(still_running, timeout=SCHEDULED_JOB_SHUTDOWN_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, + SCHEDULED_JOB_SHUTDOWN_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 index 7defd6cef6c..1f94cee04ed 100644 --- a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -7,11 +7,11 @@ from datetime import datetime, timedelta import pytest from apscheduler.schedulers.asyncio import AsyncIOScheduler -import litellm.proxy.shutdown.scheduled_jobs as scheduled_jobs +from litellm.constants import SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, - stop_in_flight_scheduler_jobs, pause_scheduled_jobs, + stop_in_flight_scheduler_jobs, ) @@ -75,9 +75,8 @@ async def test_in_flight_jobs_observe_cancellation_before_shutdown_returns(): @pytest.mark.asyncio -async def test_a_job_that_is_finishing_is_allowed_to_finish_rather_than_cancelled(monkeypatch): +async def test_a_job_that_is_finishing_is_allowed_to_finish_rather_than_cancelled(): """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): @@ -99,16 +98,18 @@ 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): +async def test_a_job_that_ignores_cancellation_is_abandoned_after_the_timeout(caplog): """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): with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): 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 + assert ( + f"1 scheduled job(s) did not finish within {SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS}s of cancellation" + in caplog.text + ) @pytest.mark.asyncio From 4e388e6aea52ad6c9c2939997f860689e69716aa Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 19:37:28 +0000 Subject: [PATCH 08/10] refactor(proxy): inject scheduled job shutdown timeouts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/shutdown/scheduled_jobs.py | 20 ++++++++++++------- .../proxy/shutdown/test_scheduled_jobs.py | 10 +++------- 2 files changed, 16 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py index cf4937780b8..e920ce19eb9 100644 --- a/litellm/proxy/shutdown/scheduled_jobs.py +++ b/litellm/proxy/shutdown/scheduled_jobs.py @@ -40,10 +40,16 @@ def pause_scheduled_jobs(scheduler: StoppableScheduler) -> None: scheduler.pause() -async def stop_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: AwaitableAsyncIOExecutor) -> None: +async def stop_in_flight_scheduler_jobs( + scheduler: StoppableScheduler, + executor: AwaitableAsyncIOExecutor, + *, + finish_timeout_seconds: float = SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS, + cancel_timeout_seconds: float = SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS, +) -> None: """ - Let in-flight jobs finish for up to SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS, then stop the scheduler and - wait, bounded by SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS, for the jobs it cancels. + Let in-flight jobs finish for up to finish_timeout_seconds, then stop the scheduler and wait, bounded by + cancel_timeout_seconds, for the jobs it cancels. 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. @@ -54,11 +60,11 @@ async def stop_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: if in_flight: verbose_proxy_logger.info( "Waiting up to %ss for %d in-flight scheduled job(s) to finish", - SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS, + finish_timeout_seconds, len(in_flight), ) still_running: Final = ( - (await asyncio.wait(in_flight, timeout=SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS))[1] + (await asyncio.wait(in_flight, timeout=finish_timeout_seconds))[1] if in_flight else frozenset() ) @@ -66,10 +72,10 @@ async def stop_in_flight_scheduler_jobs(scheduler: StoppableScheduler, executor: if not still_running: return verbose_proxy_logger.info("Cancelling %d in-flight scheduled job(s) for shutdown", len(still_running)) - _done, pending = await asyncio.wait(still_running, timeout=SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS) + _done, pending = await asyncio.wait(still_running, timeout=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), - SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS, + 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 index 1f94cee04ed..fbce38db39f 100644 --- a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -7,7 +7,6 @@ from datetime import datetime, timedelta import pytest from apscheduler.schedulers.asyncio import AsyncIOScheduler -from litellm.constants import SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, pause_scheduled_jobs, @@ -80,7 +79,7 @@ async def test_a_job_that_is_finishing_is_allowed_to_finish_rather_than_cancelle 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) + await stop_in_flight_scheduler_jobs(scheduler, executor, finish_timeout_seconds=2.0) assert write.events == ["committed", "finished"] assert stuck.events == ["cancelled", "finished"] @@ -103,13 +102,10 @@ async def test_a_job_that_ignores_cancellation_is_abandoned_after_the_timeout(ca job = _Job(swallow_cancellation=True) async with _running_scheduler(job) as (scheduler, executor): with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - await stop_in_flight_scheduler_jobs(scheduler, executor) + await stop_in_flight_scheduler_jobs(scheduler, executor, cancel_timeout_seconds=0.05) assert job.events == ["cancelled"] - assert ( - f"1 scheduled job(s) did not finish within {SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS}s of cancellation" - in caplog.text - ) + assert "1 scheduled job(s) did not finish within 0.05s of cancellation" in caplog.text @pytest.mark.asyncio From ee07f710630bdadb7035f08cfb9cd728ec4425db Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 19:43:11 +0000 Subject: [PATCH 09/10] style(proxy): format scheduled job timeout configuration Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 8 ++++++-- litellm/proxy/shutdown/scheduled_jobs.py | 4 +--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 842adf62f6b..1971c336a96 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1742,8 +1742,12 @@ SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = float( SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", "300")) SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", "30")) SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = int(os.getenv("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", "100000")) -SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS: Final = float(os.getenv("SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS", "5")) -SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS: Final = float(os.getenv("SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS", "5")) +SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS: Final = float( + os.getenv("SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS", "5") +) +SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS: Final = float( + os.getenv("SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS", "5") +) TOOL_SPEND_TOP_TOOLS: Final = 100 SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day") SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)) diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py index e920ce19eb9..7889c35cf4e 100644 --- a/litellm/proxy/shutdown/scheduled_jobs.py +++ b/litellm/proxy/shutdown/scheduled_jobs.py @@ -64,9 +64,7 @@ async def stop_in_flight_scheduler_jobs( len(in_flight), ) still_running: Final = ( - (await asyncio.wait(in_flight, timeout=finish_timeout_seconds))[1] - if in_flight - else frozenset() + (await asyncio.wait(in_flight, timeout=finish_timeout_seconds))[1] if in_flight else frozenset() ) scheduler.shutdown(wait=False) if not still_running: From 0a88658227f8e0d7e2d4928df7c5ee63bd83fd1d Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 20:25:56 +0000 Subject: [PATCH 10/10] chore: retrigger ci after docs merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>