| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import argparse |
| 5 | import asyncio |
| 6 | import json |
| 7 | import os |
| 8 | import stat |
| 9 | import sys |
| 10 | from dataclasses import dataclass |
| 11 | from pathlib import Path |
| 12 | from typing import Any |
| 13 | |
| 14 | from pydantic import ValidationError |
| 15 | |
| 16 | from vllm.snapshot.manifest import ReleaseMarker, _validation_path, _write_json_atomic |
| 17 | from vllm.snapshot.types import Oracle, oracles_match |
| 18 | |
| 19 | _CANARY_PROMPT = "The capital of France is" |
| 20 | |
| 21 | |
| 22 | class SnapshotCanaryError(RuntimeError): |
| 23 | """The initialized engine did not produce a valid snapshot oracle.""" |
| 24 | |
| 25 | |
| 26 | class SnapshotBarrierError(RuntimeError): |
| 27 | """The controller supplied an invalid snapshot release marker.""" |
| 28 | |
| 29 | |
| 30 | @dataclass(frozen=True) |
| 31 | class ListenerConfig: |
| 32 | host: str | None |
| 33 | port: int |
| 34 | |
| 35 | |
| 36 | @dataclass(frozen=True) |
| 37 | class ControlArgs: |
| 38 | ready_file: Path |
| 39 | release_file: Path |
| 40 | release_timeout_s: float |
| 41 | |
| 42 | |
| 43 | def write_ready_atomic(path: Path, oracle: Oracle) -> None: |
| 44 | _write_json_atomic( |
| 45 | path, |
| 46 | { |
| 47 | "sampled_token_logprob": oracle.sampled_token_logprob, |
| 48 | "token_ids": oracle.token_ids, |
| 49 | "text": oracle.text, |
| 50 | }, |
| 51 | ) |
| 52 | |
| 53 | |
| 54 | def read_release_marker(path: Path) -> ListenerConfig: |
| 55 | path = Path(path) |
| 56 | path_stat = path.lstat() |
| 57 | if stat.S_ISLNK(path_stat.st_mode): |
| 58 | raise SnapshotBarrierError("release marker must not be a symlink") |
| 59 | try: |
| 60 | payload = json.loads(path.read_text()) |
| 61 | except (OSError, json.JSONDecodeError) as error: |
| 62 | raise SnapshotBarrierError("release marker is not valid JSON") from error |
| 63 | try: |
| 64 | marker = ReleaseMarker.model_validate(payload) |
| 65 | except ValidationError as error: |
| 66 | raise SnapshotBarrierError( |
| 67 | f"release marker is invalid: {_validation_path(error)}" |
| 68 | ) from error |
| 69 | return ListenerConfig(host=marker.host, port=marker.port) |
| 70 | |
| 71 | |
| 72 | def parse_control_args(argv: list[str]) -> tuple[ControlArgs, list[str]]: |
| 73 | parser = argparse.ArgumentParser(add_help=False) |
| 74 | parser.add_argument("--ready-file", type=Path, required=True) |
| 75 | parser.add_argument("--release-file", type=Path, required=True) |
| 76 | parser.add_argument("--release-timeout-s", type=float, default=900.0) |
| 77 | control_args, remaining = parser.parse_known_args(argv) |
| 78 | if remaining and remaining[0] == "--": |
| 79 | remaining = remaining[1:] |
| 80 | if control_args.release_timeout_s <= 0: |
| 81 | raise ValueError("release timeout must be positive") |
| 82 | return ( |
| 83 | ControlArgs( |
| 84 | ready_file=control_args.ready_file, |
| 85 | release_file=control_args.release_file, |
| 86 | release_timeout_s=control_args.release_timeout_s, |
| 87 | ), |
| 88 | remaining, |
| 89 | ) |
| 90 | |
| 91 | |
| 92 | async def _release_reloadable_state(engine: Any) -> None: |
| 93 | """Discard model and KV state before the process image is captured.""" |
| 94 | await engine.sleep(level=2) |
| 95 | |
| 96 | |
| 97 | async def _restore_reloadable_state(engine: Any) -> None: |
| 98 | """Rebuild state discarded by ``_release_reloadable_state``.""" |
| 99 | await engine.wake_up(tags=["weights"]) |
| 100 | await engine.collective_rpc("reload_weights") |
| 101 | await engine.wake_up(tags=["kv_cache"]) |
| 102 | |
| 103 | |
| 104 | def oracle_from_request_output(request_output: Any) -> Oracle: |
| 105 | try: |
| 106 | candidate = request_output.outputs[0] |
| 107 | token_ids = tuple(candidate.token_ids) |
| 108 | (sampled_token_id,) = token_ids |
| 109 | return Oracle( |
| 110 | token_ids=token_ids, |
| 111 | text=candidate.text, |
| 112 | sampled_token_logprob=candidate.logprobs[0][sampled_token_id].logprob, |
| 113 | ) |
| 114 | except ValueError as error: |
| 115 | raise SnapshotCanaryError( |
| 116 | "snapshot canary must produce exactly one finite token logprob" |
| 117 | ) from error |
| 118 | except (AttributeError, IndexError, KeyError, TypeError) as error: |
| 119 | raise SnapshotCanaryError( |
| 120 | "snapshot canary did not return sampled token logprob" |
| 121 | ) from error |
| 122 | |
| 123 | |
| 124 | async def run_engine_canary(engine: Any) -> Oracle: |
| 125 | from vllm import SamplingParams |
| 126 | |
| 127 | final_output = None |
| 128 | sampling_params = SamplingParams( |
| 129 | temperature=0, |
| 130 | min_tokens=1, |
| 131 | max_tokens=1, |
| 132 | seed=0, |
| 133 | logprobs=0, |
| 134 | ) |
| 135 | async for output in engine.generate( |
| 136 | _CANARY_PROMPT, |
| 137 | sampling_params, |
| 138 | request_id="vllm-snapshot-canary", |
| 139 | ): |
| 140 | final_output = output |
| 141 | if final_output is None: |
| 142 | raise SnapshotCanaryError("canary generation returned no output") |
| 143 | return oracle_from_request_output(final_output) |
| 144 | |
| 145 | |
| 146 | def detach_snapshot_streams() -> None: |
| 147 | """Remove launch-log descriptors from the reusable process image.""" |
| 148 | sys.stdout.flush() |
| 149 | sys.stderr.flush() |
| 150 | sink = os.open(os.devnull, os.O_WRONLY) |
| 151 | try: |
| 152 | os.dup2(sink, sys.stdout.fileno()) |
| 153 | os.dup2(sink, sys.stderr.fileno()) |
| 154 | finally: |
| 155 | os.close(sink) |
| 156 | |
| 157 | |
| 158 | async def wait_for_release_marker( |
| 159 | path: Path, |
| 160 | *, |
| 161 | timeout_s: float, |
| 162 | poll_interval_s: float = 0.05, |
| 163 | ) -> ListenerConfig: |
| 164 | if timeout_s <= 0: |
| 165 | raise ValueError("release timeout must be positive") |
| 166 | if poll_interval_s <= 0: |
| 167 | raise ValueError("release poll interval must be positive") |
| 168 | remaining_s = timeout_s |
| 169 | while not path.exists(): |
| 170 | if remaining_s <= 0: |
| 171 | raise TimeoutError(f"release marker not found before timeout: {path}") |
| 172 | delay_s = min(poll_interval_s, remaining_s) |
| 173 | await asyncio.sleep(delay_s) |
| 174 | remaining_s -= delay_s |
| 175 | return read_release_marker(path) |
| 176 | |
| 177 | |
| 178 | def parse_vllm_args(argv: list[str]) -> Any: |
| 179 | from vllm.entrypoints.launchers.cli_args import ( |
| 180 | make_arg_parser, |
| 181 | validate_parsed_serve_args, |
| 182 | ) |
| 183 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 184 | |
| 185 | parser = FlexibleArgumentParser(prog="python -m vllm.snapshot.server") |
| 186 | args = make_arg_parser(parser).parse_args(argv) |
| 187 | if args.model_tag is not None: |
| 188 | args.model = args.model_tag |
| 189 | if args.grpc or args.headless: |
| 190 | raise ValueError("snapshot server requires the HTTP frontend") |
| 191 | if args.api_server_count not in (None, 1): |
| 192 | raise ValueError("snapshot server supports one API server") |
| 193 | args.api_server_count = None |
| 194 | validate_parsed_serve_args(args) |
| 195 | return args |
| 196 | |
| 197 | |
| 198 | async def run_vllm_snapshot_child(control: ControlArgs, args: Any) -> None: |
| 199 | from vllm.entrypoints.launchers.api_server.entry import ( |
| 200 | build_and_serve, |
| 201 | build_async_engine_client, |
| 202 | ) |
| 203 | from vllm.entrypoints.launchers.launcher import setup_server |
| 204 | |
| 205 | args.enable_sleep_mode = True |
| 206 | |
| 207 | async with build_async_engine_client(args) as engine: |
| 208 | oracle = await run_engine_canary(engine) |
| 209 | try: |
| 210 | await _release_reloadable_state(engine) |
| 211 | await _restore_reloadable_state(engine) |
| 212 | rehearsal_oracle = await run_engine_canary(engine) |
| 213 | if not oracles_match(oracle, rehearsal_oracle): |
| 214 | raise SnapshotCanaryError("snapshot rehearsal changed canary output") |
| 215 | await _release_reloadable_state(engine) |
| 216 | except SnapshotCanaryError: |
| 217 | raise |
| 218 | except Exception as error: |
| 219 | raise SnapshotCanaryError("snapshot rehearsal failed") from error |
| 220 | |
| 221 | detach_snapshot_streams() |
| 222 | write_ready_atomic(control.ready_file, oracle) |
| 223 | listener = await wait_for_release_marker( |
| 224 | control.release_file, |
| 225 | timeout_s=control.release_timeout_s, |
| 226 | ) |
| 227 | args.host = listener.host |
| 228 | args.port = listener.port |
| 229 | await _restore_reloadable_state(engine) |
| 230 | listen_address, sock = setup_server(args, reuse_port=False) |
| 231 | try: |
| 232 | shutdown_task = await build_and_serve( |
| 233 | engine, |
| 234 | listen_address, |
| 235 | sock, |
| 236 | args, |
| 237 | ) |
| 238 | await shutdown_task |
| 239 | finally: |
| 240 | sock.close() |
| 241 | |
| 242 | |
| 243 | def main(argv: list[str] | None = None) -> None: |
| 244 | control, remaining = parse_control_args(sys.argv[1:] if argv is None else argv) |
| 245 | args = parse_vllm_args(remaining) |
| 246 | import uvloop |
| 247 | |
| 248 | uvloop.run(run_vllm_snapshot_child(control, args)) |
| 249 | |
| 250 | |
| 251 | if __name__ == "__main__": |
| 252 | main() |