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

@@ -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))

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

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