| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """The `preload` subcommand for the vLLM CLI. |
| 4 | |
| 5 | Launches one weight cache daemon per GPU; each daemon holds the |
| 6 | post-quantized, TP-sharded weights of its rank in GPU memory and serves |
| 7 | CUDA IPC handles to vLLM engines over a Unix domain socket, so restarting |
| 8 | engines can map the weights via zero-copy IPC instead of reloading from disk: |
| 9 | |
| 10 | vllm preload --model /path/to/model --tensor-parallel-size 4 |
| 11 | |
| 12 | Engines then load from the daemons with: |
| 13 | |
| 14 | vllm serve /path/to/model --tensor-parallel-size 4 --load-format ipc_cache |
| 15 | |
| 16 | Tensor, expert and data parallelism are supported, including across nodes; |
| 17 | see vllm/model_executor/model_loader/weight_cache/daemon.py for the rank |
| 18 | layout and rendezvous details. |
| 19 | """ |
| 20 | |
| 21 | import argparse |
| 22 | import multiprocessing |
| 23 | import queue |
| 24 | import signal |
| 25 | import sys |
| 26 | import threading |
| 27 | import typing |
| 28 | from collections.abc import Sequence |
| 29 | from itertools import product |
| 30 | from multiprocessing.process import BaseProcess |
| 31 | |
| 32 | from fastapi import FastAPI |
| 33 | from fastapi.responses import Response |
| 34 | |
| 35 | from vllm.engine.arg_utils import EngineArgs |
| 36 | from vllm.entrypoints.cli.types import CLISubcommand |
| 37 | from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG |
| 38 | from vllm.logger import init_logger |
| 39 | from vllm.model_executor.model_loader.weight_cache.protocol import ( |
| 40 | check_ipc_platform_support, |
| 41 | ) |
| 42 | from vllm.utils.network_utils import get_distributed_init_method, get_open_port |
| 43 | |
| 44 | if typing.TYPE_CHECKING: |
| 45 | from vllm.config import VllmConfig |
| 46 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 47 | else: |
| 48 | FlexibleArgumentParser = argparse.ArgumentParser |
| 49 | |
| 50 | logger = init_logger(__name__) |
| 51 | |
| 52 | |
| 53 | class _HealthState: |
| 54 | def __init__(self, processes: Sequence[BaseProcess]) -> None: |
| 55 | self.processes = processes |
| 56 | self._ready = False |
| 57 | self._lock = threading.Lock() |
| 58 | |
| 59 | def mark_ready(self) -> None: |
| 60 | with self._lock: |
| 61 | self._ready = True |
| 62 | |
| 63 | @property |
| 64 | def is_ready(self) -> bool: |
| 65 | with self._lock: |
| 66 | ready = self._ready |
| 67 | return ready and not any(proc.exitcode is not None for proc in self.processes) |
| 68 | |
| 69 | |
| 70 | def _build_health_app(state: _HealthState) -> FastAPI: |
| 71 | app = FastAPI(openapi_url=None, docs_url=None, redoc_url=None) |
| 72 | |
| 73 | @app.get("/health", include_in_schema=False) |
| 74 | async def health() -> Response: |
| 75 | return Response(status_code=200 if state.is_ready else 503) |
| 76 | |
| 77 | return app |
| 78 | |
| 79 | |
| 80 | def _start_health_server( |
| 81 | host: str, port: int, state: _HealthState |
| 82 | ) -> tuple[typing.Any, threading.Thread]: |
| 83 | import uvicorn |
| 84 | |
| 85 | config = uvicorn.Config( |
| 86 | _build_health_app(state), |
| 87 | host=host, |
| 88 | port=port, |
| 89 | log_level="error", |
| 90 | access_log=False, |
| 91 | ) |
| 92 | server = uvicorn.Server(config) |
| 93 | thread = threading.Thread( |
| 94 | target=server.run, |
| 95 | name="vllm-weight-cache-health", |
| 96 | daemon=True, |
| 97 | ) |
| 98 | thread.start() |
| 99 | return server, thread |
| 100 | |
| 101 | |
| 102 | class PreloadSubcommand(CLISubcommand): |
| 103 | """The `preload` subcommand for the vLLM CLI.""" |
| 104 | |
| 105 | name = "preload" |
| 106 | |
| 107 | @staticmethod |
| 108 | def add_cli_args(parser: FlexibleArgumentParser) -> None: |
| 109 | EngineArgs.add_cli_args(parser) |
| 110 | parser.add_argument( |
| 111 | "--weight-cache-socket-dir", |
| 112 | type=str, |
| 113 | default=None, |
| 114 | help="Directory for the daemon Unix sockets (default: tempdir).", |
| 115 | ) |
| 116 | parser.add_argument( |
| 117 | "--weight-cache-master-port", |
| 118 | type=int, |
| 119 | default=None, |
| 120 | help="Rendezvous port for the daemon's own TP group. Must differ " |
| 121 | "from the engine's --master-port (the daemon holds its group open " |
| 122 | "while serving) and match across nodes. Required when --nnodes > " |
| 123 | "1; defaults to a free port for single-node.", |
| 124 | ) |
| 125 | parser.add_argument( |
| 126 | "--weight-cache-draft-master-port", |
| 127 | type=int, |
| 128 | default=None, |
| 129 | help="Rendezvous port for the MTP draft daemon group. Defaults to " |
| 130 | "--weight-cache-master-port + 1 for multi-node, or a free port " |
| 131 | "for single-node.", |
| 132 | ) |
| 133 | parser.add_argument( |
| 134 | "--weight-cache-health-port", |
| 135 | type=int, |
| 136 | default=None, |
| 137 | help="Port for the parent /health endpoint; disabled by default.", |
| 138 | ) |
| 139 | parser.add_argument( |
| 140 | "--weight-cache-health-host", |
| 141 | type=str, |
| 142 | default="0.0.0.0", |
| 143 | help="Host for the parent /health endpoint (default: 0.0.0.0).", |
| 144 | ) |
| 145 | |
| 146 | @staticmethod |
| 147 | def cmd(args: argparse.Namespace) -> None: |
| 148 | # Imported here so `vllm --help` and friends stay free of torch. |
| 149 | from vllm.model_executor.model_loader.weight_cache.daemon import ( |
| 150 | _reject_unsupported_parallelism, |
| 151 | _run_daemon, |
| 152 | get_draft_daemon_config, |
| 153 | plan_local_ranks, |
| 154 | ) |
| 155 | from vllm.model_executor.model_loader.weight_cache.utils import ( |
| 156 | format_daemon_role, |
| 157 | ) |
| 158 | |
| 159 | engine_args = EngineArgs.from_cli_args(args) |
| 160 | vllm_config = engine_args.create_engine_config() |
| 161 | if vllm_config.load_config.load_format == "ipc_cache": |
| 162 | raise ValueError( |
| 163 | "The weight cache daemon itself must load from disk; use the " |
| 164 | "default --load-format" |
| 165 | ) |
| 166 | # Config-only so it can fail before any model loading; the quant |
| 167 | # method check needs the created model and runs in |
| 168 | # WeightCacheDaemon.get_model. |
| 169 | check_ipc_platform_support() |
| 170 | parallel_config = vllm_config.parallel_config |
| 171 | _reject_unsupported_parallelism(parallel_config) |
| 172 | tp_size = parallel_config.tensor_parallel_size |
| 173 | pp_size = parallel_config.pipeline_parallel_size |
| 174 | dp_size = parallel_config.data_parallel_size |
| 175 | |
| 176 | placements = plan_local_ranks(parallel_config) |
| 177 | local_world_size = len(placements) |
| 178 | if dp_size == 1: |
| 179 | nnodes = parallel_config.nnodes |
| 180 | node_rank = parallel_config.node_rank |
| 181 | master_addr = parallel_config.master_addr |
| 182 | else: |
| 183 | master_addr = parallel_config.data_parallel_master_ip |
| 184 | if parallel_config.nnodes > 1: |
| 185 | nnodes = parallel_config.nnodes |
| 186 | node_rank = parallel_config.node_rank |
| 187 | else: |
| 188 | nnodes = dp_size // parallel_config.data_parallel_size_local |
| 189 | node_rank = parallel_config.data_parallel_rank // ( |
| 190 | parallel_config.data_parallel_size_local |
| 191 | ) |
| 192 | if nnodes > 1 and master_addr in ("127.0.0.1", "localhost"): |
| 193 | raise ValueError( |
| 194 | "Data parallelism across nodes requires a reachable " |
| 195 | "--data-parallel-address for the daemon rendezvous" |
| 196 | ) |
| 197 | |
| 198 | # The daemon forms its own world group and holds it open while |
| 199 | # serving, so it needs a rendezvous port distinct from the engine's. |
| 200 | # All nodes must agree on it; single-node can auto-pick a free port. |
| 201 | # The master address is the engine's --master-addr (TP across nodes) |
| 202 | # or --data-parallel-address. |
| 203 | if nnodes > 1 and args.weight_cache_master_port is None: |
| 204 | raise ValueError( |
| 205 | "--weight-cache-master-port is required when the daemons span nodes" |
| 206 | ) |
| 207 | master_port = args.weight_cache_master_port or get_open_port() |
| 208 | distributed_init_method = get_distributed_init_method(master_addr, master_port) |
| 209 | |
| 210 | # (is_draft, vllm_config, rendezvous) per daemon group. |
| 211 | groups: list[tuple[bool, VllmConfig, str]] = [ |
| 212 | (False, vllm_config, distributed_init_method) |
| 213 | ] |
| 214 | draft_vllm_config = get_draft_daemon_config(vllm_config) |
| 215 | if draft_vllm_config is not None: |
| 216 | draft_master_port = args.weight_cache_draft_master_port or ( |
| 217 | master_port + 1 if nnodes > 1 else get_open_port() |
| 218 | ) |
| 219 | if draft_master_port == master_port: |
| 220 | raise ValueError( |
| 221 | "--weight-cache-draft-master-port must differ from " |
| 222 | "--weight-cache-master-port" |
| 223 | ) |
| 224 | groups.append( |
| 225 | ( |
| 226 | True, |
| 227 | draft_vllm_config, |
| 228 | get_distributed_init_method(master_addr, draft_master_port), |
| 229 | ) |
| 230 | ) |
| 231 | |
| 232 | ctx = multiprocessing.get_context("spawn") |
| 233 | ready_queue: multiprocessing.Queue[tuple[str, int]] = ctx.Queue() |
| 234 | # Local index == device index; global rank enumerates DP, PP, then TP. |
| 235 | expected_ready = { |
| 236 | ( |
| 237 | format_daemon_role(is_draft), |
| 238 | dp_rank * pp_size * tp_size + pp_rank * tp_size + tp_rank, |
| 239 | ) |
| 240 | for (is_draft, _, _), (_, dp_rank, pp_rank, tp_rank) in product( |
| 241 | groups, placements |
| 242 | ) |
| 243 | } |
| 244 | procs = [ |
| 245 | ctx.Process( |
| 246 | target=_run_daemon, |
| 247 | args=( |
| 248 | dp_rank * pp_size * tp_size + pp_rank * tp_size + tp_rank, |
| 249 | local_rank, |
| 250 | config, |
| 251 | init_method, |
| 252 | args.weight_cache_socket_dir, |
| 253 | ready_queue, |
| 254 | is_draft, |
| 255 | dp_rank, |
| 256 | pp_rank, |
| 257 | ), |
| 258 | name=( |
| 259 | f"vllm-weight-cache-{format_daemon_role(is_draft)}-" |
| 260 | f"{(dp_rank * pp_size * tp_size + pp_rank * tp_size + tp_rank)}" |
| 261 | ), |
| 262 | ) |
| 263 | for (is_draft, config, init_method), ( |
| 264 | local_rank, |
| 265 | dp_rank, |
| 266 | pp_rank, |
| 267 | tp_rank, |
| 268 | ) in product(groups, placements) |
| 269 | ] |
| 270 | health_state = _HealthState(procs) |
| 271 | health_server = None |
| 272 | health_thread = None |
| 273 | if args.weight_cache_health_port is not None: |
| 274 | health_server, health_thread = _start_health_server( |
| 275 | args.weight_cache_health_host, |
| 276 | args.weight_cache_health_port, |
| 277 | health_state, |
| 278 | ) |
| 279 | for proc in procs: |
| 280 | proc.start() |
| 281 | |
| 282 | def _shutdown(signum, frame): |
| 283 | if health_server is not None: |
| 284 | health_server.should_exit = True |
| 285 | for proc in procs: |
| 286 | proc.terminate() |
| 287 | |
| 288 | signal.signal(signal.SIGINT, _shutdown) |
| 289 | signal.signal(signal.SIGTERM, _shutdown) |
| 290 | |
| 291 | ready: set[tuple[str, int]] = set() |
| 292 | while len(ready) < len(expected_ready): |
| 293 | try: |
| 294 | ready.add(ready_queue.get(timeout=1.0)) |
| 295 | except queue.Empty: |
| 296 | dead = [p for p in procs if p.exitcode is not None] |
| 297 | if dead: |
| 298 | logger.error( |
| 299 | "Weight cache daemon rank(s) exited during startup " |
| 300 | "(exitcodes=%s); shutting down.", |
| 301 | [p.exitcode for p in dead], |
| 302 | ) |
| 303 | for proc in procs: |
| 304 | proc.terminate() |
| 305 | for proc in procs: |
| 306 | proc.join() |
| 307 | if health_server is not None: |
| 308 | health_server.should_exit = True |
| 309 | if health_thread is not None: |
| 310 | health_thread.join() |
| 311 | sys.exit(max((p.exitcode or 0) for p in procs)) |
| 312 | health_state.mark_ready() |
| 313 | socket_dir_msg = args.weight_cache_socket_dir or "the default socket dir" |
| 314 | logger.info_once( |
| 315 | "===== Weight cache daemon READY: node %d/%d serving %d local " |
| 316 | "rank(s) x %d role(s) in %s =====", |
| 317 | node_rank, |
| 318 | nnodes, |
| 319 | local_world_size, |
| 320 | len(groups), |
| 321 | socket_dir_msg, |
| 322 | ) |
| 323 | |
| 324 | for proc in procs: |
| 325 | proc.join() |
| 326 | if health_server is not None: |
| 327 | health_server.should_exit = True |
| 328 | if health_thread is not None: |
| 329 | health_thread.join() |
| 330 | sys.exit(max(proc.exitcode or 0 for proc in procs)) |
| 331 | |
| 332 | def subparser_init( |
| 333 | self, subparsers: argparse._SubParsersAction |
| 334 | ) -> FlexibleArgumentParser: |
| 335 | preload_parser = subparsers.add_parser( |
| 336 | self.name, |
| 337 | help="Launch weight cache daemons (one per GPU) for fast engine restarts.", |
| 338 | description="Launch weight cache daemons (one per GPU) that " |
| 339 | "keep post-quantized weights in GPU memory and serve them to " |
| 340 | "engines over CUDA IPC.", |
| 341 | usage="vllm preload --model <model> [options]", |
| 342 | ) |
| 343 | self.add_cli_args(preload_parser) |
| 344 | preload_parser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(subcmd=self.name) |
| 345 | return preload_parser |
| 346 | |
| 347 | |
| 348 | def cmd_init() -> list[CLISubcommand]: |
| 349 | return [PreloadSubcommand()] |