wehub-resource-sync
781 lines
26 KiB
Python
781 lines
26 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
|
||
"""
|
||
Tests for DPSupervisor: unit tests and lifecycle integration tests.
|
||
|
||
Lifecycle integration tests replace child vLLM servers with lightweight
|
||
aiohttp "fake" servers controlled by the test, so the suite runs without GPUs.
|
||
_start_children is monkeypatched to install FakeProcess objects (with
|
||
controllable liveness/timing) alongside those fake HTTP servers.
|
||
|
||
Port allocation (kept far from default vLLM ports to avoid conflicts):
|
||
Supervisor : 19256
|
||
Children : 18000, 18001
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import asyncio
|
||
import contextlib
|
||
import os
|
||
import signal
|
||
import subprocess
|
||
import tempfile
|
||
import time
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
|
||
import aiohttp
|
||
import pytest
|
||
import uvicorn
|
||
from fastapi import FastAPI, Response
|
||
|
||
import vllm.entrypoints.openai.dp_supervisor as dp_sup
|
||
from vllm.entrypoints.openai.dp_supervisor import (
|
||
CHILD_EXIT_GRACE_S,
|
||
DPSupervisor,
|
||
_build_vllm_dp_server_args,
|
||
infer_multi_port_external_lb_start_rank,
|
||
validate_multi_port_external_lb_args,
|
||
)
|
||
from vllm.logger import init_logger
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
_SUPERVISOR_PORT = 19256
|
||
_CHILD_PORT_BASE = 18000
|
||
_N_CHILDREN = 2
|
||
_PROBE_INTERVAL = 1.0
|
||
_POLL_INTERVAL = 1.0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Args factories
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _make_unit_args(**overrides) -> argparse.Namespace:
|
||
"""Minimal args for unit tests (no real network activity)."""
|
||
base = {
|
||
"host": None,
|
||
"port": 8000,
|
||
"data_parallel_multi_port_external_lb": True,
|
||
"data_parallel_supervisor_port": 9256,
|
||
"dp_supervisor_probe_interval_s": 5.0,
|
||
"dp_supervisor_probe_timeout_s": 5.0,
|
||
"dp_supervisor_probe_failure_threshold": 3,
|
||
"data_parallel_size": 8,
|
||
"data_parallel_size_local": 4,
|
||
"data_parallel_start_rank": None,
|
||
"data_parallel_rank": None,
|
||
"data_parallel_external_lb": False,
|
||
"data_parallel_hybrid_lb": False,
|
||
"api_server_count": None,
|
||
"headless": False,
|
||
"grpc": False,
|
||
"uds": None,
|
||
"ssl_keyfile": None,
|
||
"ssl_certfile": None,
|
||
"ssl_ca_certs": None,
|
||
"ssl_cert_reqs": 0,
|
||
"ssl_ciphers": None,
|
||
"node_rank": 1,
|
||
"tensor_parallel_size": 1,
|
||
"pipeline_parallel_size": 1,
|
||
"uvicorn_log_level": "info",
|
||
"shutdown_timeout": 5.0,
|
||
}
|
||
base.update(overrides)
|
||
return argparse.Namespace(**base)
|
||
|
||
|
||
def _make_args(**overrides) -> argparse.Namespace:
|
||
"""Args for lifecycle integration tests (real loopback servers)."""
|
||
base: dict = dict(
|
||
host="127.0.0.1",
|
||
port=_CHILD_PORT_BASE,
|
||
data_parallel_multi_port_external_lb=True,
|
||
data_parallel_supervisor_port=_SUPERVISOR_PORT,
|
||
dp_supervisor_probe_interval_s=_PROBE_INTERVAL,
|
||
dp_supervisor_probe_timeout_s=1.0,
|
||
dp_supervisor_probe_failure_threshold=3,
|
||
data_parallel_size=_N_CHILDREN,
|
||
data_parallel_size_local=_N_CHILDREN,
|
||
data_parallel_start_rank=0,
|
||
data_parallel_rank=None,
|
||
data_parallel_external_lb=False,
|
||
data_parallel_hybrid_lb=False,
|
||
api_server_count=None,
|
||
headless=False,
|
||
grpc=False,
|
||
uds=None,
|
||
ssl_keyfile=None,
|
||
ssl_certfile=None,
|
||
ssl_ca_certs=None,
|
||
ssl_cert_reqs=0,
|
||
ssl_ciphers=None,
|
||
node_rank=0,
|
||
tensor_parallel_size=1,
|
||
pipeline_parallel_size=1,
|
||
uvicorn_log_level="warning",
|
||
shutdown_timeout=0.0,
|
||
)
|
||
base.update(overrides)
|
||
return argparse.Namespace(**base)
|
||
|
||
|
||
def _generate_self_signed_cert(cert_dir: Path) -> tuple[Path, Path]:
|
||
"""Generate a self-signed certificate for HTTPS lifecycle tests."""
|
||
cert_file = cert_dir / "cert.pem"
|
||
key_file = cert_dir / "key.pem"
|
||
subprocess.run(
|
||
[
|
||
"openssl",
|
||
"req",
|
||
"-x509",
|
||
"-newkey",
|
||
"rsa:2048",
|
||
"-keyout",
|
||
str(key_file),
|
||
"-out",
|
||
str(cert_file),
|
||
"-days",
|
||
"1",
|
||
"-nodes",
|
||
"-subj",
|
||
"/CN=localhost",
|
||
],
|
||
check=True,
|
||
capture_output=True,
|
||
)
|
||
return cert_file, key_file
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Unit tests
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_infer_multi_port_external_lb_start_rank_uses_node_rank():
|
||
args = _make_unit_args()
|
||
assert infer_multi_port_external_lb_start_rank(args) == 4
|
||
|
||
|
||
def test_build_multi_port_external_lb_child_args_sets_external_rank_server():
|
||
args = _make_unit_args(data_parallel_start_rank=8, api_server_count=None)
|
||
child_args = _build_vllm_dp_server_args(args, local_rank=2)
|
||
|
||
assert child_args.port == 8002
|
||
assert child_args.data_parallel_rank == 10
|
||
assert child_args.data_parallel_size_local == 1
|
||
assert child_args.data_parallel_external_lb is True
|
||
assert child_args.data_parallel_hybrid_lb is False
|
||
assert child_args.data_parallel_multi_port_external_lb is False
|
||
assert child_args.api_server_count == 1
|
||
|
||
|
||
def test_run_vllm_dp_server_uses_python_server_by_default(monkeypatch):
|
||
calls: list[str] = []
|
||
|
||
monkeypatch.setattr(dp_sup.os, "setpgrp", lambda: None)
|
||
monkeypatch.setattr(dp_sup, "set_process_title", lambda *_args: None)
|
||
monkeypatch.setattr(dp_sup, "decorate_logs", lambda *_args: None)
|
||
monkeypatch.setattr(dp_sup.envs, "VLLM_RUST_FRONTEND_PATH", None, raising=False)
|
||
monkeypatch.setattr(
|
||
dp_sup, "_run_python_vllm_dp_server", lambda _args: calls.append("python")
|
||
)
|
||
monkeypatch.setattr(
|
||
dp_sup, "_run_rust_vllm_dp_server", lambda _args: calls.append("rust")
|
||
)
|
||
|
||
dp_sup._run_vllm_dp_server(_make_unit_args(data_parallel_rank=4))
|
||
|
||
assert calls == ["python"]
|
||
|
||
|
||
def test_run_vllm_dp_server_uses_rust_frontend_when_enabled(monkeypatch):
|
||
calls: list[str] = []
|
||
|
||
monkeypatch.setattr(dp_sup.os, "setpgrp", lambda: None)
|
||
monkeypatch.setattr(dp_sup, "set_process_title", lambda *_args: None)
|
||
monkeypatch.setattr(dp_sup, "decorate_logs", lambda *_args: None)
|
||
monkeypatch.setattr(
|
||
dp_sup.envs,
|
||
"VLLM_RUST_FRONTEND_PATH",
|
||
"/tmp/vllm-rs",
|
||
raising=False,
|
||
)
|
||
monkeypatch.setattr(
|
||
dp_sup, "_run_python_vllm_dp_server", lambda _args: calls.append("python")
|
||
)
|
||
monkeypatch.setattr(
|
||
dp_sup, "_run_rust_vllm_dp_server", lambda _args: calls.append("rust")
|
||
)
|
||
|
||
dp_sup._run_vllm_dp_server(_make_unit_args(data_parallel_rank=4))
|
||
|
||
assert calls == ["rust"]
|
||
|
||
|
||
def test_validate_multi_port_external_lb_args_allows_ssl():
|
||
args = _make_unit_args(
|
||
ssl_keyfile="/tmp/server.key",
|
||
ssl_certfile="/tmp/server.crt",
|
||
ssl_ca_certs="/tmp/ca.crt",
|
||
)
|
||
validate_multi_port_external_lb_args(args)
|
||
|
||
|
||
def test_aggregates_health():
|
||
supervisor = DPSupervisor(_make_unit_args())
|
||
supervisor._is_ready = True
|
||
assert supervisor.is_ready is True
|
||
|
||
|
||
def test_handles_shutdown_event():
|
||
supervisor = DPSupervisor(_make_unit_args())
|
||
supervisor._is_ready = True
|
||
supervisor._shutdown_event.set()
|
||
assert supervisor.is_ready is False
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_handles_child_exit(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
):
|
||
supervisor = DPSupervisor(_make_unit_args())
|
||
supervisor._processes = [
|
||
SimpleNamespace(
|
||
name="APIServer_DPRank_4", exitcode=None, is_alive=lambda: True
|
||
),
|
||
SimpleNamespace(name="APIServer_DPRank_5", exitcode=17, is_alive=lambda: False),
|
||
]
|
||
|
||
async def fake_probe(*_args, **_kwargs) -> bool:
|
||
return True
|
||
|
||
monkeypatch.setattr(dp_sup, "_probe_endpoint", fake_probe)
|
||
|
||
await supervisor._monitor_children()
|
||
assert supervisor._is_ready is False
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_handles_probe_failure(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
):
|
||
supervisor = DPSupervisor(_make_unit_args(dp_supervisor_probe_interval_s=0.0))
|
||
supervisor.child_ports = [8000]
|
||
probe_results = iter([True, False])
|
||
|
||
async def fake_probe(*_args, **_kwargs) -> bool:
|
||
return next(probe_results)
|
||
|
||
monkeypatch.setattr(dp_sup, "_probe_endpoint", fake_probe)
|
||
|
||
await supervisor._monitor_children()
|
||
assert supervisor._is_ready is False
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_shutdown_if_supervisor_server_error_on_startup(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
):
|
||
class FakeLoop:
|
||
def add_signal_handler(self, *_args, **_kwargs):
|
||
pass
|
||
|
||
def remove_signal_handler(self, *_args, **_kwargs):
|
||
pass
|
||
|
||
class FakeServer:
|
||
def __init__(self, _config):
|
||
self.started = False
|
||
self.should_exit = False
|
||
|
||
async def serve(self):
|
||
raise ValueError("supervisor boom")
|
||
|
||
async def fake_shutdown_children(self):
|
||
return None
|
||
|
||
def fake_start_children(self):
|
||
return None
|
||
|
||
async def fake_monitor_children(self):
|
||
# Mark ready so the supervisor server is started, then block until
|
||
# shutdown (triggered when the failing server task exits).
|
||
self._is_ready = True
|
||
await self._shutdown_event.wait()
|
||
|
||
monkeypatch.setattr(dp_sup.asyncio, "get_running_loop", lambda: FakeLoop())
|
||
monkeypatch.setattr(dp_sup.uvicorn, "Server", FakeServer)
|
||
monkeypatch.setattr(DPSupervisor, "_shutdown_children", fake_shutdown_children)
|
||
monkeypatch.setattr(DPSupervisor, "_start_children", fake_start_children)
|
||
monkeypatch.setattr(DPSupervisor, "_monitor_children", fake_monitor_children)
|
||
|
||
supervisor = DPSupervisor(_make_unit_args())
|
||
|
||
with pytest.raises(ValueError, match="supervisor boom"):
|
||
await supervisor.run()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Lifecycle integration tests – MockVLLMServer
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class MockVLLMServer:
|
||
"""
|
||
Minimal FastAPI server that mimics one vLLM replica.
|
||
GET /health returns 200 when healthy, 503 otherwise.
|
||
Health state is toggled by the test via set_healthy().
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
port: int,
|
||
drain_seconds: float = 0.0,
|
||
ssl_keyfile: str | None = None,
|
||
ssl_certfile: str | None = None,
|
||
) -> None:
|
||
self.port = port
|
||
self._healthy = False
|
||
self._drain_seconds = drain_seconds
|
||
self._ssl_keyfile = ssl_keyfile
|
||
self._ssl_certfile = ssl_certfile
|
||
self._server: uvicorn.Server | None = None
|
||
self._serve_task: asyncio.Task | None = None
|
||
|
||
async def start(self) -> None:
|
||
app = FastAPI()
|
||
|
||
@app.get("/health")
|
||
async def health() -> Response:
|
||
print(f"MockServer {self.port}: /health: {self._healthy}")
|
||
return Response(status_code=200 if self._healthy else 503)
|
||
|
||
@app.get("/set_healthy")
|
||
async def set_healthy() -> Response:
|
||
print(f"MockServer {self.port}: /set_healthy")
|
||
self._healthy = True
|
||
return Response(status_code=200)
|
||
|
||
@app.get("/set_unhealthy")
|
||
async def set_unhealthy() -> Response:
|
||
print(f"MockServer {self.port}: /set_unhealthy")
|
||
self._healthy = False
|
||
return Response(status_code=200)
|
||
|
||
@app.get("/kill")
|
||
async def kill() -> Response:
|
||
print(f"MockServer {self.port}: /kill")
|
||
os.kill(os.getpid(), signal.SIGKILL)
|
||
|
||
config = uvicorn.Config(
|
||
app,
|
||
host="127.0.0.1",
|
||
port=self.port,
|
||
log_level="warning",
|
||
lifespan="off",
|
||
ssl_keyfile=self._ssl_keyfile,
|
||
ssl_certfile=self._ssl_certfile,
|
||
)
|
||
self._server = uvicorn.Server(config)
|
||
|
||
# Configure request draining if needed.
|
||
# Uvicorn's capture_signals() installs signal.signal(SIGTERM, self.handle_exit),
|
||
# which sets should_exit=True immediately. Override handle_exit on the instance
|
||
# so capture_signals() picks up our version that drains first.
|
||
if self._drain_seconds > 0:
|
||
self._shutdown_event = asyncio.Event()
|
||
loop = asyncio.get_running_loop()
|
||
|
||
async def _drain_and_stop() -> None:
|
||
await self._shutdown_event.wait()
|
||
print(f"MockServer {self.port}: draining for {self._drain_seconds}s.")
|
||
await asyncio.sleep(self._drain_seconds)
|
||
print("Setting should_exit")
|
||
if self._server is not None:
|
||
self._server.should_exit = True
|
||
|
||
self._drain_task = asyncio.create_task(_drain_and_stop())
|
||
|
||
def _custom_handle_exit(sig: int, frame: object) -> None:
|
||
print("Got SIGTERM, setting shutdown.")
|
||
if not self._shutdown_event.is_set():
|
||
loop.call_soon_threadsafe(self._shutdown_event.set)
|
||
|
||
self._server.handle_exit = _custom_handle_exit
|
||
|
||
self._serve_task = asyncio.create_task(self._server.serve())
|
||
while not self._server.started:
|
||
await asyncio.sleep(0.01)
|
||
print(f"Mock DP Server on port {self.port} started")
|
||
|
||
await self._serve_task
|
||
|
||
|
||
def launch_mock_vllm(child_args: argparse.Namespace):
|
||
logger.info("Launching mock vLLM on port %s", child_args.port)
|
||
mock_vllm = MockVLLMServer(
|
||
port=child_args.port,
|
||
ssl_keyfile=child_args.ssl_keyfile,
|
||
ssl_certfile=child_args.ssl_certfile,
|
||
)
|
||
asyncio.run(mock_vllm.start())
|
||
|
||
|
||
def launch_mock_vllm_with_drain(
|
||
child_args: argparse.Namespace,
|
||
):
|
||
logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port)
|
||
mock_vllm = MockVLLMServer(
|
||
port=child_args.port,
|
||
drain_seconds=10.0,
|
||
ssl_keyfile=child_args.ssl_keyfile,
|
||
ssl_certfile=child_args.ssl_certfile,
|
||
)
|
||
asyncio.run(mock_vllm.start())
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Lifecycle test helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def _poll_supervisor_health(expected_status: int, use_ssl: bool = False) -> bool:
|
||
"""
|
||
GET /health on the supervisor once and check for expected_status.
|
||
|
||
Pass expected_status=-1 to assert the supervisor is not listening yet
|
||
(a connection error is expected). The supervisor only starts its HTTP
|
||
server once every child is ready, so /health is refused until then.
|
||
"""
|
||
scheme = "https" if use_ssl else "http"
|
||
url = f"{scheme}://127.0.0.1:{_SUPERVISOR_PORT}/health"
|
||
async with aiohttp.ClientSession() as session:
|
||
try:
|
||
async with session.get(url, ssl=False if use_ssl else None) as resp:
|
||
if resp.status != expected_status:
|
||
print(f"expected: {expected_status=}, got: {resp.status=}")
|
||
return False
|
||
return True
|
||
except aiohttp.ClientError:
|
||
if expected_status != -1:
|
||
print(f"expected: {expected_status=}, got: aiohttp.ClientError")
|
||
return False
|
||
return True
|
||
|
||
|
||
async def _await_supervisor_health(
|
||
expected_status: int, use_ssl: bool = False, retries: int = 20
|
||
) -> bool:
|
||
"""Retry _poll_supervisor_health, tolerating supervisor server startup."""
|
||
for _ in range(retries):
|
||
if await _poll_supervisor_health(expected_status, use_ssl=use_ssl):
|
||
return True
|
||
await asyncio.sleep(0.5)
|
||
return False
|
||
|
||
|
||
async def _poll_until_api_server_running(
|
||
port: int, retries: int = 10, use_ssl: bool = False
|
||
) -> None:
|
||
scheme = "https" if use_ssl else "http"
|
||
url = f"{scheme}://127.0.0.1:{port}/health"
|
||
async with aiohttp.ClientSession() as session:
|
||
for _ in range(retries):
|
||
try:
|
||
async with session.get(url, ssl=False if use_ssl else None) as resp:
|
||
if resp.status != 200:
|
||
return
|
||
await asyncio.sleep(1.0)
|
||
except aiohttp.ClientError:
|
||
print("Test detected not started yet, sleeping for 1s")
|
||
await asyncio.sleep(1.0)
|
||
|
||
|
||
async def _set_healthy(port: int, use_ssl: bool = False) -> None:
|
||
scheme = "https" if use_ssl else "http"
|
||
url = f"{scheme}://127.0.0.1:{port}/set_healthy"
|
||
async with (
|
||
aiohttp.ClientSession() as session,
|
||
session.get(url, ssl=False if use_ssl else None) as resp,
|
||
):
|
||
assert resp.status == 200
|
||
|
||
|
||
async def _set_unhealthy(port: int, use_ssl: bool = False) -> None:
|
||
scheme = "https" if use_ssl else "http"
|
||
url = f"{scheme}://127.0.0.1:{port}/set_unhealthy"
|
||
async with (
|
||
aiohttp.ClientSession() as session,
|
||
session.get(url, ssl=False if use_ssl else None) as resp,
|
||
):
|
||
assert resp.status == 200
|
||
|
||
|
||
async def _kill_server(port: int, use_ssl: bool = False) -> None:
|
||
scheme = "https" if use_ssl else "http"
|
||
url = f"{scheme}://127.0.0.1:{port}/kill"
|
||
try:
|
||
async with (
|
||
aiohttp.ClientSession() as session,
|
||
session.get(url, ssl=False if use_ssl else None) as resp,
|
||
):
|
||
assert resp.status != 200
|
||
except Exception as e:
|
||
assert isinstance(e, aiohttp.ClientConnectorError)
|
||
|
||
|
||
@contextlib.asynccontextmanager
|
||
async def _run_supervisor(
|
||
args: argparse.Namespace,
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
launch_fn=None,
|
||
):
|
||
if launch_fn is None:
|
||
launch_fn = launch_mock_vllm
|
||
monkeypatch.setattr(dp_sup, "_run_vllm_dp_server", launch_fn)
|
||
supervisor = DPSupervisor(args)
|
||
task = asyncio.create_task(supervisor.run())
|
||
await asyncio.sleep(1.0)
|
||
try:
|
||
yield supervisor, task
|
||
finally:
|
||
task.cancel()
|
||
with contextlib.suppress(asyncio.CancelledError):
|
||
await task
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Lifecycle integration tests
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_basic_lifecycle(monkeypatch):
|
||
"""
|
||
A) Supervisor is not listening while children are unhealthy.
|
||
B) /health returns 200 once every child reports healthy.
|
||
C) SIGTERM and shutdown
|
||
"""
|
||
args = _make_args()
|
||
|
||
vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]
|
||
|
||
async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
|
||
assert await _poll_supervisor_health(-1)
|
||
assert not supervisor.is_ready
|
||
|
||
for port in vllm_server_ports:
|
||
assert await _poll_supervisor_health(-1)
|
||
assert not supervisor.is_ready
|
||
await _poll_until_api_server_running(port)
|
||
|
||
await _set_healthy(vllm_server_ports[0])
|
||
await asyncio.sleep(1.0)
|
||
assert await _poll_supervisor_health(-1)
|
||
assert not supervisor.is_ready
|
||
print("supervisor not listening --- expected!")
|
||
|
||
for port in vllm_server_ports:
|
||
await _set_healthy(port)
|
||
assert await _await_supervisor_health(200)
|
||
assert supervisor.is_ready
|
||
print("/health is 200 --- expected!")
|
||
|
||
await asyncio.sleep(1.0)
|
||
assert await _poll_supervisor_health(200)
|
||
assert supervisor.is_ready
|
||
print("/health is 200 --- expected!")
|
||
|
||
os.kill(os.getpid(), signal.SIGTERM)
|
||
|
||
await asyncio.wait_for(_task, timeout=5.0)
|
||
for p in supervisor._processes:
|
||
assert not p.is_alive()
|
||
print("everything was cleaned up!")
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_basic_lifecycle_with_ssl(monkeypatch):
|
||
with tempfile.TemporaryDirectory() as cert_dir:
|
||
cert_file, key_file = _generate_self_signed_cert(Path(cert_dir))
|
||
args = _make_args(
|
||
ssl_keyfile=str(key_file),
|
||
ssl_certfile=str(cert_file),
|
||
)
|
||
|
||
vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]
|
||
|
||
async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
|
||
assert await _poll_supervisor_health(-1, use_ssl=True)
|
||
assert not supervisor.is_ready
|
||
|
||
for port in vllm_server_ports:
|
||
assert await _poll_supervisor_health(-1, use_ssl=True)
|
||
assert not supervisor.is_ready
|
||
await _poll_until_api_server_running(port, use_ssl=True)
|
||
|
||
for port in vllm_server_ports:
|
||
await _set_healthy(port, use_ssl=True)
|
||
|
||
assert await _await_supervisor_health(200, use_ssl=True)
|
||
assert supervisor.is_ready
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_failed_startup(monkeypatch):
|
||
"""
|
||
A) One of the vLLM servers crashes during startup.
|
||
B) DPSupervisor detects this, and cleans up resources.
|
||
"""
|
||
args = _make_args()
|
||
|
||
vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]
|
||
|
||
async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
|
||
assert await _poll_supervisor_health(-1)
|
||
assert not supervisor.is_ready
|
||
|
||
for port in vllm_server_ports:
|
||
await _poll_until_api_server_running(port)
|
||
|
||
await _kill_server(port)
|
||
|
||
await asyncio.wait_for(_task, timeout=5.0)
|
||
for p in supervisor._processes:
|
||
assert not p.is_alive()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_becomes_unhealthy(monkeypatch):
|
||
"""
|
||
A) Supervisor is not listening while children are unhealthy.
|
||
B) /health returns 200 once every child reports healthy.
|
||
C) Child process becomes unhealthy.
|
||
D) Detected and shutdown.
|
||
"""
|
||
args = _make_args()
|
||
|
||
vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]
|
||
|
||
async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
|
||
assert await _poll_supervisor_health(-1)
|
||
assert not supervisor.is_ready
|
||
|
||
for port in vllm_server_ports:
|
||
assert await _poll_supervisor_health(-1)
|
||
assert not supervisor.is_ready
|
||
await _poll_until_api_server_running(port)
|
||
|
||
await _set_healthy(vllm_server_ports[0])
|
||
await asyncio.sleep(1.0)
|
||
assert await _poll_supervisor_health(-1)
|
||
assert not supervisor.is_ready
|
||
print("supervisor not listening --- expected!")
|
||
|
||
for port in vllm_server_ports:
|
||
await _set_healthy(port)
|
||
assert await _await_supervisor_health(200)
|
||
assert supervisor.is_ready
|
||
print("/health is 200 --- expected!")
|
||
|
||
await _set_unhealthy(port)
|
||
|
||
await asyncio.wait_for(_task, timeout=5.0)
|
||
for p in supervisor._processes:
|
||
assert not p.is_alive()
|
||
print("everything was cleaned up!")
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_dp_server_fails(monkeypatch):
|
||
"""
|
||
A) Supervisor is not listening while children are unhealthy.
|
||
B) /health returns 200 once every child reports healthy.
|
||
C) Child process fails.
|
||
D) Detected and shutdown.
|
||
"""
|
||
args = _make_args()
|
||
|
||
vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]
|
||
|
||
async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
|
||
assert await _poll_supervisor_health(-1)
|
||
assert not supervisor.is_ready
|
||
|
||
for port in vllm_server_ports:
|
||
assert await _poll_supervisor_health(-1)
|
||
assert not supervisor.is_ready
|
||
await _poll_until_api_server_running(port)
|
||
|
||
await _set_healthy(vllm_server_ports[0])
|
||
await asyncio.sleep(1.0)
|
||
assert await _poll_supervisor_health(-1)
|
||
assert not supervisor.is_ready
|
||
print("supervisor not listening --- expected!")
|
||
|
||
for port in vllm_server_ports:
|
||
await _set_healthy(port)
|
||
assert await _await_supervisor_health(200)
|
||
assert supervisor.is_ready
|
||
print("/health is 200 --- expected!")
|
||
|
||
dp_mock_server_process = supervisor._processes[0]
|
||
os.kill(dp_mock_server_process.pid, signal.SIGKILL)
|
||
await asyncio.sleep(1.0)
|
||
assert not dp_mock_server_process.is_alive()
|
||
|
||
await asyncio.wait_for(_task, timeout=5.0)
|
||
for p in supervisor._processes:
|
||
assert not p.is_alive()
|
||
print("everything was cleaned up!")
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_shutdown_timeout(monkeypatch: pytest.MonkeyPatch):
|
||
"""
|
||
Child mock servers delay shutdown by 10s on SIGTERM (simulating in-flight
|
||
request drain). The supervisor is configured with shutdown_timeout=10,
|
||
so its total wait budget is 10 + CHILD_EXIT_GRACE_S seconds. The
|
||
children exit naturally within that window, so no force-kill should occur
|
||
and the measured wall-clock time must be >= 10s.
|
||
"""
|
||
_DRAIN_SECONDS = 10.0
|
||
_SHUTDOWN_TIMEOUT = 10.0
|
||
|
||
args = _make_args(shutdown_timeout=_SHUTDOWN_TIMEOUT)
|
||
vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]
|
||
|
||
async with _run_supervisor(
|
||
args, monkeypatch, launch_fn=launch_mock_vllm_with_drain
|
||
) as (supervisor, _task):
|
||
for port in vllm_server_ports:
|
||
await _poll_until_api_server_running(port)
|
||
|
||
for port in vllm_server_ports:
|
||
await _set_healthy(port)
|
||
assert await _await_supervisor_health(200)
|
||
assert supervisor.is_ready
|
||
|
||
start_t = time.perf_counter()
|
||
os.kill(os.getpid(), signal.SIGTERM)
|
||
|
||
print(f"DRAINING FOR {_DRAIN_SECONDS}")
|
||
await asyncio.wait_for(_task, timeout=_DRAIN_SECONDS + CHILD_EXIT_GRACE_S + 5.0)
|
||
elapsed = time.perf_counter() - start_t
|
||
|
||
assert elapsed >= _DRAIN_SECONDS, (
|
||
f"Supervisor exited after only {elapsed:.1f}s; "
|
||
f"expected >= {_DRAIN_SECONDS}s for request draining"
|
||
)
|
||
|
||
for p in supervisor._processes:
|
||
assert not p.is_alive()
|
||
print(f"Supervisor waited {elapsed:.1f}s for children to drain — expected!")
|