fix(proxy): carry the graceful shutdown timeout into the reload worker pool

This commit is contained in:
mateo-berri 2026-10-20 18:33:38 -07:00
parent 0c0ea5cfc3
commit a02c0c7582
2 changed files with 23 additions and 2 deletions

View file

@ -21,6 +21,7 @@ class ServerOptions(TypedDict, total=False):
host: ReadOnly[str]
port: ReadOnly[int]
server_header: ReadOnly[bool]
timeout_graceful_shutdown: ReadOnly[float]
log_config: ReadOnly[str | dict[str, JsonValue]] # mutable-ok: uvicorn.Config requires a concrete logging dict
log_level: ReadOnly[str | int]
timeout_keep_alive: ReadOnly[int]
@ -67,7 +68,7 @@ def pool_config(options: ReloadOptions, workers: int) -> Config:
partial(load_app_or_exit, options["app"]),
factory=True,
workers=workers,
**SERVER_OPTIONS.validate_python({"timeout_worker_healthcheck": RELOAD_WORKER_HEALTHCHECK_SECONDS, **options}),
**SERVER_OPTIONS.validate_python({"timeout_worker_healthcheck": RELOAD_WORKER_HEALTHCHECK_SECONDS, **options}), # pyright: ignore[reportArgumentType] # uvicorn types the float graceful shutdown timeout as int
)
@ -102,7 +103,7 @@ def run_multi_worker_reload(options: ReloadOptions, workers: int) -> None:
reload=options["reload"],
reload_dirs=list(options["reload_dirs"]),
reload_includes=list(options["reload_includes"]),
**SERVER_OPTIONS.validate_python(options),
**SERVER_OPTIONS.validate_python(options), # pyright: ignore[reportArgumentType] # uvicorn types the float graceful shutdown timeout as int
)
shared_socket: Final = worker_config.bind_socket()
try:

View file

@ -22,6 +22,7 @@ import pytest
from uvicorn import config as uvicorn_config
from litellm.proxy.dev_reload import (
RELOAD_OPTIONS,
RELOAD_WORKER_HEALTHCHECK_SECONDS,
ReloadOptions,
load_app_or_exit,
@ -29,6 +30,7 @@ from litellm.proxy.dev_reload import (
pool_config,
stop_pool_when_reloader_exits,
)
from litellm.proxy.proxy_cli import ProxyInitializationHelpers
UNCHANGED_LOGGING: Final = {"version": 1, "incremental": True}
POOL_WORKERS: Final = 2
@ -132,6 +134,24 @@ def test_the_reload_pool_defaults_the_worker_healthcheck_and_keeps_an_explicit_o
)
def test_the_reload_pool_accepts_and_keeps_every_uvicorn_arg_the_proxy_cli_builds(tmp_path: Path) -> None:
cli_args: Final = ProxyInitializationHelpers.get_default_unvicorn_init_args(
"127.0.0.1", 4123, keepalive_timeout=61, timeout_worker_healthcheck=RELOAD_WORKER_HEALTHCHECK_SECONDS + 7
)
options: Final = RELOAD_OPTIONS.validate_python(
{
**cli_args,
"log_config": UNCHANGED_LOGGING,
"reload": True,
"reload_dirs": [str(tmp_path)],
"reload_includes": ["*.py"],
}
)
config: Final = pool_config(options, workers=POOL_WORKERS)
server_args: Final = {key: value for key, value in cli_args.items() if key not in ("app", "log_config")}
assert {key: getattr(config, key) for key in server_args} == server_args
def test_the_pool_is_told_to_stop_once_the_reloader_process_exits() -> None:
reloader: Final = multiprocessing.get_context("spawn").Process(target=time.sleep, args=(POOL_TIMEOUT_SECONDS,))
reloader.start()