| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import json |
| 5 | import math |
| 6 | import os |
| 7 | import stat |
| 8 | from pathlib import Path |
| 9 | from typing import Annotated, Any, Literal |
| 10 | |
| 11 | from pydantic import ( |
| 12 | BaseModel, |
| 13 | ConfigDict, |
| 14 | Field, |
| 15 | StrictFloat, |
| 16 | StrictInt, |
| 17 | StrictStr, |
| 18 | ValidationError, |
| 19 | field_validator, |
| 20 | model_validator, |
| 21 | ) |
| 22 | from typing_extensions import Self |
| 23 | |
| 24 | |
| 25 | class SnapshotCompatibilityError(RuntimeError): |
| 26 | """The snapshot identity does not match the requested runtime.""" |
| 27 | |
| 28 | |
| 29 | class SnapshotSecurityError(RuntimeError): |
| 30 | """The snapshot artifact does not meet the private-path contract.""" |
| 31 | |
| 32 | |
| 33 | _MAX_PID = 2**31 - 1 |
| 34 | _PositivePid = Annotated[int, Field(gt=0, le=_MAX_PID)] |
| 35 | _OracleTokenId = Annotated[int, Field(ge=0)] |
| 36 | |
| 37 | |
| 38 | class SnapshotRuntimeIdentity(BaseModel): |
| 39 | model_config = ConfigDict(extra="forbid", frozen=True, strict=True) |
| 40 | |
| 41 | vllm_version: str |
| 42 | python_version: str |
| 43 | torch_version: str |
| 44 | cuda_runtime: str |
| 45 | driver_version: str |
| 46 | criu_version: str |
| 47 | cuda_checkpoint_sha256: str |
| 48 | kernel_release: str |
| 49 | host_id: str |
| 50 | gpu_name: str |
| 51 | gpu_uuid: str |
| 52 | environment: tuple[tuple[str, str], ...] |
| 53 | |
| 54 | |
| 55 | class SnapshotManifest(SnapshotRuntimeIdentity): |
| 56 | schema_version: Literal[1] |
| 57 | boundary: Literal["post-engine-init-reloadable-state-released"] |
| 58 | created_at: str |
| 59 | artifact_bytes: Annotated[int, Field(ge=0)] |
| 60 | model: str |
| 61 | served_model_name: str |
| 62 | model_revision: str |
| 63 | tokenizer_revision: str |
| 64 | engine_argv: tuple[str, ...] |
| 65 | process_tree: Annotated[tuple[_PositivePid, ...], Field(min_length=1)] |
| 66 | cuda_holders: Annotated[tuple[_PositivePid, ...], Field(min_length=1)] |
| 67 | external_cache_files: tuple[tuple[str, str], ...] = () |
| 68 | oracle_token_ids: Annotated[ |
| 69 | tuple[_OracleTokenId, ...], Field(min_length=1, max_length=1) |
| 70 | ] |
| 71 | oracle_text: str |
| 72 | oracle_sampled_token_logprob: float |
| 73 | |
| 74 | @field_validator("schema_version", mode="before") |
| 75 | @classmethod |
| 76 | def _require_schema_version_one(cls, value: Any) -> Any: |
| 77 | if type(value) is not int or value != 1: |
| 78 | raise ValueError("schema version must be integer 1") |
| 79 | return value |
| 80 | |
| 81 | @field_validator("oracle_sampled_token_logprob", mode="before") |
| 82 | @classmethod |
| 83 | def _require_finite_float(cls, value: Any) -> Any: |
| 84 | if type(value) is not float or not math.isfinite(value): |
| 85 | raise ValueError("sampled token log probability must be a finite float") |
| 86 | return value |
| 87 | |
| 88 | @model_validator(mode="after") |
| 89 | def _validate_process_relationships(self) -> Self: |
| 90 | if len(set(self.process_tree)) != len(self.process_tree): |
| 91 | raise ValueError("process_tree contains duplicate PIDs") |
| 92 | if len(set(self.cuda_holders)) != len(self.cuda_holders): |
| 93 | raise ValueError("cuda_holders contains duplicate PIDs") |
| 94 | if not set(self.cuda_holders).issubset(self.process_tree): |
| 95 | raise ValueError("cuda_holders must be a subset of process_tree") |
| 96 | return self |
| 97 | |
| 98 | |
| 99 | class ReleaseMarker(BaseModel): |
| 100 | """Release marker the controller writes; unknown keys are ignored.""" |
| 101 | |
| 102 | release: Literal[True] |
| 103 | host: StrictStr | None = None |
| 104 | port: Annotated[StrictInt, Field(ge=1, le=65535)] |
| 105 | |
| 106 | @field_validator("release", mode="before") |
| 107 | @classmethod |
| 108 | def _require_json_true(cls, value: Any) -> Any: |
| 109 | # A literal accepts 1 as True; the barrier requires JSON true itself. |
| 110 | if value is not True: |
| 111 | raise ValueError("release must be JSON true") |
| 112 | return value |
| 113 | |
| 114 | |
| 115 | class ReadyMarker(BaseModel): |
| 116 | """Ready marker the child writes; unknown keys are ignored.""" |
| 117 | |
| 118 | token_ids: tuple[StrictInt, ...] |
| 119 | text: StrictStr |
| 120 | sampled_token_logprob: StrictFloat | StrictInt |
| 121 | |
| 122 | |
| 123 | def validate_identity( |
| 124 | expected: SnapshotRuntimeIdentity, actual: SnapshotRuntimeIdentity |
| 125 | ) -> None: |
| 126 | """Require an exact match for every field that can affect compatibility.""" |
| 127 | # ``expected`` may be a manifest, so compare the identity fields only. |
| 128 | fields = set(SnapshotRuntimeIdentity.model_fields) |
| 129 | expected_values = expected.model_dump(include=fields) |
| 130 | actual_values = actual.model_dump(include=fields) |
| 131 | differing = [ |
| 132 | name |
| 133 | for name in SnapshotRuntimeIdentity.model_fields |
| 134 | if expected_values[name] != actual_values[name] |
| 135 | ] |
| 136 | if differing: |
| 137 | raise SnapshotCompatibilityError(f"snapshot mismatch: {', '.join(differing)}") |
| 138 | |
| 139 | |
| 140 | def _existing_path_chain(path: Path): |
| 141 | current = path |
| 142 | while True: |
| 143 | if current.exists() or current.is_symlink(): |
| 144 | yield current |
| 145 | if current.parent == current: |
| 146 | break |
| 147 | current = current.parent |
| 148 | |
| 149 | |
| 150 | def _is_trusted_sticky_directory(path_stat: os.stat_result) -> bool: |
| 151 | return ( |
| 152 | path_stat.st_uid == 0 |
| 153 | and stat.S_ISDIR(path_stat.st_mode) |
| 154 | and bool(path_stat.st_mode & stat.S_ISVTX) |
| 155 | ) |
| 156 | |
| 157 | |
| 158 | def validate_artifact_root(path: Path, *, creating: bool) -> None: |
| 159 | """Validate that an artifact path cannot be replaced by another user.""" |
| 160 | path = Path(os.path.abspath(path)) |
| 161 | if path.is_symlink(): |
| 162 | raise SnapshotSecurityError(f"snapshot path is a symlink: {path}") |
| 163 | if not creating and not path.exists(): |
| 164 | raise SnapshotSecurityError(f"snapshot path does not exist: {path}") |
| 165 | |
| 166 | effective_uid = os.geteuid() |
| 167 | for component in _existing_path_chain(path): |
| 168 | component_stat = component.lstat() |
| 169 | if stat.S_ISLNK(component_stat.st_mode): |
| 170 | raise SnapshotSecurityError( |
| 171 | f"snapshot path contains a symlink: {component}" |
| 172 | ) |
| 173 | if component_stat.st_uid not in (0, effective_uid): |
| 174 | raise SnapshotSecurityError( |
| 175 | f"snapshot path has an untrusted owner: {component}" |
| 176 | ) |
| 177 | if component_stat.st_mode & ( |
| 178 | stat.S_IWGRP | stat.S_IWOTH |
| 179 | ) and not _is_trusted_sticky_directory(component_stat): |
| 180 | raise SnapshotSecurityError( |
| 181 | f"snapshot path has a group- or world-writable ancestor: {component}" |
| 182 | ) |
| 183 | |
| 184 | if path.exists(): |
| 185 | path_stat = path.lstat() |
| 186 | if not stat.S_ISDIR(path_stat.st_mode): |
| 187 | raise SnapshotSecurityError(f"snapshot path is not a directory: {path}") |
| 188 | if path_stat.st_uid != effective_uid: |
| 189 | raise SnapshotSecurityError( |
| 190 | f"snapshot directory is not owned by the current user: {path}" |
| 191 | ) |
| 192 | if stat.S_IMODE(path_stat.st_mode) & 0o077: |
| 193 | raise SnapshotSecurityError( |
| 194 | f"snapshot directory must have mode 0700: {path}" |
| 195 | ) |
| 196 | |
| 197 | |
| 198 | def _validation_path(error: ValidationError) -> str: |
| 199 | detail = error.errors(include_url=False)[0] |
| 200 | location = detail["loc"] |
| 201 | if not location: |
| 202 | return "JSON" if detail["type"] == "json_invalid" else "root" |
| 203 | return ".".join(map(str, location)) |
| 204 | |
| 205 | |
| 206 | def _fsync_directory(path: Path) -> None: |
| 207 | directory_fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)) |
| 208 | try: |
| 209 | os.fsync(directory_fd) |
| 210 | finally: |
| 211 | os.close(directory_fd) |
| 212 | |
| 213 | |
| 214 | def _write_json_atomic(path: Path, payload: object, *, overwrite: bool = False) -> None: |
| 215 | temporary = path.with_name(f"{path.name}.tmp") |
| 216 | if overwrite: |
| 217 | temporary.unlink(missing_ok=True) |
| 218 | elif path.exists() or path.is_symlink(): |
| 219 | raise FileExistsError(path) |
| 220 | flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0) |
| 221 | descriptor = os.open(temporary, flags, 0o600) |
| 222 | try: |
| 223 | with os.fdopen(descriptor, "wb") as output: |
| 224 | output.write( |
| 225 | json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() |
| 226 | ) |
| 227 | output.flush() |
| 228 | os.fsync(output.fileno()) |
| 229 | os.replace(temporary, path) |
| 230 | _fsync_directory(path.parent) |
| 231 | except BaseException: |
| 232 | temporary.unlink(missing_ok=True) |
| 233 | raise |
| 234 | |
| 235 | |
| 236 | def write_manifest_atomic(path: Path, manifest: SnapshotManifest) -> None: |
| 237 | """Write a new private manifest and make its directory entry durable.""" |
| 238 | path = Path(path) |
| 239 | validate_artifact_root(path, creating=False) |
| 240 | destination = path / "manifest.json" |
| 241 | try: |
| 242 | _write_json_atomic(destination, manifest.model_dump(mode="json")) |
| 243 | except FileExistsError as error: |
| 244 | raise SnapshotSecurityError( |
| 245 | f"manifest already exists: {destination}" |
| 246 | ) from error |
| 247 | |
| 248 | |
| 249 | def read_manifest(path: Path) -> SnapshotManifest: |
| 250 | """Read a snapshot published by atomically installing its manifest.""" |
| 251 | path = Path(path) |
| 252 | validate_artifact_root(path, creating=False) |
| 253 | manifest_path = path / "manifest.json" |
| 254 | if manifest_path.is_symlink(): |
| 255 | raise SnapshotSecurityError(f"manifest is a symlink: {manifest_path}") |
| 256 | flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) |
| 257 | try: |
| 258 | file_descriptor = os.open(manifest_path, flags) |
| 259 | except FileNotFoundError as error: |
| 260 | raise SnapshotSecurityError( |
| 261 | f"manifest does not exist: {manifest_path}" |
| 262 | ) from error |
| 263 | try: |
| 264 | with os.fdopen(file_descriptor, encoding="utf-8") as manifest_file: |
| 265 | payload = manifest_file.read() |
| 266 | except UnicodeDecodeError as error: |
| 267 | raise SnapshotCompatibilityError("invalid snapshot manifest: JSON") from error |
| 268 | try: |
| 269 | return SnapshotManifest.model_validate_json(payload, strict=True) |
| 270 | except ValidationError as error: |
| 271 | raise SnapshotCompatibilityError( |
| 272 | f"invalid snapshot manifest: {_validation_path(error)}" |
| 273 | ) from error |
| 274 | |
| 275 | |
| 276 | def inspect_snapshot(path: Path) -> dict[str, object]: |
| 277 | return read_manifest(path).model_dump(mode="json") |