| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import argparse |
| 8 | import asyncio |
| 9 | import contextlib |
| 10 | import copy |
| 11 | import multiprocessing |
| 12 | import os |
| 13 | import signal |
| 14 | import time |
| 15 | from functools import partial |
| 16 | from http import HTTPStatus |
| 17 | from multiprocessing.process import BaseProcess |
| 18 | |
| 19 | import aiohttp |
| 20 | import psutil |
| 21 | import uvicorn |
| 22 | import uvloop |
| 23 | from fastapi import FastAPI, Response |
| 24 | |
| 25 | import vllm.envs as envs |
| 26 | from vllm.logger import configure_logging_from_args, init_logger |
| 27 | from vllm.utils.system_utils import ( |
| 28 | decorate_logs, |
| 29 | kill_process_tree, |
| 30 | set_process_title, |
| 31 | ) |
| 32 | from vllm.v1.engine.utils import get_engine_process_shutdown_timeout |
| 33 | |
| 34 | from .launcher import NoSignalServer |
| 35 | from .utils.server_utils import get_uvicorn_log_config |
| 36 | |
| 37 | logger = init_logger(__name__) |
| 38 | |
| 39 | CHILD_EXIT_GRACE_S = 5.0 |
| 40 | |
| 41 | |
| 42 | def infer_multi_port_external_lb_start_rank(args: argparse.Namespace) -> int: |
| 43 | start_rank = getattr(args, "data_parallel_start_rank", None) |
| 44 | if start_rank is not None: |
| 45 | return start_rank |
| 46 | |
| 47 | node_rank = getattr(args, "node_rank", 0) or 0 |
| 48 | local_size = getattr(args, "data_parallel_size_local", 0) or 0 |
| 49 | return node_rank * local_size |
| 50 | |
| 51 | |
| 52 | def validate_multi_port_external_lb_args(args: argparse.Namespace) -> None: |
| 53 | if getattr(args, "grpc", False): |
| 54 | raise ValueError( |
| 55 | "Error: --data-parallel-multi-port-external-lb does not support --grpc" |
| 56 | ) |
| 57 | if args.uds is not None: |
| 58 | raise ValueError( |
| 59 | "Error: --data-parallel-multi-port-external-lb does not support --uds" |
| 60 | ) |
| 61 | if bool(args.ssl_keyfile) != bool(args.ssl_certfile): |
| 62 | raise ValueError( |
| 63 | "Error: --ssl-keyfile and --ssl-certfile must be provided together" |
| 64 | ) |
| 65 | if args.api_server_count not in (None, 1): |
| 66 | raise ValueError( |
| 67 | "Error: --data-parallel-multi-port-external-lb currently requires " |
| 68 | "--api-server-count=1" |
| 69 | ) |
| 70 | if args.data_parallel_rank is not None: |
| 71 | raise ValueError( |
| 72 | "Error: --data-parallel-multi-port-external-lb manages child " |
| 73 | "--data-parallel-rank values internally" |
| 74 | ) |
| 75 | if args.data_parallel_external_lb or args.data_parallel_hybrid_lb: |
| 76 | raise ValueError( |
| 77 | "Error: --data-parallel-multi-port-external-lb cannot be combined with " |
| 78 | "--data-parallel-external-lb or --data-parallel-hybrid-lb" |
| 79 | ) |
| 80 | if args.data_parallel_size < 2: |
| 81 | raise ValueError( |
| 82 | "Error: --data-parallel-multi-port-external-lb requires " |
| 83 | "--data-parallel-size > 1" |
| 84 | ) |
| 85 | |
| 86 | local_size = args.data_parallel_size_local |
| 87 | if local_size is None or local_size < 2: |
| 88 | raise ValueError( |
| 89 | "Error: --data-parallel-multi-port-external-lb requires " |
| 90 | "--data-parallel-size-local >= 2" |
| 91 | ) |
| 92 | if local_size > args.data_parallel_size: |
| 93 | raise ValueError( |
| 94 | "Error: --data-parallel-size-local cannot exceed --data-parallel-size" |
| 95 | ) |
| 96 | if args.data_parallel_size % local_size != 0: |
| 97 | raise ValueError( |
| 98 | "Error: --data-parallel-size must be divisible by " |
| 99 | "--data-parallel-size-local" |
| 100 | ) |
| 101 | |
| 102 | start_rank = infer_multi_port_external_lb_start_rank(args) |
| 103 | if start_rank + local_size > args.data_parallel_size: |
| 104 | raise ValueError( |
| 105 | "Error: multi-port supervised ranks would exceed --data-parallel-size" |
| 106 | ) |
| 107 | |
| 108 | supervisor_port = args.data_parallel_supervisor_port |
| 109 | child_port_min = args.port |
| 110 | child_port_max = args.port + local_size - 1 |
| 111 | if child_port_min <= supervisor_port <= child_port_max: |
| 112 | raise ValueError( |
| 113 | f"Error: --data-parallel-supervisor-port {supervisor_port} " |
| 114 | f"overlaps with child rank ports {child_port_min}-{child_port_max}" |
| 115 | ) |
| 116 | |
| 117 | |
| 118 | def _build_vllm_dp_server_args( |
| 119 | args: argparse.Namespace, local_rank: int |
| 120 | ) -> argparse.Namespace: |
| 121 | child_args = copy.copy(args) |
| 122 | child_args.port = args.port + local_rank |
| 123 | child_args.data_parallel_rank = ( |
| 124 | infer_multi_port_external_lb_start_rank(args) + local_rank |
| 125 | ) |
| 126 | child_args.data_parallel_start_rank = None |
| 127 | child_args.data_parallel_size_local = 1 |
| 128 | child_args.data_parallel_external_lb = True |
| 129 | child_args.data_parallel_hybrid_lb = False |
| 130 | child_args.data_parallel_multi_port_external_lb = False |
| 131 | child_args.data_parallel_supervisor_port = None |
| 132 | child_args.api_server_count = 1 |
| 133 | child_args.device_ids = _build_device_ids(args, local_rank) |
| 134 | return child_args |
| 135 | |
| 136 | |
| 137 | def _build_device_ids(args: argparse.Namespace, local_rank: int) -> list[int | str]: |
| 138 | """Build the --device-ids value for a DP child process. |
| 139 | |
| 140 | The child resolves these against its own inherited device-control env |
| 141 | var (e.g. CUDA_VISIBLE_DEVICES), so integer IDs must stay env-relative |
| 142 | here rather than being translated to physical IDs. |
| 143 | """ |
| 144 | devices_per_rank = args.tensor_parallel_size * args.pipeline_parallel_size |
| 145 | start = local_rank * devices_per_rank |
| 146 | stop = start + devices_per_rank |
| 147 | device_ids = getattr(args, "device_ids", None) |
| 148 | if device_ids is not None: |
| 149 | if stop > len(device_ids): |
| 150 | raise ValueError( |
| 151 | f"--device-ids has {len(device_ids)} entries, but DP rank " |
| 152 | f"{local_rank} needs devices [{start}, {stop})" |
| 153 | ) |
| 154 | return device_ids[start:stop] |
| 155 | return list(range(start, stop)) |
| 156 | |
| 157 | |
| 158 | def _child_base_url(args: argparse.Namespace, port: int) -> str: |
| 159 | host = args.host or "127.0.0.1" |
| 160 | if host == "0.0.0.0": |
| 161 | host = "127.0.0.1" |
| 162 | elif host == "::": |
| 163 | host = "::1" |
| 164 | scheme = "https" if args.ssl_keyfile and args.ssl_certfile else "http" |
| 165 | return f"{scheme}://{host}:{port}" |
| 166 | |
| 167 | |
| 168 | def _join_processes_with_timeout(processes: list[BaseProcess], timeout: float) -> None: |
| 169 | deadline = time.monotonic() + timeout |
| 170 | for process in processes: |
| 171 | remaining = deadline - time.monotonic() |
| 172 | if remaining <= 0: |
| 173 | break |
| 174 | process.join(timeout=remaining) |
| 175 | |
| 176 | |
| 177 | async def _probe_endpoint( |
| 178 | session: aiohttp.ClientSession, |
| 179 | args: argparse.Namespace, |
| 180 | port: int, |
| 181 | path: str, |
| 182 | conn_err_failure_threshold: int = 3, |
| 183 | conn_err_retry_delay: float = 5.0, |
| 184 | ) -> bool: |
| 185 | """Probe /health endpoint for 200 status. |
| 186 | |
| 187 | If there is a connection error, retry every N seconds. |
| 188 | """ |
| 189 | for iteration in range(conn_err_failure_threshold): |
| 190 | try: |
| 191 | probe_ssl = None |
| 192 | if args.ssl_keyfile and args.ssl_certfile: |
| 193 | # Probes target node-local child servers over loopback, so skip |
| 194 | # certificate verification to avoid SAN/hostname mismatches for |
| 195 | # localhost/127.0.0.1 deployments. |
| 196 | probe_ssl = False |
| 197 | async with session.get( |
| 198 | _child_base_url(args, port) + path, ssl=probe_ssl |
| 199 | ) as response: |
| 200 | # vLLM returns 503 on EngineDeadError, so we should return |
| 201 | # immediately if vLLM responds with a non-200 status code. |
| 202 | return response.status == HTTPStatus.OK |
| 203 | except (aiohttp.ClientError, asyncio.TimeoutError) as e: |
| 204 | # Allow retry of connection errors. |
| 205 | logger.debug( |
| 206 | "Probe attempt %d/%d failed on port %d: %r", |
| 207 | iteration + 1, |
| 208 | conn_err_failure_threshold, |
| 209 | port, |
| 210 | e, |
| 211 | ) |
| 212 | |
| 213 | if iteration < conn_err_failure_threshold - 1: |
| 214 | await asyncio.sleep(conn_err_retry_delay) |
| 215 | |
| 216 | return False |
| 217 | |
| 218 | |
| 219 | def _build_dp_supervisor_app(supervisor: DPSupervisor) -> FastAPI: |
| 220 | app = FastAPI(openapi_url=None, docs_url=None, redoc_url=None) |
| 221 | app.state.supervisor = supervisor |
| 222 | |
| 223 | def _status_response(ok: bool) -> Response: |
| 224 | return Response( |
| 225 | status_code=(HTTPStatus.OK if ok else HTTPStatus.SERVICE_UNAVAILABLE) |
| 226 | ) |
| 227 | |
| 228 | @app.get("/health", include_in_schema=False) |
| 229 | async def health() -> Response: |
| 230 | return _status_response(app.state.supervisor.is_ready) |
| 231 | |
| 232 | @app.get("/ready", include_in_schema=False) |
| 233 | @app.get("/readyz", include_in_schema=False) |
| 234 | async def ready() -> Response: |
| 235 | return _status_response(app.state.supervisor.is_ready) |
| 236 | |
| 237 | return app |
| 238 | |
| 239 | |
| 240 | def _run_python_vllm_dp_server(child_args: argparse.Namespace) -> None: |
| 241 | from vllm.entrypoints.launchers.api_server.entry import run_server |
| 242 | |
| 243 | uvloop.run(run_server(child_args)) |
| 244 | |
| 245 | |
| 246 | def _run_rust_vllm_dp_server(child_args: argparse.Namespace) -> None: |
| 247 | from vllm.entrypoints.cli.serve import run_multi_api_server |
| 248 | |
| 249 | run_multi_api_server(child_args) |
| 250 | |
| 251 | |
| 252 | def _run_vllm_dp_server(child_args: argparse.Namespace) -> None: |
| 253 | """Entrypoint function for the vLLM DP Server.""" |
| 254 | configure_logging_from_args(child_args) |
| 255 | |
| 256 | # Create a fresh process group for the vLLM DP Server, |
| 257 | # so that CTRL-C is propagated cleanly. |
| 258 | os.setpgrp() |
| 259 | |
| 260 | name = f"APIServer_DP{child_args.data_parallel_rank}" |
| 261 | set_process_title(name) |
| 262 | decorate_logs(name) |
| 263 | if envs.VLLM_USE_RUST_FRONTEND and envs.VLLM_RUST_FRONTEND_PATH: |
| 264 | _run_rust_vllm_dp_server(child_args) |
| 265 | else: |
| 266 | _run_python_vllm_dp_server(child_args) |
| 267 | |
| 268 | |
| 269 | class DPSupervisor: |
| 270 | def __init__(self, args: argparse.Namespace): |
| 271 | validate_multi_port_external_lb_args(args) |
| 272 | self.args = args |
| 273 | self.supervisor_port = args.data_parallel_supervisor_port |
| 274 | self.child_ports = [ |
| 275 | args.port + local_rank |
| 276 | for local_rank in range(args.data_parallel_size_local) |
| 277 | ] |
| 278 | self._is_ready = False |
| 279 | self._processes: list[BaseProcess] = [] |
| 280 | self._shutdown_event = asyncio.Event() |
| 281 | self._shutdown_signal = signal.SIGTERM |
| 282 | |
| 283 | @property |
| 284 | def is_ready(self) -> bool: |
| 285 | return self._is_ready and not self._shutdown_event.is_set() |
| 286 | |
| 287 | async def run(self) -> None: |
| 288 | loop = asyncio.get_running_loop() |
| 289 | decorate_logs("DPSupervisor") |
| 290 | |
| 291 | # K8s sends SIGTERM for shutdown - begin graceful termination. |
| 292 | for sig in (signal.SIGTERM, signal.SIGINT): |
| 293 | loop.add_signal_handler(sig, partial(self._handle_signal, sig)) |
| 294 | |
| 295 | supervisor_server: uvicorn.Server | None = None |
| 296 | supervisor_server_task: asyncio.Task[None] | None = None |
| 297 | try: |
| 298 | # Launch vLLM DP Servers and begin monitoring them. |
| 299 | self._start_children() |
| 300 | monitor_task = asyncio.create_task( |
| 301 | self._monitor_children(), name="dp-monitor" |
| 302 | ) |
| 303 | |
| 304 | # Only start the DPSupervisor server once the vLLM DP Servers are |
| 305 | # ready. This avoids the supervisor /health endpoint returning 503 |
| 306 | # to external load balancer probes while the engines initialize. |
| 307 | await self._wait_until_ready(monitor_task) |
| 308 | if self.is_ready and not monitor_task.done(): |
| 309 | supervisor_server, supervisor_server_task = await self._start_server() |
| 310 | |
| 311 | await monitor_task |
| 312 | finally: |
| 313 | self._is_ready = False |
| 314 | await self._shutdown_children() |
| 315 | |
| 316 | # Shutdown the DP Supervisor server. |
| 317 | if supervisor_server is not None and supervisor_server_task is not None: |
| 318 | supervisor_server.should_exit = True |
| 319 | await supervisor_server_task |
| 320 | |
| 321 | async def _start_server(self) -> tuple[uvicorn.Server, asyncio.Task[None]]: |
| 322 | """Launch the DPSupervisor HTTP server. |
| 323 | |
| 324 | Called only after the vLLM DP Servers are ready so that /health does |
| 325 | not return 503 to external probes while the engines are initializing. |
| 326 | """ |
| 327 | app = _build_dp_supervisor_app(self) |
| 328 | host = self.args.host or "0.0.0.0" |
| 329 | |
| 330 | uvicorn_kwargs: dict = {} |
| 331 | log_config = get_uvicorn_log_config(self.args) |
| 332 | if log_config is not None: |
| 333 | uvicorn_kwargs["log_config"] = log_config |
| 334 | |
| 335 | config = uvicorn.Config( |
| 336 | app, |
| 337 | host=host, |
| 338 | port=self.supervisor_port, |
| 339 | log_level=self.args.uvicorn_log_level, |
| 340 | access_log=not self.args.disable_uvicorn_access_log, |
| 341 | ssl_keyfile=self.args.ssl_keyfile, |
| 342 | ssl_certfile=self.args.ssl_certfile, |
| 343 | ssl_ca_certs=self.args.ssl_ca_certs, |
| 344 | ssl_cert_reqs=self.args.ssl_cert_reqs, |
| 345 | ssl_ciphers=self.args.ssl_ciphers, |
| 346 | **uvicorn_kwargs, |
| 347 | ) |
| 348 | supervisor_server = NoSignalServer(config) |
| 349 | supervisor_server_task = asyncio.create_task( |
| 350 | supervisor_server.serve(), |
| 351 | name="dp-supervisor", |
| 352 | ) |
| 353 | supervisor_server_task.add_done_callback( |
| 354 | lambda _task: self._shutdown_event.set() |
| 355 | ) |
| 356 | |
| 357 | # Ensure DPSupervisor task starts on the event loop. |
| 358 | while not supervisor_server.started: |
| 359 | if supervisor_server_task.done(): |
| 360 | supervisor_server_task.result() |
| 361 | raise RuntimeError("DPSupervisor exited before startup.") |
| 362 | await asyncio.sleep(0.05) |
| 363 | logger.info("Started DPSupervisor on %s:%d", host, self.supervisor_port) |
| 364 | return supervisor_server, supervisor_server_task |
| 365 | |
| 366 | async def _wait_until_ready(self, monitor_task: asyncio.Task[None]) -> None: |
| 367 | """Block until the vLLM DP Servers are ready or shutdown is triggered. |
| 368 | |
| 369 | Returns early if monitoring stops (e.g. a DP Server dies during |
| 370 | startup), in which case the supervisor server is never started. |
| 371 | """ |
| 372 | logger.info("Waiting for vLLM DP Servers to become ready.") |
| 373 | while not self._is_ready and not self._shutdown_event.is_set(): |
| 374 | if monitor_task.done(): |
| 375 | return |
| 376 | await asyncio.sleep(0.05) |
| 377 | |
| 378 | def _handle_signal(self, signum: int) -> None: |
| 379 | """Signal handler that is added to the event loop. |
| 380 | |
| 381 | This catches the SIGTERM from K8s and begins graceful shutdown, |
| 382 | by setting the _shutdown_event(), which is watched by the main |
| 383 | coroutine monitoring the vLLM DP Servers. |
| 384 | """ |
| 385 | if self._shutdown_event.is_set(): |
| 386 | return |
| 387 | |
| 388 | self._shutdown_signal = signal.Signals(signum) |
| 389 | logger.info( |
| 390 | "DPSupervisor received %s, shutting down.", |
| 391 | self._shutdown_signal.name, |
| 392 | ) |
| 393 | |
| 394 | self._shutdown_event.set() |
| 395 | self._is_ready = False |
| 396 | |
| 397 | def _start_children(self) -> None: |
| 398 | """Launch vLLM DP Servers on separate GPUs.""" |
| 399 | logger.info("Launching vLLM DP Servers") |
| 400 | context = multiprocessing.get_context("spawn") |
| 401 | for local_rank in range(self.args.data_parallel_size_local): |
| 402 | child_args = _build_vllm_dp_server_args(self.args, local_rank) |
| 403 | process = context.Process( |
| 404 | target=_run_vllm_dp_server, |
| 405 | name=f"APIServer_DPRank_{child_args.data_parallel_rank}", |
| 406 | args=(child_args,), |
| 407 | ) |
| 408 | process.start() |
| 409 | self._processes.append(process) |
| 410 | |
| 411 | async def _probe_all_children(self) -> None: |
| 412 | """Background coroutine: probes all child endpoints on each interval. |
| 413 | |
| 414 | Exits when any server becomes unhealthy after being ready, signalling |
| 415 | _monitor_children to initiate shutdown. |
| 416 | """ |
| 417 | timeout = aiohttp.ClientTimeout(total=self.args.dp_supervisor_probe_timeout_s) |
| 418 | async with aiohttp.ClientSession(timeout=timeout) as session: |
| 419 | while not self._shutdown_event.is_set(): |
| 420 | threshold = ( |
| 421 | self.args.dp_supervisor_probe_failure_threshold |
| 422 | if self._is_ready |
| 423 | else 1 |
| 424 | ) |
| 425 | results = await asyncio.gather( |
| 426 | *( |
| 427 | _probe_endpoint( |
| 428 | session, |
| 429 | self.args, |
| 430 | port, |
| 431 | "/health", |
| 432 | conn_err_failure_threshold=threshold, |
| 433 | conn_err_retry_delay=self.args.dp_supervisor_probe_interval_s, |
| 434 | ) |
| 435 | for port in self.child_ports |
| 436 | ), |
| 437 | return_exceptions=True, |
| 438 | ) |
| 439 | all_healthy = all(r is True for r in results) |
| 440 | |
| 441 | if all_healthy: |
| 442 | # If all healthy, we are ready to receive requests. |
| 443 | # This conditional avoids a potential race condition |
| 444 | # where shutdown is set, THEN the probe returns true. |
| 445 | if not self._shutdown_event.is_set(): |
| 446 | self._is_ready = True |
| 447 | elif self._is_ready: |
| 448 | # Once ready, any failure in the probe means vLLM is dead. |
| 449 | num_unhealthy = sum(1 for r in results if r is not True) |
| 450 | logger.info( |
| 451 | "DPSupervisor probe found %s unhealthy DP Servers.", |
| 452 | num_unhealthy, |
| 453 | ) |
| 454 | self._is_ready = False |
| 455 | self._shutdown_event.set() |
| 456 | return |
| 457 | |
| 458 | with contextlib.suppress(asyncio.TimeoutError): |
| 459 | logger.debug( |
| 460 | "Waiting for %s seconds before next probe", |
| 461 | self.args.dp_supervisor_probe_interval_s, |
| 462 | ) |
| 463 | await asyncio.wait_for( |
| 464 | self._shutdown_event.wait(), |
| 465 | timeout=self.args.dp_supervisor_probe_interval_s, |
| 466 | ) |
| 467 | |
| 468 | async def _monitor_children(self) -> None: |
| 469 | """Main coroutine task that monitors the children vLLM servers. |
| 470 | |
| 471 | Before the vLLM servers are /ready: |
| 472 | - if the pid is dead, we will shut down |
| 473 | - if the probe fails, we try again after dp_supervisor_probe_interval_s |
| 474 | |
| 475 | After the vLLM servers are /ready: |
| 476 | - if the pid is dead, we will shut down |
| 477 | - if the probe fails, we will shut down |
| 478 | """ |
| 479 | probe_task = asyncio.create_task( |
| 480 | self._probe_all_children(), name="dp-health-probe" |
| 481 | ) |
| 482 | |
| 483 | try: |
| 484 | while not self._shutdown_event.is_set(): |
| 485 | # 1. Check for dead processes |
| 486 | n_failed = len([p for p in self._processes if not p.is_alive()]) |
| 487 | if n_failed > 0: |
| 488 | logger.info("DPSupervisor found %s exited DP Servers.", n_failed) |
| 489 | break |
| 490 | |
| 491 | # 2. Check if the probe background task crashed or failed. |
| 492 | if probe_task.done(): |
| 493 | # Extract exception if it crashed, or log failure |
| 494 | exc = probe_task.exception() if not probe_task.cancelled() else None |
| 495 | logger.info("DPSupervisor probe task stopped. Exception: %s", exc) |
| 496 | break |
| 497 | |
| 498 | # Sleep for probe_interval seconds or until a shutdown. |
| 499 | with contextlib.suppress(asyncio.TimeoutError): |
| 500 | logger.debug( |
| 501 | "Waiting for %s seconds before next monitor", |
| 502 | self.args.dp_supervisor_probe_interval_s, |
| 503 | ) |
| 504 | await asyncio.wait_for( |
| 505 | self._shutdown_event.wait(), |
| 506 | timeout=self.args.dp_supervisor_probe_interval_s, |
| 507 | ) |
| 508 | |
| 509 | finally: |
| 510 | # Cleanup probe task if needed. |
| 511 | if not probe_task.done(): |
| 512 | probe_task.cancel() |
| 513 | with contextlib.suppress(asyncio.CancelledError): |
| 514 | await probe_task |
| 515 | |
| 516 | async def _shutdown_children(self) -> None: |
| 517 | """Terminate the vLLM DP servers.""" |
| 518 | process_timeout = get_engine_process_shutdown_timeout( |
| 519 | self.args.shutdown_timeout, self.args.shutdown_timeout |
| 520 | ) |
| 521 | assert process_timeout is not None |
| 522 | timeout = process_timeout + CHILD_EXIT_GRACE_S |
| 523 | |
| 524 | try: |
| 525 | logger.info( |
| 526 | "DPSupervisor forwarding %s to DP Servers.", |
| 527 | self._shutdown_signal.name, |
| 528 | ) |
| 529 | for process in self._processes: |
| 530 | pid = process.pid |
| 531 | if not process.is_alive() or pid is None: |
| 532 | continue |
| 533 | with contextlib.suppress(ProcessLookupError, OSError): |
| 534 | os.kill(pid, self._shutdown_signal) |
| 535 | |
| 536 | try: |
| 537 | await asyncio.to_thread( |
| 538 | _join_processes_with_timeout, self._processes, timeout |
| 539 | ) |
| 540 | except asyncio.CancelledError: |
| 541 | logger.warning("Shutdown await cancelled") |
| 542 | raise |
| 543 | finally: |
| 544 | for process in self._processes: |
| 545 | pid = process.pid |
| 546 | if not process.is_alive() or pid is None: |
| 547 | continue |
| 548 | logger.warning( |
| 549 | "DP server %s did not exit within %.1fs; force killing.", |
| 550 | process.name, |
| 551 | timeout, |
| 552 | ) |
| 553 | with contextlib.suppress( |
| 554 | ProcessLookupError, |
| 555 | OSError, |
| 556 | psutil.NoSuchProcess, |
| 557 | psutil.AccessDenied, |
| 558 | ): |
| 559 | kill_process_tree(pid) |
| 560 | |
| 561 | |
| 562 | def run_dp_supervisor(args: argparse.Namespace) -> None: |
| 563 | configure_logging_from_args(args) |
| 564 | uvloop.run(DPSupervisor(args).run()) |