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()