test(proxy): cover the reload worker pool end to end and edit the config the proxy loads

tests/integration/configuration/test_multi_worker_reload.py

@@ -17,6 +17,11 @@ def _worker_pids(owned: OwnedProxy) -> frozenset[int]:
     return frozenset(int(match.group(1)) for match in STARTED_WORKER.finditer(owned.log.read_text()))
 
 
+def _launched_config(owned: OwnedProxy) -> Path:
+    arguments: Final = psutil.Process(owned.process.pid).cmdline()
+    return Path(arguments[arguments.index("--config") + 1])
+
+
 def _config(model: str) -> str:
     return json.dumps(
         {
@@ -74,7 +79,7 @@ def test_reload_replaces_every_worker_and_shutdown_leaves_none(
         initial: Final = eventually(
             lambda: _worker_pids(owned), lambda pids: len(pids) == workers, seconds=graceful_stop_seconds()
         )
-        config.write_text(_config("after-reload"))
+        _launched_config(owned).write_text(_config("after-reload"))
         after_config: Final = _assert_replaced(owned, initial, workers)
         for _ in range(2 * workers):
             _assert_models(owned.gateway)

tests/unit/proxy/test_dev_reload.py

@@ -1,13 +1,58 @@
+import contextlib
 import importlib
+import os
+import re
+import signal
+import socket
+import subprocess
+import sys
+import time
+from collections.abc import Callable
 from pathlib import Path
 from typing import Final
 
+import httpx
+import psutil
 import pytest
 from uvicorn import config as uvicorn_config
 
-from litellm.proxy.dev_reload import ReloadOptions, load_app_or_exit, pool_config
+from litellm.proxy.dev_reload import ReloadOptions, load_app_or_exit, multi_worker_reload_supported, pool_config
 
 UNCHANGED_LOGGING: Final = {"version": 1, "incremental": True}
+POOL_WORKERS: Final = 2
+POOL_TIMEOUT_SECONDS: Final = 120
+WORKER_HEALTHCHECK_SECONDS: Final = 30
+WORKER_STARTED: Final = re.compile(r"Started server process \[(\d+)\]")
+POOL_STOPPED: Final = "failed to start, stopping the parent process"
+APP_LOAD_ERROR: Final = "Error loading ASGI app"
+PID_APP: Final = """import os
+
+
+async def app(scope, receive, send):
+    if scope["type"] != "http":
+        return
+    await send({"type": "http.response.start", "status": 200, "headers": []})
+    await send({"type": "http.response.body", "body": str(os.getpid()).encode()})
+"""
+POOL_RUNNER: Final = """import sys
+
+sys.path.insert(0, sys.argv[1])
+
+from litellm.proxy.dev_reload import run_multi_worker_reload
+
+run_multi_worker_reload(
+    {
+        "app": "reload_pool_app:app",
+        "host": "127.0.0.1",
+        "port": int(sys.argv[2]),
+        "timeout_worker_healthcheck": int(sys.argv[3]),
+        "reload": True,
+        "reload_dirs": (sys.argv[1],),
+        "reload_includes": ("*.py",),
+    },
+    workers=int(sys.argv[4]),
+)
+"""
 
 
 def _options(app: str) -> ReloadOptions:
@@ -43,3 +88,77 @@ def test_a_working_app_loads_as_the_module_defines_it(tmp_path: Path, monkeypatc
     (tmp_path / "reload_working_app.py").write_text("async def app(scope, receive, send):\n    return None\n")
     monkeypatch.syspath_prepend(str(tmp_path))
     assert load_app_or_exit("reload_working_app:app") is importlib.import_module("reload_working_app").app
+
+
+def _free_port() -> int:
+    with socket.socket() as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()[1]
+
+
+def _started_workers(log_text: str) -> frozenset[int]:
+    return frozenset(int(match.group(1)) for match in WORKER_STARTED.finditer(log_text))
+
+
+def _wait_for_log(log: Path, process: subprocess.Popen[bytes], seen: Callable[[str], bool]) -> str:
+    deadline: Final = time.monotonic() + POOL_TIMEOUT_SECONDS
+    while time.monotonic() < deadline:
+        log_text: Final = log.read_text()
+        if seen(log_text):
+            return log_text
+        assert process.poll() is None, log_text
+        time.sleep(0.2)
+    raise AssertionError(log.read_text())
+
+
+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
+    assert not any(psutil.pid_exists(pid) for pid in previous), log_text
+    served: Final = httpx.get(f"http://127.0.0.1:{port}/", timeout=POOL_TIMEOUT_SECONDS)
+    assert int(served.text) in pool, log_text
+    return pool
+
+
+@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(5 * POOL_TIMEOUT_SECONDS)
+def test_a_save_replaces_every_worker_and_a_broken_save_stops_the_pool_until_the_next_good_one(
+    tmp_path: Path,
+) -> None:
+    app: Final = tmp_path / "reload_pool_app.py"
+    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,
+        )
+    try:
+        first_pool: Final = _next_pool(log, process, port, frozenset())
+        app.write_text("async def app(\n")
+        _wait_for_log(log, process, lambda text: POOL_STOPPED in text)
+        assert not any(psutil.pid_exists(pid) for pid in first_pool), log.read_text()
+        app.write_text(PID_APP)
+        last_pool: Final = _next_pool(log, process, port, first_pool)
+        process.send_signal(signal.SIGINT)
+        assert process.wait(timeout=POOL_TIMEOUT_SECONDS) == 0, log.read_text()
+        assert not any(psutil.pid_exists(pid) for pid in last_pool), log.read_text()
+        assert 1 <= log.read_text().count(APP_LOAD_ERROR) <= POOL_WORKERS, log.read_text()
+    finally:
+        with contextlib.suppress(ProcessLookupError):
+            os.killpg(process.pid, signal.SIGKILL)
+        process.wait()