mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
Merge dc68ca731b into 045fb7ba84
This commit is contained in:
commit
332b0d4fc1
5 changed files with 567 additions and 3 deletions
|
|
@ -39,7 +39,9 @@ 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
|
||||
|
||||
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`. The Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into. I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it
|
||||
Run QA proxies with the command below, which passes `--num_workers 2 --reload`, 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
|
||||
|
||||
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it 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`. The Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into. I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it
|
||||
|
||||
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
|
||||
- don't use emojis
|
||||
|
|
|
|||
122
litellm/proxy/dev_reload.py
Normal file
122
litellm/proxy/dev_reload.py
Normal file
|
|
@ -0,0 +1,122 @@
|
|||
import importlib
|
||||
import logging
|
||||
import multiprocessing
|
||||
import sys
|
||||
import threading
|
||||
from collections.abc import Callable, 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
|
||||
|
||||
|
||||
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]
|
||||
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="forbid"))
|
||||
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
|
||||
|
||||
|
||||
def multi_worker_reload_supported() -> bool:
|
||||
return hasattr(uvicorn_config, "STARTUP_FAILURE")
|
||||
|
||||
|
||||
def load_app_or_exit(import_path: str) -> Callable[..., object]:
|
||||
module_name: Final = import_path.partition(":")[0]
|
||||
try:
|
||||
importlib.import_module(module_name)
|
||||
return ASGI_APP.validate_python(import_path)
|
||||
except Exception:
|
||||
LOGGER.exception("Error loading ASGI app %s", import_path)
|
||||
sys.exit(uvicorn_config.STARTUP_FAILURE)
|
||||
|
||||
|
||||
def pool_config(options: ReloadOptions, workers: int) -> Config:
|
||||
return Config(
|
||||
partial(load_app_or_exit, options["app"]),
|
||||
factory=True,
|
||||
workers=workers,
|
||||
**SERVER_OPTIONS.validate_python({"timeout_worker_healthcheck": RELOAD_WORKER_HEALTHCHECK_SECONDS, **options}), # 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):
|
||||
def handle_quit(self) -> None:
|
||||
LOGGER.info("Received SIGQUIT, exiting.")
|
||||
self.should_exit.set()
|
||||
|
||||
|
||||
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) -> None:
|
||||
worker_config: Final = pool_config(options, workers)
|
||||
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.validate_python(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()
|
||||
|
|
@ -15,7 +15,7 @@ import click
|
|||
import httpx
|
||||
from click.core import ParameterSource
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import ConfigDict
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.constants import DEFAULT_NUM_WORKERS_LITELLM_PROXY
|
||||
|
|
@ -1070,7 +1070,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. Also reloads when the --config YAML or .env changes. Incompatible with --run_gunicorn, --run_hypercorn, and --run_granian.",
|
||||
)
|
||||
@click.option(
|
||||
"--prometheus_metrics_port",
|
||||
|
|
@ -1641,6 +1641,23 @@ def run_server(
|
|||
|
||||
if num_workers > 1:
|
||||
start_query_engine_reaper()
|
||||
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),
|
||||
)
|
||||
return
|
||||
print(
|
||||
"\033[1;33mLiteLLM Proxy: --reload with --num_workers>1 requires uvicorn>=0.50.0, "
|
||||
"so only a single worker reloads.\033[0m"
|
||||
)
|
||||
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()
|
||||
328
tests/unit/proxy/test_dev_reload.py
Normal file
328
tests/unit/proxy/test_dev_reload.py
Normal file
|
|
@ -0,0 +1,328 @@
|
|||
import contextlib
|
||||
import importlib
|
||||
import logging
|
||||
import logging.handlers
|
||||
import multiprocessing
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
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 (
|
||||
RELOAD_OPTIONS,
|
||||
RELOAD_WORKER_HEALTHCHECK_SECONDS,
|
||||
ReloadOptions,
|
||||
load_app_or_exit,
|
||||
multi_worker_reload_supported,
|
||||
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
|
||||
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 = 3
|
||||
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 _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_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_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)
|
||||
reloader.kill()
|
||||
assert pool_should_exit.wait(POOL_TIMEOUT_SECONDS)
|
||||
finally:
|
||||
reloader.kill()
|
||||
reloader.join()
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
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 _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
|
||||
|
||||
|
||||
@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 = _free_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)
|
||||
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()
|
||||
|
||||
|
||||
@REQUIRES_UVICORN_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()
|
||||
|
||||
|
||||
@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 = _free_port()
|
||||
process: Final = _start_pool(tmp_path, port, log, runner=SLOW_CORE_DUMP_ON_QUIT + POOL_RUNNER)
|
||||
try:
|
||||
workers: Final = _next_pool(log, process, port, frozenset())
|
||||
pool_parent: Final = psutil.Process(min(workers)).ppid()
|
||||
os.killpg(process.pid, signal.SIGQUIT)
|
||||
deadline: Final = time.monotonic() + WORKER_HEALTHCHECK_SECONDS / 2
|
||||
while process.poll() is None or 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()
|
||||
Loading…
Add table
Reference in a new issue