From 82c0ed923e6f4d55925f33a125d5b3b7431f709b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:47:37 -0700 Subject: [PATCH] fix(proxy): stop the reload worker pool when the reloader exits and give reloads a longer worker healthcheck Closing the terminal or pressing Ctrl-\ killed the reloader and left the pool parent holding the port while it respawned workers that could not start. The pool now stops once its parent reloader is gone. The reload pool also defaults the worker healthcheck to 30 seconds unless --timeout_worker_healthcheck is set, since uvicorn's 5 second default killed workers that were still booting on a reload. --- litellm/proxy/dev_reload.py | 14 ++++- tests/unit/proxy/test_dev_reload.py | 93 +++++++++++++++++++++++------ 2 files changed, 89 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/dev_reload.py b/litellm/proxy/dev_reload.py index 693ba44578a..a29319b8542 100644 --- a/litellm/proxy/dev_reload.py +++ b/litellm/proxy/dev_reload.py @@ -1,6 +1,8 @@ import importlib import logging +import os import sys +import threading from collections.abc import Callable, Sequence from functools import partial from inspect import signature @@ -42,6 +44,8 @@ RELOAD_OPTIONS: Final = TypeAdapter(ReloadOptions) SERVER_OPTIONS: Final = TypeAdapter(ServerOptions) ASGI_APP: Final = TypeAdapter(ImportString[Callable[..., object]]) LOGGER: Final = logging.getLogger("uvicorn.error") +RELOAD_WORKER_HEALTHCHECK_SECONDS: Final = 30 +RELOADER_CHECK_SECONDS: Final = 0.5 def multi_worker_reload_supported() -> bool: @@ -63,10 +67,17 @@ def pool_config(options: ReloadOptions, workers: int) -> Config: partial(load_app_or_exit, options["app"]), factory=True, workers=workers, - **SERVER_OPTIONS.validate_python(options), + **SERVER_OPTIONS.validate_python({"timeout_worker_healthcheck": RELOAD_WORKER_HEALTHCHECK_SECONDS, **options}), ) +def stop_pool_when_reloader_exits(pool: Multiprocess, reloader_pid: int) -> None: + while not pool.should_exit.wait(RELOADER_CHECK_SECONDS): + if os.getppid() != reloader_pid: + LOGGER.info("Reloader process [%d] exited, stopping the worker pool.", reloader_pid) + pool.should_exit.set() + + def _run_worker_pool(config: Config, sockets: Sequence[socket] | None) -> None: shared_sockets: Final = list(sockets or ()) pool: Final = ( @@ -78,6 +89,7 @@ def _run_worker_pool(config: Config, sockets: Sequence[socket] | None) -> None: if "target" in signature(Multiprocess).parameters else Multiprocess(config=config, sockets=shared_sockets) ) + threading.Thread(target=stop_pool_when_reloader_exits, args=(pool, os.getppid()), daemon=True).start() pool.run() diff --git a/tests/unit/proxy/test_dev_reload.py b/tests/unit/proxy/test_dev_reload.py index 5ebf7c152d9..5f7b4093855 100644 --- a/tests/unit/proxy/test_dev_reload.py +++ b/tests/unit/proxy/test_dev_reload.py @@ -19,7 +19,13 @@ import psutil import pytest from uvicorn import config as uvicorn_config -from litellm.proxy.dev_reload import ReloadOptions, load_app_or_exit, multi_worker_reload_supported, pool_config +from litellm.proxy.dev_reload import ( + RELOAD_WORKER_HEALTHCHECK_SECONDS, + ReloadOptions, + load_app_or_exit, + multi_worker_reload_supported, + pool_config, +) UNCHANGED_LOGGING: Final = {"version": 1, "incremental": True} POOL_WORKERS: Final = 2 @@ -114,6 +120,15 @@ def test_a_broken_import_logs_the_line_that_raised( assert (str(app_file), 2) in raised_at, record.exc_text +def test_the_reload_pool_defaults_the_worker_healthcheck_and_keeps_an_explicit_one() -> None: + explicit: Final = RELOAD_WORKER_HEALTHCHECK_SECONDS + 7 + assert pool_config(_options("app:app"), workers=2).timeout_worker_healthcheck == RELOAD_WORKER_HEALTHCHECK_SECONDS + assert ( + pool_config({**_options("app:app"), "timeout_worker_healthcheck": explicit}, workers=2).timeout_worker_healthcheck + == explicit + ) + + def test_a_working_app_loads_as_the_module_defines_it(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: (tmp_path / "reload_working_app.py").write_text("async def app(scope, receive, send):\n return None\n") monkeypatch.syspath_prepend(str(tmp_path)) @@ -141,6 +156,30 @@ def _wait_for_log(log: Path, process: subprocess.Popen[bytes], seen: Callable[[s raise AssertionError(log.read_text()) +def _start_pool(app_dir: Path, port: int, log: Path) -> subprocess.Popen[bytes]: + with log.open("wb") as output: + return subprocess.Popen( + ( + sys.executable, + "-I", + "-c", + POOL_RUNNER, + str(app_dir), + str(port), + str(WORKER_HEALTHCHECK_SECONDS), + str(POOL_WORKERS), + ), + stdout=output, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + + +def _port_refuses_connections(port: int) -> bool: + with socket.socket() as sock: + return sock.connect_ex(("127.0.0.1", port)) != 0 + + def _next_pool(log: Path, process: subprocess.Popen[bytes], port: int, previous: frozenset[int]) -> frozenset[int]: log_text: Final = _wait_for_log(log, process, lambda text: len(_started_workers(text) - previous) == POOL_WORKERS) pool: Final = _started_workers(log_text) - previous @@ -161,22 +200,7 @@ def test_a_save_replaces_every_worker_and_a_broken_save_stops_the_pool_until_the app.write_text(PID_APP) log: Final = tmp_path / "reload.log" port: Final = _free_port() - with log.open("wb") as output: - process: Final = subprocess.Popen( - ( - sys.executable, - "-I", - "-c", - POOL_RUNNER, - str(tmp_path), - str(port), - str(WORKER_HEALTHCHECK_SECONDS), - str(POOL_WORKERS), - ), - stdout=output, - stderr=subprocess.STDOUT, - start_new_session=True, - ) + process: Final = _start_pool(tmp_path, port, log) try: first_pool: Final = _next_pool(log, process, port, frozenset()) app.write_text("async def app(\n") @@ -192,3 +216,38 @@ def test_a_save_replaces_every_worker_and_a_broken_save_stops_the_pool_until_the with contextlib.suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGKILL) process.wait() + + +@pytest.mark.skipif( + not multi_worker_reload_supported(), reason="uvicorn<0.50.0 cannot stop a worker pool on a startup failure" +) +@pytest.mark.timeout(3 * POOL_TIMEOUT_SECONDS) +@pytest.mark.parametrize( + "kill_reloader", + ( + pytest.param(lambda reloader: os.kill(reloader, signal.SIGKILL), id="sigkill-reloader"), + pytest.param(lambda reloader: os.killpg(reloader, signal.SIGHUP), id="terminal-closed"), + pytest.param(lambda reloader: os.killpg(reloader, signal.SIGQUIT), id="ctrl-backslash"), + ), +) +def test_the_pool_stops_and_frees_the_port_when_the_reloader_dies( + tmp_path: Path, kill_reloader: Callable[[int], None] +) -> None: + (tmp_path / "reload_pool_app.py").write_text(PID_APP) + log: Final = tmp_path / "reload.log" + port: Final = _free_port() + process: Final = _start_pool(tmp_path, port, log) + try: + workers: Final = _next_pool(log, process, port, frozenset()) + pool_parent: Final = psutil.Process(min(workers)).ppid() + assert psutil.Process(pool_parent).ppid() == process.pid, log.read_text() + kill_reloader(process.pid) + process.wait(timeout=POOL_TIMEOUT_SECONDS) + deadline: Final = time.monotonic() + POOL_TIMEOUT_SECONDS + while psutil.pid_exists(pool_parent) or not _port_refuses_connections(port): + assert time.monotonic() < deadline, log.read_text() + time.sleep(0.2) + finally: + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + process.wait()