scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 131

refactor(proxy): inject scheduled job shutdown timeouts

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-21 19:37:28 +00:00
parent 9c411dd6f2
commit 4e388e6aea
2 changed files with 16 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/proxy/shutdown/scheduled_jobs.py
View file

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

10
tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py
View file

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