From 9c411dd6f2e569ea7ac54c08f03847d5c70cddba Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 19:35:29 +0000 Subject: [PATCH] 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