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