| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Create and restore transaction orchestration for local snapshots.""" |
| 4 | |
| 5 | import argparse |
| 6 | import shutil |
| 7 | import sys |
| 8 | from pathlib import Path |
| 9 | |
| 10 | from vllm.snapshot.manifest import ( |
| 11 | read_manifest, |
| 12 | validate_artifact_root, |
| 13 | validate_identity, |
| 14 | ) |
| 15 | from vllm.snapshot.runtime import ( |
| 16 | LocalSnapshotTools, |
| 17 | SnapshotCreateError, |
| 18 | SnapshotRestoreError, |
| 19 | _error_detail, |
| 20 | ) |
| 21 | from vllm.snapshot.types import Oracle, oracles_match |
| 22 | |
| 23 | |
| 24 | def _child_engine_argv( |
| 25 | args: argparse.Namespace, engine_argv: tuple[str, ...] | None |
| 26 | ) -> tuple[str, ...]: |
| 27 | """Build the child argv from the first argv source that is present. |
| 28 | |
| 29 | In order: an explicit ``engine_argv`` from the caller, the process argv |
| 30 | tail when invoked as ``vllm snapshot create ...``, then the bare model from |
| 31 | the parsed args. Whichever source wins, snapshot options are stripped and |
| 32 | sleep mode is appended. |
| 33 | """ |
| 34 | if not engine_argv: |
| 35 | process_argv = tuple(sys.argv[1:]) |
| 36 | if process_argv[:2] == ("snapshot", "create"): |
| 37 | engine_argv = process_argv[2:] |
| 38 | else: |
| 39 | model = getattr(args, "model_tag", None) or getattr(args, "model", None) |
| 40 | if not model: |
| 41 | raise SnapshotCreateError("snapshot create requires a model argument") |
| 42 | engine_argv = (str(model),) |
| 43 | child_argv: list[str] = [] |
| 44 | iterator = iter(engine_argv) |
| 45 | for item in iterator: |
| 46 | option, separator, _ = item.partition("=") |
| 47 | option = option.replace("_", "-") |
| 48 | if len(option) > 2 and "--snapshot-dir".startswith(option): |
| 49 | if not separator: |
| 50 | next(iterator, None) |
| 51 | else: |
| 52 | child_argv.append(item) |
| 53 | if "--enable-sleep-mode" not in child_argv: |
| 54 | child_argv.append("--enable-sleep-mode") |
| 55 | return tuple(child_argv) |
| 56 | |
| 57 | |
| 58 | def create_snapshot( |
| 59 | args: argparse.Namespace, |
| 60 | *, |
| 61 | engine_argv: tuple[str, ...] | None = None, |
| 62 | tools: LocalSnapshotTools | None = None, |
| 63 | ) -> None: |
| 64 | target = Path(args.snapshot_dir).absolute() |
| 65 | toolset = tools or LocalSnapshotTools() |
| 66 | toolset.preflight("create", target) |
| 67 | if target.exists() or target.is_symlink(): |
| 68 | raise SnapshotCreateError(f"snapshot target already exists: {target}") |
| 69 | validate_artifact_root(target, creating=True) |
| 70 | target.parent.mkdir(mode=0o700, parents=False, exist_ok=True) |
| 71 | validate_artifact_root(target.parent, creating=False) |
| 72 | target.mkdir(mode=0o700) |
| 73 | validate_artifact_root(target, creating=False) |
| 74 | published = False |
| 75 | root_pid: int | None = None |
| 76 | try: |
| 77 | child_argv = _child_engine_argv(args, engine_argv) |
| 78 | root_pid = toolset.launch_child(target, child_argv) |
| 79 | oracle = toolset.wait_ready(target, root_pid) |
| 80 | inventory = toolset.inventory(root_pid, target) |
| 81 | toolset.dump(target, inventory) |
| 82 | toolset.verify_dead(inventory) |
| 83 | manifest = toolset.make_manifest(args, child_argv, inventory, oracle, target) |
| 84 | toolset.publish(target, manifest) |
| 85 | published = True |
| 86 | except BaseException: |
| 87 | if root_pid is not None: |
| 88 | toolset.abort_create(root_pid) |
| 89 | raise |
| 90 | finally: |
| 91 | # Runs even when abort_create raised, so no partial artifact is left. |
| 92 | if not published and target.exists(): |
| 93 | shutil.rmtree(target) |
| 94 | |
| 95 | |
| 96 | def restore_snapshot( |
| 97 | args: argparse.Namespace, *, tools: LocalSnapshotTools | None = None |
| 98 | ) -> None: |
| 99 | artifact = Path(args.snapshot_dir).absolute() |
| 100 | toolset = tools or LocalSnapshotTools() |
| 101 | toolset.preflight("restore", artifact) |
| 102 | manifest = read_manifest(artifact) |
| 103 | validate_identity(manifest, toolset.current_identity(manifest.gpu_uuid)) |
| 104 | |
| 105 | root_pid: int | None = None |
| 106 | try: |
| 107 | root_pid = toolset.restore(artifact, manifest) |
| 108 | toolset.release(artifact, args.host, args.port) |
| 109 | toolset.wait_listener(root_pid, args.host, args.port) |
| 110 | actual = toolset.request_oracle(args.host, args.port, manifest) |
| 111 | expected = Oracle( |
| 112 | token_ids=manifest.oracle_token_ids, |
| 113 | text=manifest.oracle_text, |
| 114 | sampled_token_logprob=manifest.oracle_sampled_token_logprob, |
| 115 | ) |
| 116 | if not oracles_match(expected, actual): |
| 117 | raise SnapshotRestoreError( |
| 118 | f"snapshot oracle mismatch: expected={expected!r}, actual={actual!r}" |
| 119 | ) |
| 120 | toolset.complete_restore(root_pid) |
| 121 | except BaseException as error: |
| 122 | if root_pid is not None: |
| 123 | try: |
| 124 | toolset.cleanup(root_pid) |
| 125 | except BaseException as cleanup_error: |
| 126 | raise SnapshotRestoreError( |
| 127 | f"snapshot restore failed: {_error_detail(error)}; " |
| 128 | f"cleanup failed: {_error_detail(cleanup_error)}" |
| 129 | ) from error |
| 130 | if isinstance(error, SnapshotRestoreError): |
| 131 | raise |
| 132 | raise SnapshotRestoreError(str(error)) from error |