| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import argparse |
| 5 | import json |
| 6 | import os |
| 7 | import platform |
| 8 | import typing |
| 9 | from pathlib import Path |
| 10 | |
| 11 | from vllm.entrypoints.cli.types import CLISubcommand |
| 12 | |
| 13 | if typing.TYPE_CHECKING: |
| 14 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 15 | else: |
| 16 | FlexibleArgumentParser = argparse.ArgumentParser |
| 17 | |
| 18 | |
| 19 | def validate_create_args(args: argparse.Namespace) -> None: |
| 20 | model = str(getattr(args, "model_tag", None) or args.model) |
| 21 | revision = str(getattr(args, "revision", None) or "") |
| 22 | if ( |
| 23 | Path(model).exists() |
| 24 | or len(revision) != 40 |
| 25 | or any(character not in "0123456789abcdefABCDEF" for character in revision) |
| 26 | ): |
| 27 | raise ValueError( |
| 28 | "snapshot create requires an immutable remote model and " |
| 29 | "40-character --revision" |
| 30 | ) |
| 31 | parallel_sizes = ( |
| 32 | ("tensor parallel", args.tensor_parallel_size), |
| 33 | ("pipeline parallel", args.pipeline_parallel_size), |
| 34 | ("data parallel", args.data_parallel_size), |
| 35 | ("prefill context parallel", args.prefill_context_parallel_size), |
| 36 | ) |
| 37 | for name, size in parallel_sizes: |
| 38 | if size != 1: |
| 39 | raise ValueError(f"snapshot create currently requires {name} size 1") |
| 40 | unsupported_frontend = ( |
| 41 | args.api_server_count not in (None, 1) |
| 42 | or args.grpc |
| 43 | or args.headless |
| 44 | or args.api_key |
| 45 | or os.environ.get("VLLM_API_KEY") |
| 46 | or args.uds |
| 47 | or args.middleware |
| 48 | ) |
| 49 | if unsupported_frontend or any( |
| 50 | ( |
| 51 | args.ssl_keyfile, |
| 52 | args.ssl_certfile, |
| 53 | args.ssl_ca_certs, |
| 54 | args.enable_ssl_refresh, |
| 55 | args.ssl_ciphers, |
| 56 | ) |
| 57 | ): |
| 58 | raise ValueError( |
| 59 | "snapshot create requires one unauthenticated plaintext HTTP server" |
| 60 | ) |
| 61 | if args.logprobs_mode in {"raw_logits", "processed_logits"}: |
| 62 | raise ValueError("snapshot canary requires a log-probability mode") |
| 63 | if any( |
| 64 | getattr(args, name, None) is not None |
| 65 | for name in ("speculative_config", "spec_method", "spec_model", "spec_tokens") |
| 66 | ): |
| 67 | raise ValueError("snapshot create does not support speculative decoding") |
| 68 | if platform.system() != "Linux" or platform.machine() != "x86_64": |
| 69 | raise ValueError("snapshot create requires Linux x86_64") |
| 70 | |
| 71 | |
| 72 | def run_create(args: argparse.Namespace) -> None: |
| 73 | validate_create_args(args) |
| 74 | from vllm.snapshot.controller import create_snapshot |
| 75 | |
| 76 | create_snapshot(args) |
| 77 | |
| 78 | |
| 79 | def run_inspect(args: argparse.Namespace) -> None: |
| 80 | from vllm.snapshot.manifest import inspect_snapshot |
| 81 | |
| 82 | print(json.dumps(inspect_snapshot(Path(args.snapshot_dir)), indent=2)) |
| 83 | |
| 84 | |
| 85 | def run_restore(args: argparse.Namespace) -> None: |
| 86 | from vllm.snapshot.controller import restore_snapshot |
| 87 | |
| 88 | restore_snapshot(args) |
| 89 | |
| 90 | |
| 91 | class SnapshotSubcommand(CLISubcommand): |
| 92 | """The `snapshot` subcommand for the vLLM CLI.""" |
| 93 | |
| 94 | name = "snapshot" |
| 95 | |
| 96 | def __init__(self, *, create_requested: bool = False) -> None: |
| 97 | self.create_requested = create_requested |
| 98 | |
| 99 | @staticmethod |
| 100 | def cmd(args: argparse.Namespace) -> None: |
| 101 | args.snapshot_dispatch(args) |
| 102 | |
| 103 | def subparser_init( |
| 104 | self, subparsers: argparse._SubParsersAction |
| 105 | ) -> FlexibleArgumentParser: |
| 106 | parser = subparsers.add_parser( |
| 107 | self.name, |
| 108 | help="Create, inspect, or restore an initialized vLLM snapshot.", |
| 109 | usage="vllm snapshot <create|inspect|restore> [options]", |
| 110 | ) |
| 111 | actions = parser.add_subparsers(required=True, dest="snapshot_action") |
| 112 | |
| 113 | create_parser = actions.add_parser( |
| 114 | "create", help="Create a snapshot from an initialized TP1 engine." |
| 115 | ) |
| 116 | if self.create_requested: |
| 117 | from vllm.entrypoints.launchers.cli_args import make_arg_parser |
| 118 | |
| 119 | create_parser = make_arg_parser(create_parser) |
| 120 | else: |
| 121 | create_parser.add_argument("model_tag", nargs="?") |
| 122 | create_parser.add_argument("--snapshot-dir", required=True) |
| 123 | create_parser.set_defaults(snapshot_dispatch=run_create) |
| 124 | |
| 125 | inspect_parser = actions.add_parser( |
| 126 | "inspect", help="Inspect a snapshot without restoring it." |
| 127 | ) |
| 128 | inspect_parser.add_argument("snapshot_dir") |
| 129 | inspect_parser.set_defaults(snapshot_dispatch=run_inspect) |
| 130 | |
| 131 | restore_parser = actions.add_parser( |
| 132 | "restore", help="Restore a same-host TP1 snapshot." |
| 133 | ) |
| 134 | restore_parser.add_argument("snapshot_dir") |
| 135 | restore_parser.add_argument("--host", default=None) |
| 136 | restore_parser.add_argument("--port", type=int, default=8000) |
| 137 | restore_parser.set_defaults(snapshot_dispatch=run_restore) |
| 138 | |
| 139 | return parser |
| 140 | |
| 141 | |
| 142 | def cmd_init(*, create_requested: bool = False) -> list[CLISubcommand]: |
| 143 | return [SnapshotSubcommand(create_requested=create_requested)] |