mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
fix(proxy)!: support reload with multiple workers (#45909)
* fix(proxy): support reload with multiple workers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): stop the reload worker pool on a broken edit instead of respawning it Workers now load the app through a factory that exits with uvicorn's STARTUP_FAILURE code on any import error, so the pool stops and the reloader waits for the next save, the same as single-worker reload. Multi-worker reload needs uvicorn>=0.50.0 for that, so older versions keep the single-worker reload with a warning. * test(integration): stop the reload test's proxy through the harness shutdown helper Main's TQ010 rule rejects a raw .terminate() under tests/integration, so the multi-worker reload test now calls OwnedProxy.begin_shutdown(), which sends the same SIGTERM * test(proxy): cover the reload worker pool end to end and edit the config the proxy loads * docs(agents): scope the two-worker reload QA rule to uvicorn proxies that use Redis * fix(proxy): keep the original traceback when a reloaded worker fails to import the app * docs(agents): point both router and coordination Redis at the QA Redis and cover single-worker trees * 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. * fix(proxy): watch the reloader through its process sentinel so the pool also stops on Windows * fix(proxy): carry the graceful shutdown timeout into the reload worker pool * fix(proxy): stop the reload worker pool on Ctrl-\ instead of waiting out a dying worker's healthcheck * test(proxy): skip the startup-failure reload tests on uvicorn without a startup failure code * docs(agents): keep the QA proxy rule to two workers, Redis, and reload * fix(proxy): keep Windows on single-worker reload and pass unlisted uvicorn args to the reload pool * fix(proxy): stop reload workers when their pool parent dies * test(proxy): move the reload pool process tests to the integration suite The unit test guide keeps tests/unit in-process, so the subprocess pool cases now live in tests/integration/configuration and use the harness port allocator and process helpers * docs(agents): keep main's proof rules and inline the multi-worker QA command Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): exit reload workers as startup failures when the app factory raises Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): check the reload pool serves the overridden worker app Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
parent
a1c5829f6b
commit
187f936cb3
6 changed files with 753 additions and 3 deletions
|
|
@ -41,6 +41,8 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM
|
|||
|
||||
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just drop the section
|
||||
|
||||
Run QA proxies with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --num_workers 2 --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`, and swap its `--config` for a copy of `litellm/proxy/dev_config.yaml` that points both `router_settings` (`redis_host`, `redis_port`) and `general_settings.coordination_redis` (`host`, `port`) at one dedicated Redis instance, with credentials when required, so the workers share router state and spend, budget, and rate-limit counters. Check that two worker PIDs are serving, and use the same worker count and Redis setup for the before and after proof. When the PR is about `--run_gunicorn`, `--run_hypercorn`, or `--run_granian`, which `--reload` does not support, drop `--reload`. When it is about behavior without Redis, drop Redis. When the proxy prints `"workers" flag is ignored when reloading is enabled`, that checkout or its uvicorn reloads only one worker, so drop the two-worker check for it. Say in the proof which part you dropped and why
|
||||
|
||||
Simple PRs may skip before/after e2e when the full diff is mechanical and focused checks completely prove it, such as #49030 restoring an existing helper name. Trace affected callers and run relevant tests/static checks at the tip. For a bug fix, establish that the regression catches the original defect. Size alone never qualifies. New behavior, API/provider semantics, auth, billing, persistence, concurrency or uncertain integration risk need live e2e. In Screenshots / Proof of Fix, record `E2E exception at <tip SHA>`, scope, rationale, commands and results. Tests count as proof only for this exception. Reassess changed diffs, retain CI/bot gates, and honor explicit e2e requests (2026-10-28, unnecessary live QA for #49030)
|
||||
|
||||
For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce
|
||||
|
|
|
|||
172
litellm/proxy/dev_reload.py
Normal file
172
litellm/proxy/dev_reload.py
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
import importlib
|
||||
import logging
|
||||
import multiprocessing
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import threading
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from functools import partial
|
||||
from inspect import signature
|
||||
from multiprocessing.process import BaseProcess
|
||||
from socket import socket
|
||||
from typing import Final, Literal
|
||||
|
||||
from pydantic import ConfigDict, ImportString, JsonValue, TypeAdapter, with_config
|
||||
from typing_extensions import ReadOnly, Required, TypedDict
|
||||
from uvicorn import Config, Server
|
||||
from uvicorn import config as uvicorn_config
|
||||
from uvicorn.supervisors import ChangeReload, Multiprocess
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow"))
|
||||
class ServerOptions(TypedDict, total=False):
|
||||
host: ReadOnly[str]
|
||||
port: ReadOnly[int]
|
||||
server_header: ReadOnly[bool]
|
||||
timeout_graceful_shutdown: ReadOnly[float]
|
||||
log_config: ReadOnly[ # mutable-ok: uvicorn.Config requires a concrete logging dict
|
||||
str | dict[str, JsonValue] | None
|
||||
]
|
||||
log_level: ReadOnly[str | int]
|
||||
timeout_keep_alive: ReadOnly[int]
|
||||
timeout_worker_healthcheck: ReadOnly[int]
|
||||
limit_max_requests: ReadOnly[int]
|
||||
limit_max_requests_jitter: ReadOnly[int]
|
||||
limit_concurrency: ReadOnly[int]
|
||||
ssl_keyfile: ReadOnly[str]
|
||||
ssl_certfile: ReadOnly[str]
|
||||
loop: ReadOnly[Literal["none", "auto", "asyncio", "uvloop"]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow"))
|
||||
class ReloadOptions(ServerOptions):
|
||||
app: Required[ReadOnly[str]]
|
||||
reload: ReadOnly[bool]
|
||||
reload_dirs: ReadOnly[tuple[str, ...]]
|
||||
reload_includes: ReadOnly[tuple[str, ...]]
|
||||
|
||||
|
||||
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_ONLY_KEYS: Final = frozenset(ReloadOptions.__annotations__) - frozenset(ServerOptions.__annotations__)
|
||||
|
||||
|
||||
def multi_worker_reload_supported(platform: str = sys.platform) -> bool:
|
||||
return platform != "win32" and hasattr(uvicorn_config, "STARTUP_FAILURE")
|
||||
|
||||
|
||||
def server_options(options: ReloadOptions) -> ServerOptions:
|
||||
return SERVER_OPTIONS.validate_python(
|
||||
{key: value for key, value in options.items() if key not in RELOADER_ONLY_KEYS}
|
||||
)
|
||||
|
||||
|
||||
def load_app_or_exit(import_path: str, factory: bool = False) -> Callable[..., object]:
|
||||
module_name: Final = import_path.partition(":")[0]
|
||||
try:
|
||||
importlib.import_module(module_name)
|
||||
app: Final = ASGI_APP.validate_python(import_path)
|
||||
return ASGI_APP.validate_python(app()) if factory else app
|
||||
except Exception:
|
||||
LOGGER.exception("Error loading ASGI app %s", import_path)
|
||||
sys.exit(uvicorn_config.STARTUP_FAILURE)
|
||||
|
||||
|
||||
def stop_worker_when_pool_exits(pool_parent: BaseProcess) -> None:
|
||||
pool_parent.join()
|
||||
LOGGER.info("Worker pool process [%s] exited, stopping worker [%s].", pool_parent.pid, os.getpid())
|
||||
os.kill(os.getpid(), signal.SIGTERM)
|
||||
|
||||
|
||||
def load_worker_app(import_path: str, worker_app_factory: bool) -> Callable[..., object]:
|
||||
pool_parent: Final = multiprocessing.parent_process()
|
||||
if pool_parent is not None:
|
||||
threading.Thread(target=stop_worker_when_pool_exits, args=(pool_parent,), daemon=True).start()
|
||||
return load_app_or_exit(import_path, worker_app_factory)
|
||||
|
||||
|
||||
def pool_config(options: ReloadOptions, workers: int, worker_overrides: Mapping[str, object] | None = None) -> Config:
|
||||
overrides: Final = worker_overrides or {}
|
||||
worker_app: Final = TypeAdapter(str).validate_python(overrides.get("app", options["app"]))
|
||||
worker_app_factory: Final = TypeAdapter(bool).validate_python(overrides.get("factory", False))
|
||||
return Config(
|
||||
partial(load_worker_app, worker_app, worker_app_factory),
|
||||
factory=True,
|
||||
workers=workers,
|
||||
**SERVER_OPTIONS.validate_python(
|
||||
{
|
||||
"timeout_worker_healthcheck": RELOAD_WORKER_HEALTHCHECK_SECONDS,
|
||||
**server_options(options),
|
||||
**{key: value for key, value in overrides.items() if key not in ("app", "factory")},
|
||||
}
|
||||
), # pyright: ignore[reportArgumentType] # uvicorn types the float graceful shutdown timeout as int
|
||||
)
|
||||
|
||||
|
||||
def stop_pool_when_reloader_exits(reloader: BaseProcess, pool_should_exit: threading.Event) -> None:
|
||||
reloader.join()
|
||||
LOGGER.info("Reloader process [%s] exited, stopping the worker pool.", reloader.pid)
|
||||
pool_should_exit.set()
|
||||
|
||||
|
||||
class ReloadWorkerPool(Multiprocess):
|
||||
quit_received: bool = False
|
||||
|
||||
def handle_quit(self) -> None:
|
||||
LOGGER.info("Received SIGQUIT, exiting.")
|
||||
self.quit_received = True
|
||||
self.should_exit.set()
|
||||
|
||||
def terminate_all(self) -> None:
|
||||
quit_signal: Final = getattr(signal, "SIGQUIT", None)
|
||||
if quit_signal is not None and quit_signal in self.signal_queue:
|
||||
self.handle_quit()
|
||||
if not self.quit_received:
|
||||
super().terminate_all()
|
||||
return
|
||||
for process in self.processes:
|
||||
process.kill()
|
||||
|
||||
|
||||
def _run_worker_pool(config: Config, sockets: Sequence[socket] | None) -> None:
|
||||
shared_sockets: Final = list(sockets or ())
|
||||
pool: Final = (
|
||||
ReloadWorkerPool(
|
||||
config=config,
|
||||
target=Server(config=config).run, # pyright: ignore[reportCallIssue] # Older uvicorn versions require target
|
||||
sockets=shared_sockets,
|
||||
)
|
||||
if "target" in signature(Multiprocess).parameters
|
||||
else ReloadWorkerPool(config=config, sockets=shared_sockets)
|
||||
)
|
||||
reloader: Final = multiprocessing.parent_process()
|
||||
if reloader is not None:
|
||||
threading.Thread(target=stop_pool_when_reloader_exits, args=(reloader, pool.should_exit), daemon=True).start()
|
||||
pool.run()
|
||||
|
||||
|
||||
def run_multi_worker_reload(
|
||||
options: ReloadOptions, workers: int, worker_overrides: Mapping[str, object] | None = None
|
||||
) -> None:
|
||||
worker_config: Final = pool_config(options, workers, worker_overrides)
|
||||
reload_config: Final = Config(
|
||||
options["app"],
|
||||
workers=1,
|
||||
reload=options["reload"],
|
||||
reload_dirs=list(options["reload_dirs"]),
|
||||
reload_includes=list(options["reload_includes"]),
|
||||
**server_options(options), # pyright: ignore[reportArgumentType] # uvicorn types the float graceful shutdown timeout as int
|
||||
)
|
||||
shared_socket: Final = worker_config.bind_socket()
|
||||
try:
|
||||
ChangeReload(
|
||||
config=reload_config,
|
||||
target=partial(_run_worker_pool, worker_config),
|
||||
sockets=[shared_socket],
|
||||
).run()
|
||||
finally:
|
||||
shared_socket.close()
|
||||
|
|
@ -1193,7 +1193,7 @@ class ProxyInitializationHelpers:
|
|||
"--reload",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help="Enable uvicorn hot reload (dev only). Also reloads when the --config YAML file changes. Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.",
|
||||
help="Enable uvicorn hot reload (dev only), including every --num_workers worker on uvicorn>=0.50.0 outside Windows. Also reloads when the --config YAML or .env changes. Incompatible with --run_gunicorn, --run_hypercorn, and --run_granian.",
|
||||
)
|
||||
@click.option(
|
||||
"--prometheus_metrics_port",
|
||||
|
|
@ -1785,9 +1785,28 @@ def run_server(
|
|||
if num_workers > 1:
|
||||
start_query_engine_reaper()
|
||||
if num_workers > 1 or reload:
|
||||
uvicorn_args.update(
|
||||
ProxyInitializationHelpers.child_process_uvicorn_overrides(uvicorn_args.get("log_config"))
|
||||
worker_overrides: Final = ProxyInitializationHelpers.child_process_uvicorn_overrides(
|
||||
uvicorn_args.get("log_config")
|
||||
)
|
||||
if reload and num_workers > 1:
|
||||
from litellm.proxy.dev_reload import (
|
||||
RELOAD_OPTIONS,
|
||||
multi_worker_reload_supported,
|
||||
run_multi_worker_reload,
|
||||
)
|
||||
|
||||
if multi_worker_reload_supported():
|
||||
run_multi_worker_reload(
|
||||
RELOAD_OPTIONS.validate_python(uvicorn_args),
|
||||
workers=TypeAdapter(int).validate_python(num_workers),
|
||||
worker_overrides=worker_overrides,
|
||||
)
|
||||
return
|
||||
print(
|
||||
"\033[1;33mLiteLLM Proxy: --reload with --num_workers>1 requires uvicorn>=0.50.0 "
|
||||
"and an OS other than Windows, so only a single worker reloads.\033[0m"
|
||||
)
|
||||
uvicorn_args.update(worker_overrides)
|
||||
uvicorn.run(
|
||||
**uvicorn_args,
|
||||
workers=num_workers,
|
||||
|
|
|
|||
95
tests/integration/configuration/test_multi_worker_reload.py
Normal file
95
tests/integration/configuration/test_multi_worker_reload.py
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
import json
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import psutil
|
||||
import pytest
|
||||
|
||||
from tests.integration._support.client import Gateway, eventually, object_value
|
||||
from tests.integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
|
||||
|
||||
STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
|
||||
|
||||
|
||||
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(
|
||||
{
|
||||
"model_list": [{"model_name": model, "litellm_params": {"model": "openai/reload-test"}}],
|
||||
"general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY"},
|
||||
"router_settings": {"redis_host": "os.environ/REDIS_HOST", "redis_port": "os.environ/REDIS_PORT"},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _assert_replaced(owned: OwnedProxy, previous: frozenset[int], workers: int) -> frozenset[int]:
|
||||
started: Final = eventually(
|
||||
lambda: _worker_pids(owned),
|
||||
lambda pids: len(pids - previous) == workers,
|
||||
seconds=graceful_stop_seconds(),
|
||||
)
|
||||
assert all(not psutil.pid_exists(pid) for pid in previous), owned.log.read_text()
|
||||
assert owned.process.poll() is None, owned.log.read_text()
|
||||
assert len(started) == len(previous) + workers, owned.log.read_text()
|
||||
return started
|
||||
|
||||
|
||||
def _assert_models(gateway: Gateway) -> None:
|
||||
data: Final = gateway.get("/v1/models")["data"]
|
||||
assert isinstance(data, list), data
|
||||
assert tuple(object_value(model)["id"] for model in data) == ("after-reload",)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("workers", (1, 2))
|
||||
@pytest.mark.timeout(6 * graceful_stop_seconds())
|
||||
def test_reload_replaces_every_worker_and_shutdown_leaves_none(
|
||||
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, workers: int
|
||||
) -> None:
|
||||
root: Final = Path(__file__).resolve().parents[3]
|
||||
monkeypatch.setenv("INTEGRATION_PROXY_ROOT", str(tmp_path))
|
||||
config: Final = tmp_path / "config.yaml"
|
||||
source: Final = tmp_path / "reload_source.py"
|
||||
env: Final = tmp_path / ".env"
|
||||
config.write_text(_config("before-reload"))
|
||||
source.write_text("VALUE = 1\n")
|
||||
env.write_text("RELOAD_TEST_VALUE=before\n")
|
||||
with owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{
|
||||
"PYTHONPATH": os.pathsep.join((str(root), str(root / "tests"), os.environ.get("PYTHONPATH", ""))),
|
||||
"STORE_MODEL_IN_DB": "False",
|
||||
},
|
||||
config=config,
|
||||
remove_environment=("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "PYTHON_DOTENV_DISABLED"),
|
||||
workers=workers,
|
||||
database_setup=(),
|
||||
extra_arguments=("--reload",),
|
||||
) as owned:
|
||||
initial: Final = eventually(
|
||||
lambda: _worker_pids(owned), lambda pids: len(pids) == workers, seconds=graceful_stop_seconds()
|
||||
)
|
||||
_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)
|
||||
source.write_text("VALUE = 2\n")
|
||||
after_source: Final = _assert_replaced(owned, after_config, workers)
|
||||
env.write_text("RELOAD_TEST_VALUE=after\n")
|
||||
after_env: Final = _assert_replaced(owned, after_source, workers)
|
||||
response: Final = owned.gateway.request("GET", "/health/readiness")
|
||||
assert response.status_code == 200, response.text
|
||||
owned.begin_shutdown()
|
||||
owned.process.wait(timeout=graceful_stop_seconds())
|
||||
assert all(not psutil.pid_exists(pid) for pid in after_env), owned.log.read_text()
|
||||
assert owned.process.returncode == 0, owned.log.read_text()
|
||||
246
tests/integration/configuration/test_reload_worker_pool.py
Normal file
246
tests/integration/configuration/test_reload_worker_pool.py
Normal file
|
|
@ -0,0 +1,246 @@
|
|||
import multiprocessing
|
||||
import re
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import psutil
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.dev_reload import multi_worker_reload_supported, stop_pool_when_reloader_exits
|
||||
from tests.integration._support.client import eventually
|
||||
from tests.integration._support.ports import allocate_port
|
||||
from tests.integration._support.process import group_members, kill_stray_process, signal_group, stop_processes
|
||||
|
||||
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()})
|
||||
"""
|
||||
REQUIRES_UVICORN_STARTUP_FAILURE: Final = pytest.mark.skipif(
|
||||
not multi_worker_reload_supported(), reason="uvicorn<0.50.0 has no startup failure exit code"
|
||||
)
|
||||
CORE_DUMP_SECONDS: Final = 15
|
||||
SLOW_CORE_DUMP_ON_QUIT: Final = f"""
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
|
||||
|
||||
def freeze_until_killed(signum, frame):
|
||||
subprocess.Popen(("sh", "-c", f"sleep {CORE_DUMP_SECONDS}; kill -9 {{os.getpid()}}"))
|
||||
os.kill(os.getpid(), signal.SIGSTOP)
|
||||
|
||||
|
||||
signal.signal(signal.SIGQUIT, freeze_until_killed)
|
||||
"""
|
||||
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 _started_workers(log_text: str) -> frozenset[int]:
|
||||
return frozenset(int(match.group(1)) for match in WORKER_STARTED.finditer(log_text))
|
||||
|
||||
|
||||
def _until(done: Callable[[], bool], log: Path, seconds: float = POOL_TIMEOUT_SECONDS) -> None:
|
||||
eventually(lambda: (done(), log.read_text()), lambda observed: observed[0], seconds=seconds)
|
||||
|
||||
|
||||
def _wait_for_log(log: Path, process: subprocess.Popen[bytes], seen: Callable[[str], bool]) -> str:
|
||||
log_text, _ = eventually(
|
||||
lambda: (log.read_text(), process.poll()),
|
||||
lambda observed: seen(observed[0]) or observed[1] is not None,
|
||||
seconds=POOL_TIMEOUT_SECONDS,
|
||||
)
|
||||
assert seen(log_text), log_text
|
||||
return log_text
|
||||
|
||||
|
||||
def _start_pool(app_dir: Path, port: int, log: Path, runner: str = POOL_RUNNER) -> subprocess.Popen[bytes]:
|
||||
with log.open("wb") as output:
|
||||
return subprocess.Popen(
|
||||
(
|
||||
sys.executable,
|
||||
"-I",
|
||||
"-c",
|
||||
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
|
||||
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
|
||||
|
||||
|
||||
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()
|
||||
pool_should_exit: Final = threading.Event()
|
||||
threading.Thread(target=stop_pool_when_reloader_exits, args=(reloader, pool_should_exit), daemon=True).start()
|
||||
try:
|
||||
assert not pool_should_exit.wait(0.5)
|
||||
kill_stray_process(reloader.pid)
|
||||
assert pool_should_exit.wait(POOL_TIMEOUT_SECONDS)
|
||||
finally:
|
||||
if reloader.is_alive():
|
||||
kill_stray_process(reloader.pid)
|
||||
reloader.join()
|
||||
|
||||
|
||||
@REQUIRES_UVICORN_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 = allocate_port()
|
||||
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")
|
||||
_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)
|
||||
signal_group(process.pid, 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:
|
||||
stop_processes(group_members(process.pid))
|
||||
process.wait()
|
||||
|
||||
|
||||
@REQUIRES_UVICORN_STARTUP_FAILURE
|
||||
@pytest.mark.timeout(3 * POOL_TIMEOUT_SECONDS)
|
||||
@pytest.mark.parametrize(
|
||||
("kill_reloader", "slow_core_dump", "shutdown_timeout"),
|
||||
(
|
||||
pytest.param(kill_stray_process, False, POOL_TIMEOUT_SECONDS, id="sigkill-reloader"),
|
||||
pytest.param(
|
||||
lambda reloader: signal_group(reloader, signal.SIGHUP), False, POOL_TIMEOUT_SECONDS, id="terminal-closed"
|
||||
),
|
||||
pytest.param(lambda reloader: signal_group(reloader, signal.SIGQUIT), True, 10, id="ctrl-backslash"),
|
||||
),
|
||||
)
|
||||
def test_the_pool_stops_and_frees_the_port_when_the_reloader_dies(
|
||||
tmp_path: Path, kill_reloader: Callable[[int], None], slow_core_dump: bool, shutdown_timeout: float
|
||||
) -> None:
|
||||
app_source: Final = PID_APP + (SLOW_CORE_DUMP_ON_QUIT if slow_core_dump else "")
|
||||
(tmp_path / "reload_pool_app.py").write_text(app_source)
|
||||
log: Final = tmp_path / "reload.log"
|
||||
port: Final = allocate_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)
|
||||
_until(
|
||||
lambda: not psutil.pid_exists(pool_parent) and _port_refuses_connections(port),
|
||||
log,
|
||||
seconds=shutdown_timeout,
|
||||
)
|
||||
finally:
|
||||
stop_processes(group_members(process.pid))
|
||||
process.wait()
|
||||
|
||||
|
||||
@REQUIRES_UVICORN_STARTUP_FAILURE
|
||||
@pytest.mark.timeout(3 * POOL_TIMEOUT_SECONDS)
|
||||
def test_a_killed_pool_parent_takes_its_workers_with_it_so_the_next_save_serves_only_new_workers(
|
||||
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 = allocate_port()
|
||||
process: Final = _start_pool(tmp_path, port, log)
|
||||
try:
|
||||
orphans: Final = _next_pool(log, process, port, frozenset())
|
||||
kill_stray_process(psutil.Process(min(orphans)).ppid())
|
||||
_until(lambda: not any(psutil.pid_exists(pid) for pid in orphans), log)
|
||||
app.write_text(f"{PID_APP}\n")
|
||||
_next_pool(log, process, port, orphans)
|
||||
finally:
|
||||
stop_processes(group_members(process.pid))
|
||||
process.wait()
|
||||
|
||||
|
||||
@REQUIRES_UVICORN_STARTUP_FAILURE
|
||||
@pytest.mark.timeout(3 * POOL_TIMEOUT_SECONDS)
|
||||
def test_ctrl_backslash_frees_the_port_without_waiting_out_the_worker_healthcheck_when_core_dumps_are_slow(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
(tmp_path / "reload_pool_app.py").write_text(PID_APP + SLOW_CORE_DUMP_ON_QUIT)
|
||||
log: Final = tmp_path / "reload.log"
|
||||
port: Final = allocate_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()
|
||||
signal_group(process.pid, signal.SIGQUIT)
|
||||
_until(
|
||||
lambda: (
|
||||
process.poll() is not None and not psutil.pid_exists(pool_parent) and _port_refuses_connections(port)
|
||||
),
|
||||
log,
|
||||
seconds=WORKER_HEALTHCHECK_SECONDS / 2,
|
||||
)
|
||||
finally:
|
||||
stop_processes(group_members(process.pid))
|
||||
process.wait()
|
||||
216
tests/unit/proxy/test_dev_reload.py
Normal file
216
tests/unit/proxy/test_dev_reload.py
Normal file
|
|
@ -0,0 +1,216 @@
|
|||
import asyncio
|
||||
import importlib
|
||||
import logging
|
||||
import logging.handlers
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from starlette.types import Message, Scope
|
||||
from uvicorn import config as uvicorn_config
|
||||
|
||||
from litellm.proxy.dev_reload import (
|
||||
RELOAD_OPTIONS,
|
||||
RELOAD_WORKER_HEALTHCHECK_SECONDS,
|
||||
ReloadOptions,
|
||||
load_app_or_exit,
|
||||
load_worker_app,
|
||||
multi_worker_reload_supported,
|
||||
pool_config,
|
||||
)
|
||||
from litellm.proxy.proxy_cli import ProxyInitializationHelpers
|
||||
|
||||
UNCHANGED_LOGGING: Final = {"version": 1, "incremental": True}
|
||||
POOL_WORKERS: Final = 2
|
||||
REQUIRES_UVICORN_STARTUP_FAILURE: Final = pytest.mark.skipif(
|
||||
not multi_worker_reload_supported(), reason="uvicorn<0.50.0 has no startup failure exit code"
|
||||
)
|
||||
|
||||
|
||||
def _options(app: str) -> ReloadOptions:
|
||||
return {
|
||||
"app": app,
|
||||
"log_config": UNCHANGED_LOGGING,
|
||||
"reload": True,
|
||||
"reload_dirs": (),
|
||||
"reload_includes": ("*.py",),
|
||||
}
|
||||
|
||||
|
||||
@REQUIRES_UVICORN_STARTUP_FAILURE
|
||||
@pytest.mark.parametrize(
|
||||
("module", "source"),
|
||||
(
|
||||
("reload_syntax_error_app", "def app(\n"),
|
||||
("reload_name_error_app", "app = undefined_name\n"),
|
||||
("reload_runtime_error_app", "raise RuntimeError('broken edit')\n"),
|
||||
("reload_missing_attribute_app", "value = 1\n"),
|
||||
),
|
||||
)
|
||||
def test_a_broken_app_exits_the_worker_as_a_startup_failure(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, module: str, source: str
|
||||
) -> None:
|
||||
(tmp_path / f"{module}.py").write_text(source)
|
||||
monkeypatch.syspath_prepend(str(tmp_path))
|
||||
with pytest.raises(SystemExit) as exit_info:
|
||||
pool_config(_options(f"{module}:app"), workers=2).load()
|
||||
assert exit_info.value.code == uvicorn_config.STARTUP_FAILURE
|
||||
|
||||
|
||||
@REQUIRES_UVICORN_STARTUP_FAILURE
|
||||
@pytest.mark.parametrize(
|
||||
("module", "source"),
|
||||
(
|
||||
(
|
||||
"reload_factory_missing_module_app",
|
||||
"def factory():\n import a_module_that_does_not_exist\n",
|
||||
),
|
||||
(
|
||||
"reload_factory_runtime_error_app",
|
||||
"def factory():\n raise RuntimeError('broken edit')\n",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_a_broken_factory_exits_the_worker_as_a_startup_failure(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, module: str, source: str
|
||||
) -> None:
|
||||
(tmp_path / f"{module}.py").write_text(source)
|
||||
monkeypatch.syspath_prepend(str(tmp_path))
|
||||
with pytest.raises(SystemExit) as exit_info:
|
||||
pool_config(
|
||||
_options(f"{module}:app"),
|
||||
workers=POOL_WORKERS,
|
||||
worker_overrides={"app": f"{module}:factory", "factory": True},
|
||||
).load()
|
||||
assert exit_info.value.code == uvicorn_config.STARTUP_FAILURE
|
||||
|
||||
|
||||
@REQUIRES_UVICORN_STARTUP_FAILURE
|
||||
@pytest.mark.parametrize(
|
||||
("module", "source"),
|
||||
(
|
||||
("reload_import_error_app", "import os\nfrom os import missing_name\n"),
|
||||
("reload_value_error_app", "import os\nraise ValueError('broken edit')\n"),
|
||||
),
|
||||
)
|
||||
def test_a_broken_import_logs_the_line_that_raised(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, module: str, source: str
|
||||
) -> None:
|
||||
app_file: Final = tmp_path / f"{module}.py"
|
||||
app_file.write_text(source)
|
||||
monkeypatch.syspath_prepend(str(tmp_path))
|
||||
uvicorn_logger: Final = logging.getLogger("uvicorn.error")
|
||||
captured: Final = logging.handlers.BufferingHandler(capacity=POOL_WORKERS)
|
||||
uvicorn_logger.addHandler(captured)
|
||||
try:
|
||||
with pytest.raises(SystemExit):
|
||||
load_app_or_exit(f"{module}:app")
|
||||
finally:
|
||||
uvicorn_logger.removeHandler(captured)
|
||||
(record,) = captured.buffer
|
||||
assert record.exc_info is not None
|
||||
raised_at: Final = frozenset((frame.filename, frame.lineno) for frame in traceback.extract_tb(record.exc_info[2]))
|
||||
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_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_reload_pool_passes_a_uvicorn_arg_it_does_not_list_through_to_the_workers() -> None:
|
||||
config: Final = pool_config(RELOAD_OPTIONS.validate_python({**_options("app:app"), "backlog": 123}), workers=POOL_WORKERS)
|
||||
assert (config.backlog, config.reload) == (123, False)
|
||||
|
||||
|
||||
def test_the_reload_pool_serves_the_overridden_worker_app(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
module: Final = "reload_worker_override_app"
|
||||
(tmp_path / f"{module}.py").write_text(
|
||||
"async def base_app(scope, receive, send):\n"
|
||||
" await receive()\n"
|
||||
" await send({'type': 'http.response.start', 'status': 200, 'headers': []})\n"
|
||||
" await send({'type': 'http.response.body', 'body': b'base app', 'more_body': False})\n"
|
||||
"\n"
|
||||
"async def overridden_app(scope, receive, send):\n"
|
||||
" await receive()\n"
|
||||
" await send({'type': 'http.response.start', 'status': 200, 'headers': []})\n"
|
||||
" await send({'type': 'http.response.body', 'body': b'overridden app', 'more_body': False})\n"
|
||||
"\n"
|
||||
"def base_factory():\n"
|
||||
" return base_app\n"
|
||||
"\n"
|
||||
"def factory():\n"
|
||||
" return overridden_app\n"
|
||||
)
|
||||
monkeypatch.syspath_prepend(str(tmp_path))
|
||||
config: Final = pool_config(
|
||||
_options(f"{module}:base_factory"),
|
||||
workers=POOL_WORKERS,
|
||||
worker_overrides={"app": f"{module}:factory", "factory": True, "log_config": None},
|
||||
)
|
||||
assert config.log_config is None
|
||||
config.load()
|
||||
loaded_app: Final = config.loaded_app
|
||||
assert loaded_app is not None
|
||||
scope: Final[Scope] = {"type": "http", "method": "GET", "path": "/", "headers": []}
|
||||
|
||||
async def run_app() -> tuple[Message, ...]:
|
||||
messages: Final[asyncio.Queue[Message]] = asyncio.Queue()
|
||||
|
||||
async def receive() -> Message:
|
||||
return {"type": "http.request", "body": b"", "more_body": False}
|
||||
|
||||
async def send(message: Message) -> None:
|
||||
await messages.put(message)
|
||||
|
||||
await loaded_app(scope, receive, send)
|
||||
return tuple(messages.get_nowait() for _ in range(messages.qsize()))
|
||||
|
||||
messages: Final = asyncio.run(run_app())
|
||||
assert messages == (
|
||||
{"type": "http.response.start", "status": 200, "headers": []},
|
||||
{"type": "http.response.body", "body": b"overridden app", "more_body": False},
|
||||
)
|
||||
|
||||
|
||||
def test_load_worker_app_calls_the_overridden_factory(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
module: Final = "reload_worker_app_factory"
|
||||
(tmp_path / f"{module}.py").write_text(
|
||||
"def app(scope, receive, send):\n return None\n\ndef factory():\n return app\n"
|
||||
)
|
||||
monkeypatch.syspath_prepend(str(tmp_path))
|
||||
app: Final = importlib.import_module(module).app
|
||||
assert load_worker_app(f"{module}:factory", worker_app_factory=True) is app
|
||||
|
||||
|
||||
def test_windows_keeps_single_worker_reload_whatever_the_uvicorn_version() -> None:
|
||||
assert not multi_worker_reload_supported("win32")
|
||||
|
||||
|
||||
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))
|
||||
assert load_app_or_exit("reload_working_app:app") is importlib.import_module("reload_working_app").app
|
||||
Loading…
Add table
Reference in a new issue