wehub-resource-sync
557 lines
20 KiB
Python
557 lines
20 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import contextlib
|
|
import copy
|
|
import multiprocessing
|
|
import os
|
|
import signal
|
|
import time
|
|
from functools import partial
|
|
from http import HTTPStatus
|
|
from multiprocessing.process import BaseProcess
|
|
|
|
import aiohttp
|
|
import psutil
|
|
import uvicorn
|
|
import uvloop
|
|
from fastapi import FastAPI, Response
|
|
|
|
import vllm.envs as envs
|
|
from vllm.logger import init_logger
|
|
from vllm.utils.system_utils import (
|
|
decorate_logs,
|
|
kill_process_tree,
|
|
set_process_title,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
CHILD_EXIT_GRACE_S = 5.0
|
|
|
|
|
|
def infer_multi_port_external_lb_start_rank(args: argparse.Namespace) -> int:
|
|
start_rank = getattr(args, "data_parallel_start_rank", None)
|
|
if start_rank is not None:
|
|
return start_rank
|
|
|
|
node_rank = getattr(args, "node_rank", 0) or 0
|
|
local_size = getattr(args, "data_parallel_size_local", 0) or 0
|
|
return node_rank * local_size
|
|
|
|
|
|
def validate_multi_port_external_lb_args(args: argparse.Namespace) -> None:
|
|
if getattr(args, "grpc", False):
|
|
raise ValueError(
|
|
"Error: --data-parallel-multi-port-external-lb does not support --grpc"
|
|
)
|
|
if args.uds is not None:
|
|
raise ValueError(
|
|
"Error: --data-parallel-multi-port-external-lb does not support --uds"
|
|
)
|
|
if bool(args.ssl_keyfile) != bool(args.ssl_certfile):
|
|
raise ValueError(
|
|
"Error: --ssl-keyfile and --ssl-certfile must be provided together"
|
|
)
|
|
if args.api_server_count not in (None, 1):
|
|
raise ValueError(
|
|
"Error: --data-parallel-multi-port-external-lb currently requires "
|
|
"--api-server-count=1"
|
|
)
|
|
if args.data_parallel_rank is not None:
|
|
raise ValueError(
|
|
"Error: --data-parallel-multi-port-external-lb manages child "
|
|
"--data-parallel-rank values internally"
|
|
)
|
|
if args.data_parallel_external_lb or args.data_parallel_hybrid_lb:
|
|
raise ValueError(
|
|
"Error: --data-parallel-multi-port-external-lb cannot be combined with "
|
|
"--data-parallel-external-lb or --data-parallel-hybrid-lb"
|
|
)
|
|
if args.data_parallel_size < 2:
|
|
raise ValueError(
|
|
"Error: --data-parallel-multi-port-external-lb requires "
|
|
"--data-parallel-size > 1"
|
|
)
|
|
|
|
local_size = args.data_parallel_size_local
|
|
if local_size is None or local_size < 2:
|
|
raise ValueError(
|
|
"Error: --data-parallel-multi-port-external-lb requires "
|
|
"--data-parallel-size-local >= 2"
|
|
)
|
|
if local_size > args.data_parallel_size:
|
|
raise ValueError(
|
|
"Error: --data-parallel-size-local cannot exceed --data-parallel-size"
|
|
)
|
|
if args.data_parallel_size % local_size != 0:
|
|
raise ValueError(
|
|
"Error: --data-parallel-size must be divisible by "
|
|
"--data-parallel-size-local"
|
|
)
|
|
|
|
start_rank = infer_multi_port_external_lb_start_rank(args)
|
|
if start_rank + local_size > args.data_parallel_size:
|
|
raise ValueError(
|
|
"Error: multi-port supervised ranks would exceed --data-parallel-size"
|
|
)
|
|
|
|
supervisor_port = args.data_parallel_supervisor_port
|
|
child_port_min = args.port
|
|
child_port_max = args.port + local_size - 1
|
|
if child_port_min <= supervisor_port <= child_port_max:
|
|
raise ValueError(
|
|
f"Error: --data-parallel-supervisor-port {supervisor_port} "
|
|
f"overlaps with child rank ports {child_port_min}-{child_port_max}"
|
|
)
|
|
|
|
|
|
def _build_vllm_dp_server_args(
|
|
args: argparse.Namespace, local_rank: int
|
|
) -> argparse.Namespace:
|
|
child_args = copy.copy(args)
|
|
child_args.port = args.port + local_rank
|
|
child_args.data_parallel_rank = (
|
|
infer_multi_port_external_lb_start_rank(args) + local_rank
|
|
)
|
|
child_args.data_parallel_start_rank = None
|
|
child_args.data_parallel_size_local = 1
|
|
child_args.data_parallel_external_lb = True
|
|
child_args.data_parallel_hybrid_lb = False
|
|
child_args.data_parallel_multi_port_external_lb = False
|
|
child_args.data_parallel_supervisor_port = None
|
|
child_args.api_server_count = 1
|
|
child_args.device_ids = _build_device_ids(args, local_rank)
|
|
return child_args
|
|
|
|
|
|
def _build_device_ids(args: argparse.Namespace, local_rank: int) -> list[int | str]:
|
|
"""Build the --device-ids value for a DP child process.
|
|
|
|
The child resolves these against its own inherited device-control env
|
|
var (e.g. CUDA_VISIBLE_DEVICES), so integer IDs must stay env-relative
|
|
here rather than being translated to physical IDs.
|
|
"""
|
|
devices_per_rank = args.tensor_parallel_size * args.pipeline_parallel_size
|
|
start = local_rank * devices_per_rank
|
|
stop = start + devices_per_rank
|
|
device_ids = getattr(args, "device_ids", None)
|
|
if device_ids is not None:
|
|
if stop > len(device_ids):
|
|
raise ValueError(
|
|
f"--device-ids has {len(device_ids)} entries, but DP rank "
|
|
f"{local_rank} needs devices [{start}, {stop})"
|
|
)
|
|
return device_ids[start:stop]
|
|
return list(range(start, stop))
|
|
|
|
|
|
def _child_base_url(args: argparse.Namespace, port: int) -> str:
|
|
host = args.host or "127.0.0.1"
|
|
if host == "0.0.0.0":
|
|
host = "127.0.0.1"
|
|
elif host == "::":
|
|
host = "::1"
|
|
scheme = "https" if args.ssl_keyfile and args.ssl_certfile else "http"
|
|
return f"{scheme}://{host}:{port}"
|
|
|
|
|
|
def _join_processes_with_timeout(processes: list[BaseProcess], timeout: float) -> None:
|
|
deadline = time.monotonic() + timeout
|
|
for process in processes:
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
break
|
|
process.join(timeout=remaining)
|
|
|
|
|
|
async def _probe_endpoint(
|
|
session: aiohttp.ClientSession,
|
|
args: argparse.Namespace,
|
|
port: int,
|
|
path: str,
|
|
conn_err_failure_threshold: int = 3,
|
|
conn_err_retry_delay: float = 5.0,
|
|
) -> bool:
|
|
"""
|
|
Probe /health endpoint for 200 status.
|
|
|
|
If there is a connection error, retry every N seconds.
|
|
"""
|
|
for iteration in range(conn_err_failure_threshold):
|
|
try:
|
|
probe_ssl = None
|
|
if args.ssl_keyfile and args.ssl_certfile:
|
|
# Probes target node-local child servers over loopback, so skip
|
|
# certificate verification to avoid SAN/hostname mismatches for
|
|
# localhost/127.0.0.1 deployments.
|
|
probe_ssl = False
|
|
async with session.get(
|
|
_child_base_url(args, port) + path, ssl=probe_ssl
|
|
) as response:
|
|
# vLLM returns 503 on EngineDeadError, so we should return
|
|
# immediately if vLLM responds with a non-200 status code.
|
|
return response.status == HTTPStatus.OK
|
|
except (aiohttp.ClientError, asyncio.TimeoutError) as e:
|
|
# Allow retry of connection errors.
|
|
logger.debug(
|
|
"Probe attempt %d/%d failed on port %d: %r",
|
|
iteration + 1,
|
|
conn_err_failure_threshold,
|
|
port,
|
|
e,
|
|
)
|
|
|
|
if iteration < conn_err_failure_threshold - 1:
|
|
await asyncio.sleep(conn_err_retry_delay)
|
|
|
|
return False
|
|
|
|
|
|
def _build_dp_supervisor_app(supervisor: DPSupervisor) -> FastAPI:
|
|
app = FastAPI(openapi_url=None, docs_url=None, redoc_url=None)
|
|
app.state.supervisor = supervisor
|
|
|
|
def _status_response(ok: bool) -> Response:
|
|
return Response(
|
|
status_code=(HTTPStatus.OK if ok else HTTPStatus.SERVICE_UNAVAILABLE)
|
|
)
|
|
|
|
@app.get("/health", include_in_schema=False)
|
|
async def health() -> Response:
|
|
return _status_response(app.state.supervisor.is_ready)
|
|
|
|
@app.get("/ready", include_in_schema=False)
|
|
@app.get("/readyz", include_in_schema=False)
|
|
async def ready() -> Response:
|
|
return _status_response(app.state.supervisor.is_ready)
|
|
|
|
return app
|
|
|
|
|
|
def _run_python_vllm_dp_server(child_args: argparse.Namespace) -> None:
|
|
from vllm.entrypoints.openai.api_server import run_server
|
|
|
|
uvloop.run(run_server(child_args))
|
|
|
|
|
|
def _run_rust_vllm_dp_server(child_args: argparse.Namespace) -> None:
|
|
from vllm.entrypoints.cli.serve import run_multi_api_server
|
|
|
|
run_multi_api_server(child_args)
|
|
|
|
|
|
def _run_vllm_dp_server(child_args: argparse.Namespace) -> None:
|
|
"""
|
|
Entrypoint function for the vLLM DP Server.
|
|
"""
|
|
# Create a fresh process group for the vLLM DP Server,
|
|
# so that CTRL-C is propagated cleanly.
|
|
os.setpgrp()
|
|
|
|
name = f"APIServer_DP{child_args.data_parallel_rank}"
|
|
set_process_title(name)
|
|
decorate_logs(name)
|
|
if envs.VLLM_RUST_FRONTEND_PATH:
|
|
_run_rust_vllm_dp_server(child_args)
|
|
else:
|
|
_run_python_vllm_dp_server(child_args)
|
|
|
|
|
|
class DPSupervisor:
|
|
def __init__(self, args: argparse.Namespace):
|
|
validate_multi_port_external_lb_args(args)
|
|
self.args = args
|
|
self.supervisor_port = args.data_parallel_supervisor_port
|
|
self.child_ports = [
|
|
args.port + local_rank
|
|
for local_rank in range(args.data_parallel_size_local)
|
|
]
|
|
self._is_ready = False
|
|
self._processes: list[BaseProcess] = []
|
|
self._shutdown_event = asyncio.Event()
|
|
self._shutdown_signal = signal.SIGTERM
|
|
|
|
@property
|
|
def is_ready(self) -> bool:
|
|
return self._is_ready and not self._shutdown_event.is_set()
|
|
|
|
async def run(self) -> None:
|
|
loop = asyncio.get_running_loop()
|
|
decorate_logs("DPSupervisor")
|
|
|
|
# K8s sends SIGTERM for shutdown - begin graceful termination.
|
|
for sig in (signal.SIGTERM, signal.SIGINT):
|
|
loop.add_signal_handler(sig, partial(self._handle_signal, sig))
|
|
|
|
supervisor_server: uvicorn.Server | None = None
|
|
supervisor_server_task: asyncio.Task[None] | None = None
|
|
try:
|
|
# Launch vLLM DP Servers and begin monitoring them.
|
|
self._start_children()
|
|
monitor_task = asyncio.create_task(
|
|
self._monitor_children(), name="dp-monitor"
|
|
)
|
|
|
|
# Only start the DPSupervisor server once the vLLM DP Servers are
|
|
# ready. This avoids the supervisor /health endpoint returning 503
|
|
# to external load balancer probes while the engines initialize.
|
|
await self._wait_until_ready(monitor_task)
|
|
if self.is_ready and not monitor_task.done():
|
|
supervisor_server, supervisor_server_task = await self._start_server()
|
|
|
|
await monitor_task
|
|
finally:
|
|
self._is_ready = False
|
|
await self._shutdown_children()
|
|
|
|
# Shutdown the DP Supervisor server.
|
|
if supervisor_server is not None and supervisor_server_task is not None:
|
|
supervisor_server.should_exit = True
|
|
await supervisor_server_task
|
|
|
|
async def _start_server(self) -> tuple[uvicorn.Server, asyncio.Task[None]]:
|
|
"""
|
|
Launch the DPSupervisor HTTP server.
|
|
|
|
Called only after the vLLM DP Servers are ready so that /health does
|
|
not return 503 to external probes while the engines are initializing.
|
|
"""
|
|
app = _build_dp_supervisor_app(self)
|
|
host = self.args.host or "0.0.0.0"
|
|
config = uvicorn.Config(
|
|
app,
|
|
host=host,
|
|
port=self.supervisor_port,
|
|
log_level=self.args.uvicorn_log_level,
|
|
ssl_keyfile=self.args.ssl_keyfile,
|
|
ssl_certfile=self.args.ssl_certfile,
|
|
ssl_ca_certs=self.args.ssl_ca_certs,
|
|
ssl_cert_reqs=self.args.ssl_cert_reqs,
|
|
ssl_ciphers=self.args.ssl_ciphers,
|
|
)
|
|
supervisor_server = uvicorn.Server(config)
|
|
supervisor_server_task = asyncio.create_task(
|
|
supervisor_server.serve(),
|
|
name="dp-supervisor",
|
|
)
|
|
supervisor_server_task.add_done_callback(
|
|
lambda _task: self._shutdown_event.set()
|
|
)
|
|
|
|
# Ensure DPSupervisor task starts on the event loop.
|
|
while not supervisor_server.started:
|
|
if supervisor_server_task.done():
|
|
supervisor_server_task.result()
|
|
raise RuntimeError("DPSupervisor exited before startup.")
|
|
await asyncio.sleep(0.05)
|
|
logger.info("Started DPSupervisor on %s:%d", host, self.supervisor_port)
|
|
return supervisor_server, supervisor_server_task
|
|
|
|
async def _wait_until_ready(self, monitor_task: asyncio.Task[None]) -> None:
|
|
"""
|
|
Block until the vLLM DP Servers are ready or shutdown is triggered.
|
|
|
|
Returns early if monitoring stops (e.g. a DP Server dies during
|
|
startup), in which case the supervisor server is never started.
|
|
"""
|
|
logger.info("Waiting for vLLM DP Servers to become ready.")
|
|
while not self._is_ready and not self._shutdown_event.is_set():
|
|
if monitor_task.done():
|
|
return
|
|
await asyncio.sleep(0.05)
|
|
|
|
def _handle_signal(self, signum: int) -> None:
|
|
"""
|
|
Signal handler that is added to the event loop.
|
|
|
|
This catches the SIGTERM from K8s and begins graceful shutdown,
|
|
by setting the _shutdown_event(), which is watched by the main
|
|
coroutine monitoring the vLLM DP Servers.
|
|
"""
|
|
|
|
if self._shutdown_event.is_set():
|
|
return
|
|
|
|
self._shutdown_signal = signal.Signals(signum)
|
|
logger.info(
|
|
"DPSupervisor received %s, shutting down.",
|
|
self._shutdown_signal.name,
|
|
)
|
|
|
|
self._shutdown_event.set()
|
|
self._is_ready = False
|
|
|
|
def _start_children(self) -> None:
|
|
"""
|
|
Launch vLLM DP Servers on separate GPUs.
|
|
"""
|
|
logger.info("Launching vLLM DP Servers")
|
|
context = multiprocessing.get_context("spawn")
|
|
for local_rank in range(self.args.data_parallel_size_local):
|
|
child_args = _build_vllm_dp_server_args(self.args, local_rank)
|
|
process = context.Process(
|
|
target=_run_vllm_dp_server,
|
|
name=f"APIServer_DPRank_{child_args.data_parallel_rank}",
|
|
args=(child_args,),
|
|
)
|
|
process.start()
|
|
self._processes.append(process)
|
|
|
|
async def _probe_all_children(self) -> None:
|
|
"""
|
|
Background coroutine: probes all child endpoints on each interval.
|
|
|
|
Exits when any server becomes unhealthy after being ready, signalling
|
|
_monitor_children to initiate shutdown.
|
|
"""
|
|
timeout = aiohttp.ClientTimeout(total=self.args.dp_supervisor_probe_timeout_s)
|
|
async with aiohttp.ClientSession(timeout=timeout) as session:
|
|
while not self._shutdown_event.is_set():
|
|
threshold = (
|
|
self.args.dp_supervisor_probe_failure_threshold
|
|
if self._is_ready
|
|
else 1
|
|
)
|
|
results = await asyncio.gather(
|
|
*(
|
|
_probe_endpoint(
|
|
session,
|
|
self.args,
|
|
port,
|
|
"/health",
|
|
conn_err_failure_threshold=threshold,
|
|
conn_err_retry_delay=self.args.dp_supervisor_probe_interval_s,
|
|
)
|
|
for port in self.child_ports
|
|
),
|
|
return_exceptions=True,
|
|
)
|
|
all_healthy = all(r is True for r in results)
|
|
|
|
if all_healthy:
|
|
# If all healthy, we are ready to receive requests.
|
|
# This conditional avoids a potential race condition
|
|
# where shutdown is set, THEN the probe returns true.
|
|
if not self._shutdown_event.is_set():
|
|
self._is_ready = True
|
|
elif self._is_ready:
|
|
# Once ready, any failure in the probe means vLLM is dead.
|
|
num_unhealthy = sum(1 for r in results if r is not True)
|
|
logger.info(
|
|
"DPSupervisor probe found %s unhealthy DP Servers.",
|
|
num_unhealthy,
|
|
)
|
|
self._is_ready = False
|
|
self._shutdown_event.set()
|
|
return
|
|
|
|
with contextlib.suppress(asyncio.TimeoutError):
|
|
logger.debug(
|
|
"Waiting for %s seconds before next probe",
|
|
self.args.dp_supervisor_probe_interval_s,
|
|
)
|
|
await asyncio.wait_for(
|
|
self._shutdown_event.wait(),
|
|
timeout=self.args.dp_supervisor_probe_interval_s,
|
|
)
|
|
|
|
async def _monitor_children(self) -> None:
|
|
"""
|
|
Main coroutine task that monitors the children vLLM servers.
|
|
|
|
Before the vLLM servers are /ready:
|
|
- if the pid is dead, we will shut down
|
|
- if the probe fails, we try again after dp_supervisor_probe_interval_s
|
|
|
|
After the vLLM servers are /ready:
|
|
- if the pid is dead, we will shut down
|
|
- if the probe fails, we will shut down
|
|
"""
|
|
probe_task = asyncio.create_task(
|
|
self._probe_all_children(), name="dp-health-probe"
|
|
)
|
|
|
|
try:
|
|
while not self._shutdown_event.is_set():
|
|
# 1. Check for dead processes
|
|
n_failed = len([p for p in self._processes if not p.is_alive()])
|
|
if n_failed > 0:
|
|
logger.info("DPSupervisor found %s exited DP Servers.", n_failed)
|
|
break
|
|
|
|
# 2. Check if the probe background task crashed or failed.
|
|
if probe_task.done():
|
|
# Extract exception if it crashed, or log failure
|
|
exc = probe_task.exception() if not probe_task.cancelled() else None
|
|
logger.info("DPSupervisor probe task stopped. Exception: %s", exc)
|
|
break
|
|
|
|
# Sleep for probe_interval seconds or until a shutdown.
|
|
with contextlib.suppress(asyncio.TimeoutError):
|
|
logger.debug(
|
|
"Waiting for %s seconds before next monitor",
|
|
self.args.dp_supervisor_probe_interval_s,
|
|
)
|
|
await asyncio.wait_for(
|
|
self._shutdown_event.wait(),
|
|
timeout=self.args.dp_supervisor_probe_interval_s,
|
|
)
|
|
|
|
finally:
|
|
# Cleanup probe task if needed.
|
|
if not probe_task.done():
|
|
probe_task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await probe_task
|
|
|
|
async def _shutdown_children(self) -> None:
|
|
"""Terminate the vLLM DP servers."""
|
|
timeout = self.args.shutdown_timeout + CHILD_EXIT_GRACE_S
|
|
|
|
try:
|
|
logger.info(
|
|
"DPSupervisor forwarding %s to DP Servers.",
|
|
self._shutdown_signal.name,
|
|
)
|
|
for process in self._processes:
|
|
pid = process.pid
|
|
if not process.is_alive() or pid is None:
|
|
continue
|
|
with contextlib.suppress(ProcessLookupError, OSError):
|
|
os.kill(pid, self._shutdown_signal)
|
|
|
|
try:
|
|
await asyncio.to_thread(
|
|
_join_processes_with_timeout, self._processes, timeout
|
|
)
|
|
except asyncio.CancelledError:
|
|
logger.warning("Shutdown await cancelled")
|
|
raise
|
|
finally:
|
|
for process in self._processes:
|
|
pid = process.pid
|
|
if not process.is_alive() or pid is None:
|
|
continue
|
|
logger.warning(
|
|
"DP server %s did not exit within %.1fs; force killing.",
|
|
process.name,
|
|
timeout,
|
|
)
|
|
with contextlib.suppress(
|
|
ProcessLookupError,
|
|
OSError,
|
|
psutil.NoSuchProcess,
|
|
psutil.AccessDenied,
|
|
):
|
|
kill_process_tree(pid)
|
|
|
|
|
|
def run_dp_supervisor(args: argparse.Namespace) -> None:
|
|
uvloop.run(DPSupervisor(args).run())
|