diff --git a/litellm/proxy/dev_reload.py b/litellm/proxy/dev_reload.py index 2ccd8e9ec45..f2494bd41bd 100644 --- a/litellm/proxy/dev_reload.py +++ b/litellm/proxy/dev_reload.py @@ -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: diff --git a/tests/unit/proxy/test_dev_reload.py b/tests/unit/proxy/test_dev_reload.py index fd6926c153f..39f9de17101 100644 --- a/tests/unit/proxy/test_dev_reload.py +++ b/tests/unit/proxy/test_dev_reload.py @@ -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()