From 4e388e6aea52ad6c9c2939997f860689e69716aa Mon Sep 17 00:00:00 2001
From: yucheng <yucheng@berri.ai>
Date: Mon, 21 Sep 2026 19:37:28 +0000
Subject: [PATCH] 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