| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for the `vllm launch` CLI subcommand.""" |
| 4 | |
| 5 | import argparse |
| 6 | import hashlib |
| 7 | import json |
| 8 | import os |
| 9 | import signal |
| 10 | import subprocess |
| 11 | import sys |
| 12 | from contextlib import asynccontextmanager |
| 13 | from pathlib import Path |
| 14 | from typing import Any |
| 15 | from unittest.mock import create_autospec, patch |
| 16 | |
| 17 | import pytest |
| 18 | |
| 19 | import vllm.snapshot.runtime as snapshot_runtime |
| 20 | import vllm.snapshot.server as snapshot_server |
| 21 | from vllm.entrypoints.cli import snapshot as snapshot_cli |
| 22 | from vllm.entrypoints.cli.launch import ( |
| 23 | LaunchSubcommand, |
| 24 | RenderSubcommand, |
| 25 | cmd_init, |
| 26 | ) |
| 27 | from vllm.snapshot.controller import ( |
| 28 | LocalSnapshotTools, |
| 29 | SnapshotCreateError, |
| 30 | SnapshotRestoreError, |
| 31 | _child_engine_argv, |
| 32 | create_snapshot, |
| 33 | restore_snapshot, |
| 34 | ) |
| 35 | from vllm.snapshot.manifest import ( |
| 36 | SnapshotCompatibilityError, |
| 37 | SnapshotManifest, |
| 38 | SnapshotRuntimeIdentity, |
| 39 | SnapshotSecurityError, |
| 40 | read_manifest, |
| 41 | validate_artifact_root, |
| 42 | validate_identity, |
| 43 | write_manifest_atomic, |
| 44 | ) |
| 45 | from vllm.snapshot.runtime import ProcessInventory, _decode_endpoint, _TcpSocketRecord |
| 46 | from vllm.snapshot.server import SnapshotCanaryError |
| 47 | from vllm.snapshot.types import Oracle, oracles_match |
| 48 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 49 | |
| 50 | |
| 51 | @pytest.fixture |
| 52 | def launch_parser(): |
| 53 | parser = FlexibleArgumentParser(description="test") |
| 54 | subparsers = parser.add_subparsers(required=False, dest="subparser") |
| 55 | LaunchSubcommand().subparser_init(subparsers) |
| 56 | return parser |
| 57 | |
| 58 | |
| 59 | def test_subcommand_name(): |
| 60 | assert LaunchSubcommand().name == "launch" |
| 61 | |
| 62 | |
| 63 | def test_cmd_init_returns_subcommand(): |
| 64 | result = cmd_init() |
| 65 | assert len(result) == 1 |
| 66 | assert isinstance(result[0], LaunchSubcommand) |
| 67 | |
| 68 | |
| 69 | # -- Parsing: `vllm launch render` -- |
| 70 | |
| 71 | |
| 72 | def test_parse_launch_render(launch_parser): |
| 73 | args = launch_parser.parse_args(["launch", "render", "--model", "test-model"]) |
| 74 | assert args.launch_component == "render" |
| 75 | |
| 76 | |
| 77 | def test_parse_launch_requires_component(launch_parser): |
| 78 | with pytest.raises(SystemExit): |
| 79 | launch_parser.parse_args(["launch", "--model", "test-model"]) |
| 80 | |
| 81 | |
| 82 | def test_parse_launch_invalid_component(launch_parser): |
| 83 | with pytest.raises(SystemExit): |
| 84 | launch_parser.parse_args(["launch", "unknown", "--model", "test-model"]) |
| 85 | |
| 86 | |
| 87 | # -- Dispatch -- |
| 88 | |
| 89 | |
| 90 | def test_cmd_launch_render_calls_run(): |
| 91 | args = argparse.Namespace(model_tag=None, model="test-model") |
| 92 | with patch("vllm.entrypoints.cli.launch.uvloop.run") as mock_uvloop_run: |
| 93 | RenderSubcommand.cmd(args) |
| 94 | mock_uvloop_run.assert_called_once() |
| 95 | |
| 96 | |
| 97 | def test_cmd_launch_model_tag_overrides(): |
| 98 | args = argparse.Namespace( |
| 99 | model_tag="tag-model", |
| 100 | model="original-model", |
| 101 | launch_command=lambda a: None, |
| 102 | ) |
| 103 | LaunchSubcommand.cmd(args) |
| 104 | assert args.model == "tag-model" |
| 105 | |
| 106 | |
| 107 | def test_cmd_launch_model_tag_none(): |
| 108 | args = argparse.Namespace( |
| 109 | model_tag=None, |
| 110 | model="original-model", |
| 111 | launch_command=lambda a: None, |
| 112 | ) |
| 113 | LaunchSubcommand.cmd(args) |
| 114 | assert args.model == "original-model" |
| 115 | |
| 116 | |
| 117 | def test_cmd_dispatches(): |
| 118 | called = {} |
| 119 | |
| 120 | def fake_dispatch(args): |
| 121 | called["args"] = args |
| 122 | |
| 123 | args = argparse.Namespace(launch_command=fake_dispatch) |
| 124 | LaunchSubcommand.cmd(args) |
| 125 | assert "args" in called |
| 126 | |
| 127 | |
| 128 | # -- Module registration -- |
| 129 | |
| 130 | |
| 131 | def test_subparser_init_returns_parser(): |
| 132 | parser = FlexibleArgumentParser(description="test") |
| 133 | subparsers = parser.add_subparsers(required=False, dest="subparser") |
| 134 | result = LaunchSubcommand().subparser_init(subparsers) |
| 135 | assert isinstance(result, FlexibleArgumentParser) |
| 136 | |
| 137 | |
| 138 | def test_launch_registered_in_main(): |
| 139 | """Verify that launch module is importable as a CLI module.""" |
| 140 | import vllm.entrypoints.cli.launch as launch_module |
| 141 | |
| 142 | assert hasattr(launch_module, "cmd_init") |
| 143 | subcmds = launch_module.cmd_init() |
| 144 | assert any(s.name == "launch" for s in subcmds) |
| 145 | |
| 146 | |
| 147 | _MODEL_REVISION = "c1899de289a04d12100db370d81485cdf75e47ca" |
| 148 | |
| 149 | |
| 150 | def _run_cli(argv: list[str]) -> None: |
| 151 | from vllm.entrypoints.cli import main as cli_main |
| 152 | |
| 153 | with patch.object(sys, "argv", ["vllm", *argv]): |
| 154 | cli_main.main() |
| 155 | |
| 156 | |
| 157 | def test_snapshot_environment_contract( |
| 158 | tmp_path: Path, |
| 159 | monkeypatch: pytest.MonkeyPatch, |
| 160 | capsys: pytest.CaptureFixture[str], |
| 161 | ): |
| 162 | from vllm.entrypoints.serve.utils import api_utils |
| 163 | |
| 164 | monkeypatch.setattr(api_utils, "cli_env_setup", pytest.fail) |
| 165 | secret = "snapshot-secret" |
| 166 | monkeypatch.setenv("VLLM_API_KEY", secret) |
| 167 | monkeypatch.setenv("VLLM_USER_SETTING", secret) |
| 168 | monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", secret) |
| 169 | environment = dict(LocalSnapshotTools()._environment_identity()) |
| 170 | assert ( |
| 171 | environment["VLLM_USER_SETTING"] |
| 172 | == hashlib.sha256(b"VLLM_USER_SETTING\0snapshot-secret").hexdigest() |
| 173 | ) |
| 174 | assert ( |
| 175 | environment["VLLM_USER_SETTING"] != environment["VLLM_WORKER_MULTIPROC_METHOD"] |
| 176 | ) |
| 177 | assert secret not in str(environment) |
| 178 | assert "VLLM_API_KEY" not in environment |
| 179 | |
| 180 | monkeypatch.delenv("VLLM_WORKER_MULTIPROC_METHOD") |
| 181 | implicit_worker_default = dict(LocalSnapshotTools()._environment_identity()) |
| 182 | monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") |
| 183 | assert dict(LocalSnapshotTools()._environment_identity()) == implicit_worker_default |
| 184 | |
| 185 | monkeypatch.setenv("VLLM_USER_SETTING", f"{secret}-changed") |
| 186 | changed_environment = dict(LocalSnapshotTools()._environment_identity()) |
| 187 | assert changed_environment["VLLM_USER_SETTING"] != environment["VLLM_USER_SETTING"] |
| 188 | |
| 189 | artifact = tmp_path / "snapshot" |
| 190 | artifact.mkdir(mode=0o700) |
| 191 | write_manifest_atomic(artifact, _manifest(environment=tuple(environment.items()))) |
| 192 | before = os.environ.copy() |
| 193 | |
| 194 | _run_cli(["snapshot", "inspect", str(artifact)]) |
| 195 | |
| 196 | assert os.environ == before |
| 197 | output = capsys.readouterr().out |
| 198 | assert secret not in output |
| 199 | inspected = json.loads(output) |
| 200 | inspected_environment = dict(inspected["environment"]) |
| 201 | assert inspected_environment == environment |
| 202 | |
| 203 | tools = LocalSnapshotTools() |
| 204 | monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity()) |
| 205 | manifest = tools.make_manifest( |
| 206 | argparse.Namespace( |
| 207 | hf_token=secret, |
| 208 | model_tag="Qwen/Qwen3-0.6B", |
| 209 | revision=_MODEL_REVISION, |
| 210 | served_model_name=None, |
| 211 | tokenizer_revision=None, |
| 212 | ), |
| 213 | ( |
| 214 | "Qwen/Qwen3-0.6B", |
| 215 | "--api-key", |
| 216 | secret, |
| 217 | f"{secret}-alternate", |
| 218 | f"--api-key={secret}", |
| 219 | "--hf-token", |
| 220 | secret, |
| 221 | f"--hf-token={secret}", |
| 222 | ), |
| 223 | ProcessInventory(100, (100, 101), (101,), "GPU-abc"), |
| 224 | _oracle(), |
| 225 | tmp_path, |
| 226 | ) |
| 227 | |
| 228 | assert manifest.engine_argv == ( |
| 229 | "Qwen/Qwen3-0.6B", |
| 230 | "--api-key", |
| 231 | "***", |
| 232 | "***", |
| 233 | "--api-key=***", |
| 234 | "--hf-token", |
| 235 | "***", |
| 236 | "--hf-token=***", |
| 237 | ) |
| 238 | |
| 239 | |
| 240 | @pytest.mark.parametrize( |
| 241 | "case", |
| 242 | [ |
| 243 | (("--hf_token", "SECRET"), ("--hf_token", "***")), |
| 244 | (("--hf_token=SECRET",), ("--hf_token=***",)), |
| 245 | (("--api_key", "SECRET"), ("--api_key", "***")), |
| 246 | (("--api_key=SECRET",), ("--api_key=***",)), |
| 247 | (("--hf-tok", "SECRET"), ("--hf-tok", "***")), |
| 248 | (("--hf-tok=SECRET",), ("--hf-tok=***",)), |
| 249 | (("--hf_tok", "SECRET"), ("--hf_tok", "***")), |
| 250 | (("--api-k", "SECRET", "SECOND"), ("--api-k", "***", "***")), |
| 251 | (("--api-k=SECRET",), ("--api-k=***",)), |
| 252 | (("--hf-overrides", "{}"), ("--hf-overrides", "{}")), |
| 253 | ], |
| 254 | ) |
| 255 | def test_snapshot_manifest_redacts_accepted_secret_options(case): |
| 256 | engine_argv, expected = case |
| 257 | tools = create_autospec(LocalSnapshotTools, instance=True) |
| 258 | tools.current_identity.return_value = _runtime_identity() |
| 259 | tools._artifact_bytes.return_value = 0 |
| 260 | manifest = LocalSnapshotTools.make_manifest( |
| 261 | tools, |
| 262 | snapshot_server.parse_vllm_args(["model", *engine_argv]), |
| 263 | engine_argv, |
| 264 | ProcessInventory(100, (100, 101), (101,), "GPU-abc"), |
| 265 | _oracle(), |
| 266 | Path(), |
| 267 | ) |
| 268 | |
| 269 | assert manifest.engine_argv == expected |
| 270 | if "SECRET" in " ".join(engine_argv): |
| 271 | assert "***" in " ".join(manifest.engine_argv) |
| 272 | assert "SECRET" not in " ".join(manifest.engine_argv) |
| 273 | |
| 274 | |
| 275 | def parse_snapshot(*argv: str): |
| 276 | parser = FlexibleArgumentParser() |
| 277 | subparsers = parser.add_subparsers(dest="subparser", required=True) |
| 278 | snapshot_cli.SnapshotSubcommand( |
| 279 | create_requested=argv[:1] == ("create",) |
| 280 | ).subparser_init(subparsers) |
| 281 | return parser.parse_args(["snapshot", *argv]) |
| 282 | |
| 283 | |
| 284 | @pytest.mark.parametrize("option", ["--snapshot-dir", "--snapshot_dir", "--snapshot-d"]) |
| 285 | @pytest.mark.parametrize("joined", [False, True]) |
| 286 | def test_snapshot_create_forwards_only_engine_options(option, joined, monkeypatch): |
| 287 | directory = [f"{option}=/tmp/snapshot"] if joined else [option, "/tmp/snapshot"] |
| 288 | argv = ["create", "Qwen/Qwen3-0.6B", "--revision", _MODEL_REVISION, *directory] |
| 289 | args = parse_snapshot(*argv) |
| 290 | monkeypatch.setattr(sys, "argv", ["vllm", "snapshot", *argv]) |
| 291 | |
| 292 | child = snapshot_server.parse_vllm_args(list(_child_engine_argv(args, None))) |
| 293 | |
| 294 | assert child.model == "Qwen/Qwen3-0.6B" |
| 295 | assert child.revision == _MODEL_REVISION |
| 296 | assert child.enable_sleep_mode |
| 297 | |
| 298 | |
| 299 | @pytest.mark.parametrize("model_in_cli", [False, True]) |
| 300 | def test_snapshot_create_config_preserves_cli_precedence( |
| 301 | tmp_path, monkeypatch, model_in_cli |
| 302 | ): |
| 303 | config = tmp_path / "engine.yaml" |
| 304 | config.write_text("model: config-model\nmax-model-len: 512\ndtype: float16\n") |
| 305 | model = ["Qwen/Qwen3-0.6B"] if model_in_cli else [] |
| 306 | argv = [ |
| 307 | "create", |
| 308 | *model, |
| 309 | "--snapshot-dir", |
| 310 | str(tmp_path / "snapshot"), |
| 311 | "--config", |
| 312 | str(config), |
| 313 | "--revision", |
| 314 | _MODEL_REVISION, |
| 315 | "--max-model-len", |
| 316 | "1024", |
| 317 | ] |
| 318 | args = parse_snapshot(*argv) |
| 319 | monkeypatch.setattr(sys, "argv", ["vllm", "snapshot", *argv]) |
| 320 | child = snapshot_server.parse_vllm_args(list(_child_engine_argv(args, None))) |
| 321 | |
| 322 | assert ( |
| 323 | (args.model_tag or args.model) |
| 324 | == child.model |
| 325 | == ("Qwen/Qwen3-0.6B" if model_in_cli else "config-model") |
| 326 | ) |
| 327 | assert args.max_model_len == child.max_model_len == 1024 |
| 328 | assert args.dtype == child.dtype == "float16" |
| 329 | |
| 330 | |
| 331 | @pytest.mark.parametrize( |
| 332 | ("probe", "status", "output"), |
| 333 | [ |
| 334 | (None, 0, ""), |
| 335 | ("nvidia-smi", 1, ""), |
| 336 | ("find", 1, ""), |
| 337 | ("ss", 1, ""), |
| 338 | ("nvidia-smi", 0, "12345"), |
| 339 | ("find", 0, "link_remap.fixture"), |
| 340 | ("ss", 0, "LISTEN"), |
| 341 | ], |
| 342 | ) |
| 343 | def test_snapshot_e2e_cleanup_requires_successful_empty_probes(probe, status, output): |
| 344 | script = ( |
| 345 | Path(__file__).resolve().parents[3] |
| 346 | / ".buildkite/scripts/initialized-snapshot-e2e.sh" |
| 347 | ) |
| 348 | if not script.exists(): |
| 349 | pytest.skip("snapshot E2E script is not packaged in this test image") |
| 350 | source = script.read_text() |
| 351 | # Exercise the real probes without entering Docker/CRIU setup. |
| 352 | functions = source[source.index("link_remaps() {") : source.index("wait_clean() {")] |
| 353 | shell = ( |
| 354 | """set -Eeuo pipefail |
| 355 | GPU_UUID=fixture PORT_ONE=18001 PORT_TWO=18002 LINK_REMAP_BASELINE='' |
| 356 | SNAPSHOT_PIDS=(999999999) |
| 357 | nvidia-smi() { :; } |
| 358 | find() { :; } |
| 359 | ss() { :; } |
| 360 | """ |
| 361 | + functions |
| 362 | ) |
| 363 | if probe: |
| 364 | shell += f'\n{probe}() {{ printf %s "{output}"; return {status}; }}\n' |
| 365 | shell += "\nif state_clean; then exit 0; else exit 1; fi\n" |
| 366 | |
| 367 | result = subprocess.run(["bash", "-c", shell], capture_output=True, timeout=10) |
| 368 | |
| 369 | assert result.returncode == (0 if probe is None else 1), result.stderr |
| 370 | |
| 371 | |
| 372 | def test_snapshot_create_cli_accepts_only_pinned_compact_mode( |
| 373 | monkeypatch: pytest.MonkeyPatch, |
| 374 | ): |
| 375 | monkeypatch.setattr(snapshot_cli.platform, "system", lambda: "Linux") |
| 376 | monkeypatch.setattr(snapshot_cli.platform, "machine", lambda: "x86_64") |
| 377 | base = ("create", "Qwen/Qwen3-0.6B", "--snapshot-dir=/tmp/snapshot") |
| 378 | compact = parse_snapshot(*base, "--revision", _MODEL_REVISION) |
| 379 | |
| 380 | snapshot_cli.validate_create_args(compact) |
| 381 | assert compact.model_tag == "Qwen/Qwen3-0.6B" |
| 382 | |
| 383 | dispatched: list[argparse.Namespace] = [] |
| 384 | monkeypatch.setattr(snapshot_cli, "run_create", dispatched.append) |
| 385 | _run_cli(["snapshot", *base, "--revision", _MODEL_REVISION]) |
| 386 | assert dispatched[0].revision == _MODEL_REVISION |
| 387 | |
| 388 | |
| 389 | @pytest.mark.parametrize( |
| 390 | ("attribute", "value", "message"), |
| 391 | [ |
| 392 | ("prefill_context_parallel_size", 2, "prefill context parallel size 1"), |
| 393 | ("model_tag", None, "immutable remote model"), |
| 394 | ("logprobs_mode", "raw_logits", "log-probability mode"), |
| 395 | ("speculative_config", {}, "does not support speculative decoding"), |
| 396 | ("spec_method", "ngram", "does not support speculative decoding"), |
| 397 | ("spec_model", "draft-model", "does not support speculative decoding"), |
| 398 | ("spec_tokens", 3, "does not support speculative decoding"), |
| 399 | ], |
| 400 | ) |
| 401 | def test_snapshot_create_rejects_invalid_config( |
| 402 | attribute: str, |
| 403 | value: object, |
| 404 | message: str, |
| 405 | tmp_path: Path, |
| 406 | ): |
| 407 | args = parse_snapshot( |
| 408 | "create", |
| 409 | "Qwen/Qwen3-0.6B", |
| 410 | "--snapshot-dir=/tmp/snapshot", |
| 411 | "--revision", |
| 412 | _MODEL_REVISION, |
| 413 | ) |
| 414 | if attribute == "model_tag": |
| 415 | local_model = tmp_path / "model" |
| 416 | local_model.mkdir() |
| 417 | value = str(local_model) |
| 418 | setattr(args, attribute, value) |
| 419 | with pytest.raises(ValueError, match=message): |
| 420 | snapshot_cli.validate_create_args(args) |
| 421 | |
| 422 | |
| 423 | def _oracle( |
| 424 | token_ids: tuple[int, ...] = (12095,), |
| 425 | text: str = " Paris", |
| 426 | logprob: float = -0.125, |
| 427 | ) -> Oracle: |
| 428 | return Oracle(token_ids, text, logprob) |
| 429 | |
| 430 | |
| 431 | def _manifest(**changes: object) -> SnapshotManifest: |
| 432 | manifest = SnapshotManifest( |
| 433 | schema_version=1, |
| 434 | boundary="post-engine-init-reloadable-state-released", |
| 435 | created_at="2026-08-06T00:00:00Z", |
| 436 | artifact_bytes=8, |
| 437 | vllm_version="0.1.dev1+snapshot", |
| 438 | python_version="3.12.3", |
| 439 | torch_version="2.9.0", |
| 440 | cuda_runtime="12.9", |
| 441 | driver_version="575.57.08", |
| 442 | criu_version="4.1", |
| 443 | cuda_checkpoint_sha256="a" * 64, |
| 444 | kernel_release="6.8.0", |
| 445 | host_id="host-a", |
| 446 | gpu_name="NVIDIA A10", |
| 447 | gpu_uuid="GPU-abc", |
| 448 | model="Qwen/Qwen3-0.6B", |
| 449 | served_model_name="Qwen/Qwen3-0.6B", |
| 450 | model_revision="model-sha", |
| 451 | tokenizer_revision="tokenizer-sha", |
| 452 | engine_argv=("Qwen/Qwen3-0.6B",), |
| 453 | environment=(("VLLM_USE_V1", "1"),), |
| 454 | process_tree=(100, 101), |
| 455 | cuda_holders=(101,), |
| 456 | oracle_token_ids=(12095,), |
| 457 | oracle_text=" Paris", |
| 458 | oracle_sampled_token_logprob=-0.125, |
| 459 | ) |
| 460 | return SnapshotManifest.model_validate( |
| 461 | {**manifest.model_dump(mode="python"), **changes}, strict=True |
| 462 | ) |
| 463 | |
| 464 | |
| 465 | def _runtime_identity(**changes: object) -> Any: |
| 466 | fields = SnapshotRuntimeIdentity.model_fields.keys() |
| 467 | payload = _manifest().model_dump(mode="python", include=fields) |
| 468 | return SnapshotRuntimeIdentity.model_validate({**payload, **changes}, strict=True) |
| 469 | |
| 470 | |
| 471 | def _fake_snapshot_tools() -> Any: |
| 472 | tools = create_autospec(LocalSnapshotTools, instance=True) |
| 473 | tools.launch_child.return_value = 100 |
| 474 | tools.wait_ready.return_value = _oracle() |
| 475 | tools.inventory.return_value = ProcessInventory(100, (100, 101), (101,), "GPU-abc") |
| 476 | |
| 477 | tools.current_identity.side_effect = lambda _gpu_uuid: _runtime_identity() |
| 478 | tools.restore.return_value = 100 |
| 479 | tools.request_oracle.return_value = _oracle() |
| 480 | return tools |
| 481 | |
| 482 | |
| 483 | @pytest.mark.parametrize("value", ["0", "-1", "nan", "inf"]) |
| 484 | def test_snapshot_timeout_must_be_positive_and_finite( |
| 485 | value: str, monkeypatch: pytest.MonkeyPatch |
| 486 | ): |
| 487 | monkeypatch.setenv("VLLM_SNAPSHOT_TIMEOUT_S", value) |
| 488 | with pytest.raises(ValueError, match="positive finite number"): |
| 489 | LocalSnapshotTools() |
| 490 | |
| 491 | |
| 492 | def test_snapshot_create_abort_reaps_after_process_group_cleanup( |
| 493 | tmp_path: Path, monkeypatch: pytest.MonkeyPatch |
| 494 | ): |
| 495 | tools = LocalSnapshotTools() |
| 496 | process = subprocess.Popen( |
| 497 | [sys.executable, "-c", "raise SystemExit(1)"], |
| 498 | stdin=subprocess.DEVNULL, |
| 499 | stdout=subprocess.DEVNULL, |
| 500 | stderr=subprocess.DEVNULL, |
| 501 | start_new_session=True, |
| 502 | ) |
| 503 | tools._children[process.pid] = process |
| 504 | (tmp_path / "child.log").write_text("root failed") |
| 505 | killed: list[tuple[int, signal.Signals]] = [] |
| 506 | monkeypatch.setattr( |
| 507 | os, |
| 508 | "killpg", |
| 509 | lambda pid, sig: killed.append((pid, sig)), |
| 510 | ) |
| 511 | monkeypatch.setattr(os, "waitid", lambda *_args: object(), raising=False) |
| 512 | |
| 513 | try: |
| 514 | with pytest.raises(SnapshotCreateError, match="root failed"): |
| 515 | tools.wait_ready(tmp_path, process.pid) |
| 516 | assert process.returncode is None |
| 517 | tools.abort_create(process.pid) |
| 518 | finally: |
| 519 | if process.pid in tools._children: |
| 520 | tools.abort_create(process.pid) |
| 521 | |
| 522 | assert killed == [(process.pid, signal.SIGKILL)] |
| 523 | assert process.pid not in tools._children |
| 524 | |
| 525 | |
| 526 | def _supply_waitid(monkeypatch: pytest.MonkeyPatch) -> None: |
| 527 | """`wait_ready` binds Linux-only `os.waitid` before it reads the marker.""" |
| 528 | monkeypatch.setattr(os, "waitid", lambda *_args: None, raising=False) |
| 529 | |
| 530 | |
| 531 | def test_snapshot_wait_ready_reads_back_the_oracle_the_child_wrote( |
| 532 | tmp_path: Path, monkeypatch: pytest.MonkeyPatch |
| 533 | ): |
| 534 | _supply_waitid(monkeypatch) |
| 535 | snapshot_server.write_ready_atomic(tmp_path / "ready.json", _oracle()) |
| 536 | |
| 537 | assert LocalSnapshotTools().wait_ready(tmp_path, os.getpid()) == _oracle() |
| 538 | |
| 539 | |
| 540 | @pytest.mark.parametrize( |
| 541 | "payload", |
| 542 | [ |
| 543 | {"token_ids": [12095], "sampled_token_logprob": -0.125}, |
| 544 | {"token_ids": ["12095"], "text": " Paris", "sampled_token_logprob": -0.125}, |
| 545 | ], |
| 546 | ) |
| 547 | def test_snapshot_wait_ready_rejects_a_malformed_marker( |
| 548 | tmp_path: Path, monkeypatch: pytest.MonkeyPatch, payload: dict[str, object] |
| 549 | ): |
| 550 | _supply_waitid(monkeypatch) |
| 551 | (tmp_path / "ready.json").write_text(json.dumps(payload)) |
| 552 | |
| 553 | with pytest.raises(SnapshotCreateError, match="ready marker invalid"): |
| 554 | LocalSnapshotTools().wait_ready(tmp_path, os.getpid()) |
| 555 | |
| 556 | |
| 557 | def test_snapshot_create_abort_rejects_surviving_process_group( |
| 558 | monkeypatch: pytest.MonkeyPatch, |
| 559 | ): |
| 560 | tools = LocalSnapshotTools() |
| 561 | process = create_autospec(subprocess.Popen, instance=True) |
| 562 | process.wait.side_effect = subprocess.TimeoutExpired("snapshot child", 10) |
| 563 | tools._children[100] = process |
| 564 | monkeypatch.setattr(os, "killpg", lambda _pid, _sig: None) |
| 565 | |
| 566 | with pytest.raises(SnapshotCreateError, match="survived SIGKILL"): |
| 567 | tools.abort_create(100) |
| 568 | |
| 569 | |
| 570 | @pytest.mark.asyncio |
| 571 | async def test_snapshot_child_rejects_mismatch_before_publication( |
| 572 | tmp_path: Path, monkeypatch: pytest.MonkeyPatch |
| 573 | ): |
| 574 | from vllm.entrypoints.launchers.api_server import entry as api_server |
| 575 | |
| 576 | canaries = iter((_oracle(), _oracle((42,), " Lyon"))) |
| 577 | |
| 578 | @asynccontextmanager |
| 579 | async def engine_context(): |
| 580 | yield object() |
| 581 | |
| 582 | async def next_canary(_engine: object) -> Oracle: |
| 583 | return next(canaries) |
| 584 | |
| 585 | async def return_none(_engine: object) -> None: |
| 586 | return None |
| 587 | |
| 588 | monkeypatch.setattr( |
| 589 | api_server, "build_async_engine_client", lambda _args: engine_context() |
| 590 | ) |
| 591 | monkeypatch.setattr(snapshot_server, "run_engine_canary", next_canary) |
| 592 | monkeypatch.setattr(snapshot_server, "_release_reloadable_state", return_none) |
| 593 | monkeypatch.setattr(snapshot_server, "_restore_reloadable_state", return_none) |
| 594 | monkeypatch.setattr(snapshot_server, "detach_snapshot_streams", lambda: None) |
| 595 | monkeypatch.setattr( |
| 596 | snapshot_server, |
| 597 | "write_ready_atomic", |
| 598 | lambda *_args: pytest.fail("mismatched snapshot was published"), |
| 599 | ) |
| 600 | control = snapshot_server.ControlArgs( |
| 601 | ready_file=tmp_path / "ready.json", |
| 602 | release_file=tmp_path / "release.json", |
| 603 | release_timeout_s=1, |
| 604 | ) |
| 605 | with pytest.raises(SnapshotCanaryError, match="rehearsal changed"): |
| 606 | await snapshot_server.run_vllm_snapshot_child(control, argparse.Namespace()) |
| 607 | |
| 608 | |
| 609 | def test_snapshot_release_marker_ignores_unknown_keys(tmp_path: Path): |
| 610 | marker = tmp_path / "release.json" |
| 611 | marker.write_text(json.dumps({"release": True, "port": 8000, "unknown": 1})) |
| 612 | |
| 613 | listener = snapshot_server.read_release_marker(marker) |
| 614 | |
| 615 | assert listener == snapshot_server.ListenerConfig(host=None, port=8000) |
| 616 | |
| 617 | |
| 618 | @pytest.mark.parametrize( |
| 619 | "payload", |
| 620 | [ |
| 621 | ["release"], |
| 622 | {"port": 8000}, |
| 623 | {"release": "true", "port": 8000}, |
| 624 | {"release": 1, "port": 8000}, |
| 625 | {"release": True, "port": 0}, |
| 626 | {"release": True, "port": 70000}, |
| 627 | {"release": True, "port": True}, |
| 628 | {"release": True, "port": "8000"}, |
| 629 | {"release": True, "port": 8000, "host": 5}, |
| 630 | ], |
| 631 | ) |
| 632 | def test_snapshot_release_marker_rejects_an_invalid_payload( |
| 633 | tmp_path: Path, payload: object |
| 634 | ): |
| 635 | marker = tmp_path / "release.json" |
| 636 | marker.write_text(json.dumps(payload)) |
| 637 | |
| 638 | with pytest.raises(snapshot_server.SnapshotBarrierError, match="release marker"): |
| 639 | snapshot_server.read_release_marker(marker) |
| 640 | |
| 641 | |
| 642 | def test_snapshot_create_rolls_back_failed_dump(tmp_path: Path): |
| 643 | target = tmp_path / "snapshot" |
| 644 | tools = _fake_snapshot_tools() |
| 645 | tools.dump.side_effect = RuntimeError("dump failed") |
| 646 | with pytest.raises(RuntimeError, match="dump failed"): |
| 647 | create_snapshot( |
| 648 | argparse.Namespace(snapshot_dir=str(target), model_tag="Qwen/Qwen3-0.6B"), |
| 649 | tools=tools, |
| 650 | ) |
| 651 | tools.abort_create.assert_called_once_with(100) |
| 652 | assert not target.exists() |
| 653 | |
| 654 | |
| 655 | def _restore_fixture(tmp_path: Path) -> argparse.Namespace: |
| 656 | artifact = tmp_path / "snapshot" |
| 657 | artifact.mkdir(mode=0o700) |
| 658 | write_manifest_atomic(artifact, _manifest()) |
| 659 | return argparse.Namespace(snapshot_dir=str(artifact), host="127.0.0.1", port=9000) |
| 660 | |
| 661 | |
| 662 | def test_snapshot_restore_rejects_identity_mismatch(tmp_path: Path): |
| 663 | args = _restore_fixture(tmp_path) |
| 664 | |
| 665 | identity_mismatch = _fake_snapshot_tools() |
| 666 | identity_mismatch.current_identity.side_effect = lambda _gpu_uuid: ( |
| 667 | _runtime_identity(vllm_version="different") |
| 668 | ) |
| 669 | with pytest.raises(SnapshotCompatibilityError, match="vllm_version"): |
| 670 | restore_snapshot(args, tools=identity_mismatch) |
| 671 | identity_mismatch.restore.assert_not_called() |
| 672 | |
| 673 | |
| 674 | def test_snapshot_identity_mismatch_names_every_differing_field(): |
| 675 | with pytest.raises(SnapshotCompatibilityError) as mismatch: |
| 676 | validate_identity( |
| 677 | _manifest(), |
| 678 | _runtime_identity(vllm_version="different", gpu_uuid="GPU-other"), |
| 679 | ) |
| 680 | |
| 681 | assert str(mismatch.value) == "snapshot mismatch: vllm_version, gpu_uuid" |
| 682 | |
| 683 | |
| 684 | def test_snapshot_restore_cleans_oracle_mismatch(tmp_path: Path): |
| 685 | args = _restore_fixture(tmp_path) |
| 686 | oracle_mismatch = _fake_snapshot_tools() |
| 687 | oracle_mismatch.request_oracle.return_value = _oracle(logprob=-0.5) |
| 688 | with pytest.raises(SnapshotRestoreError, match="oracle mismatch"): |
| 689 | restore_snapshot(args, tools=oracle_mismatch) |
| 690 | oracle_mismatch.cleanup.assert_called_once() |
| 691 | oracle_mismatch.complete_restore.assert_not_called() |
| 692 | assert oracles_match(_oracle(logprob=-100.0), _oracle(logprob=-100.0005)) |
| 693 | assert not oracles_match(_oracle(logprob=-100.0), _oracle(logprob=-100.002)) |
| 694 | |
| 695 | |
| 696 | @pytest.mark.parametrize( |
| 697 | ("invalid_update", "diagnostic"), |
| 698 | [ |
| 699 | ({"unexpected": True}, "unexpected"), |
| 700 | ({"schema_version": True}, "schema_version"), |
| 701 | ({"boundary": "unknown"}, "boundary"), |
| 702 | ({"process_tree": [100, 100]}, "root"), |
| 703 | ({"cuda_holders": [999]}, "root"), |
| 704 | ({"oracle_token_ids": [1, 2]}, "oracle_token_ids"), |
| 705 | ({"oracle_sampled_token_logprob": 0}, "oracle_sampled_token_logprob"), |
| 706 | ( |
| 707 | {"oracle_sampled_token_logprob": float("nan")}, |
| 708 | "oracle_sampled_token_logprob", |
| 709 | ), |
| 710 | ], |
| 711 | ) |
| 712 | def test_snapshot_manifest_validation( |
| 713 | tmp_path: Path, |
| 714 | invalid_update: dict[str, object], |
| 715 | diagnostic: str, |
| 716 | ): |
| 717 | artifact = tmp_path / "snapshot" |
| 718 | artifact.mkdir(mode=0o700) |
| 719 | write_manifest_atomic(artifact, _manifest()) |
| 720 | invalid = _manifest().model_dump(mode="json") |
| 721 | invalid.update(invalid_update) |
| 722 | (artifact / "manifest.json").write_text(json.dumps(invalid)) |
| 723 | with pytest.raises( |
| 724 | SnapshotCompatibilityError, |
| 725 | match=f"invalid snapshot manifest: {diagnostic}", |
| 726 | ): |
| 727 | read_manifest(artifact) |
| 728 | |
| 729 | |
| 730 | def _local_restore_artifact(tmp_path: Path, name: str = "snapshot") -> Path: |
| 731 | artifact = tmp_path / name |
| 732 | artifact.mkdir(mode=0o700) |
| 733 | (artifact / "child.log").write_bytes(b"startup log") |
| 734 | (artifact / "child.log.snapshot-size").write_text("11\n") |
| 735 | return artifact |
| 736 | |
| 737 | |
| 738 | def _mark_snapshot_pids_free(monkeypatch: pytest.MonkeyPatch) -> list[int]: |
| 739 | probed_pids: list[int] = [] |
| 740 | |
| 741 | def missing(pid: int, _signal: int) -> None: |
| 742 | probed_pids.append(pid) |
| 743 | raise ProcessLookupError |
| 744 | |
| 745 | monkeypatch.setattr(snapshot_runtime.os, "kill", missing) |
| 746 | return probed_pids |
| 747 | |
| 748 | |
| 749 | @pytest.mark.parametrize( |
| 750 | ("probe_error", "diagnostic"), |
| 751 | [ |
| 752 | (None, "already occupied: 100"), |
| 753 | (PermissionError(), "already occupied: 100"), |
| 754 | (OSError("probe failed"), "availability probe failed: 100"), |
| 755 | ], |
| 756 | ) |
| 757 | def test_snapshot_restore_rejects_unavailable_pid_before_mutation( |
| 758 | tmp_path: Path, |
| 759 | monkeypatch: pytest.MonkeyPatch, |
| 760 | probe_error: OSError | None, |
| 761 | diagnostic: str, |
| 762 | ): |
| 763 | args = _restore_fixture(tmp_path) |
| 764 | artifact = Path(args.snapshot_dir) |
| 765 | (artifact / "child.log").write_bytes(b"startup log") |
| 766 | (artifact / "release.json").write_bytes(b"release sentinel") |
| 767 | before = {path.name: path.read_bytes() for path in artifact.iterdir()} |
| 768 | tools = LocalSnapshotTools() |
| 769 | |
| 770 | def fail_probe(_pid: int, _signal: int) -> None: |
| 771 | if probe_error is not None: |
| 772 | raise probe_error |
| 773 | |
| 774 | monkeypatch.setattr(snapshot_runtime.os, "kill", fail_probe) |
| 775 | monkeypatch.setattr(tools, "preflight", lambda *_args: None) |
| 776 | monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity()) |
| 777 | monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) |
| 778 | |
| 779 | with pytest.raises(SnapshotRestoreError, match=diagnostic): |
| 780 | restore_snapshot(args, tools=tools) |
| 781 | assert {path.name: path.read_bytes() for path in artifact.iterdir()} == before |
| 782 | |
| 783 | |
| 784 | @pytest.mark.parametrize("mode", ["pinned", "unpinned", "cleanup-fails"]) |
| 785 | def test_snapshot_restore_terminates_only_a_pinned_tree( |
| 786 | tmp_path: Path, |
| 787 | monkeypatch: pytest.MonkeyPatch, |
| 788 | mode: str, |
| 789 | ): |
| 790 | _mark_snapshot_pids_free(monkeypatch) |
| 791 | artifact = _local_restore_artifact(tmp_path) |
| 792 | manifest = _manifest() |
| 793 | root_pid = manifest.process_tree[0] |
| 794 | tools = LocalSnapshotTools() |
| 795 | terminated: list[int] = [] |
| 796 | |
| 797 | def terminate_and_wait(pid: int) -> None: |
| 798 | terminated.append(pid) |
| 799 | if mode == "cleanup-fails": |
| 800 | raise SnapshotRestoreError("restored process cleanup is incomplete") |
| 801 | tools._restored_processes.pop(pid) |
| 802 | |
| 803 | def pin( |
| 804 | _artifact: Path, process_tree: tuple[int, ...], _holders: tuple[int, ...] |
| 805 | ) -> None: |
| 806 | if mode == "unpinned": |
| 807 | raise SnapshotRestoreError( |
| 808 | "restored session does not match the captured process tree" |
| 809 | ) |
| 810 | tools._restored_processes[process_tree[0]] = () |
| 811 | |
| 812 | monkeypatch.setattr(tools, "cleanup", terminate_and_wait) |
| 813 | monkeypatch.setattr(tools, "_pin_restored_tree", pin) |
| 814 | monkeypatch.setattr(tools, "_criu", lambda *_args: None) |
| 815 | monkeypatch.setattr( |
| 816 | tools, |
| 817 | "_run", |
| 818 | lambda command, **_kwargs: subprocess.CompletedProcess(command, 0, "333\n", ""), |
| 819 | ) |
| 820 | |
| 821 | with pytest.raises(SnapshotRestoreError, match="does not match") as failure: |
| 822 | tools.restore(artifact, manifest) |
| 823 | |
| 824 | # A pinned tree is ours: it must be terminated and waited for, not left to |
| 825 | # honour the abort marker. An unpinned PID is never signalled. |
| 826 | assert terminated == ([] if mode == "unpinned" else [root_pid]) |
| 827 | if mode == "cleanup-fails": |
| 828 | # The failure to terminate is reported beside the primary error, and a |
| 829 | # tree that may still be running stays pinned. |
| 830 | assert "restore cleanup failed" in str(failure.value) |
| 831 | assert tools._restored_processes == {root_pid: ()} |
| 832 | else: |
| 833 | assert tools._restored_processes == {} |
| 834 | assert json.loads((artifact / "release.json").read_text())["release"] is False |
| 835 | |
| 836 | |
| 837 | def test_snapshot_private_path_and_link_remap_security( |
| 838 | tmp_path: Path, monkeypatch: pytest.MonkeyPatch |
| 839 | ): |
| 840 | _mark_snapshot_pids_free(monkeypatch) |
| 841 | real = tmp_path / "real" |
| 842 | real.mkdir(mode=0o700) |
| 843 | linked = tmp_path / "linked" |
| 844 | linked.symlink_to(real, target_is_directory=True) |
| 845 | with pytest.raises(SnapshotSecurityError, match="symlink"): |
| 846 | validate_artifact_root(linked, creating=False) |
| 847 | |
| 848 | unsafe = tmp_path / "unsafe" |
| 849 | unsafe.mkdir() |
| 850 | unsafe.chmod(0o777) |
| 851 | with pytest.raises(SnapshotSecurityError, match="world-writable"): |
| 852 | validate_artifact_root(unsafe / "snapshot", creating=True) |
| 853 | |
| 854 | artifact = _local_restore_artifact(tmp_path) |
| 855 | saved_remaps = artifact / "link-remaps" |
| 856 | saved_remaps.mkdir() |
| 857 | (saved_remaps / "link_remap.270").write_bytes(b"semaphore state") |
| 858 | shm_dir = tmp_path / "shm" |
| 859 | shm_dir.mkdir() |
| 860 | target = shm_dir / "link_remap.270" |
| 861 | replacement = tmp_path / "replacement" |
| 862 | replacement.write_bytes(b"semaphore state") |
| 863 | target.symlink_to(replacement) |
| 864 | tools = LocalSnapshotTools() |
| 865 | tools.shm_dir = shm_dir |
| 866 | monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) |
| 867 | with pytest.raises(SnapshotRestoreError, match="conflicting CRIU link remap"): |
| 868 | tools.restore(artifact, _manifest()) |
| 869 | |
| 870 | |
| 871 | @pytest.mark.parametrize( |
| 872 | ("blocked_state", "message"), |
| 873 | [("tcp", "external established TCP"), ("io_uring", "kernel.io_uring_disabled=1")], |
| 874 | ) |
| 875 | def test_snapshot_rejects_unsafe_process_state_before_criu( |
| 876 | tmp_path: Path, monkeypatch: pytest.MonkeyPatch, blocked_state: str, message: str |
| 877 | ): |
| 878 | tools = LocalSnapshotTools() |
| 879 | monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101)) |
| 880 | monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",)) |
| 881 | monkeypatch.setattr( |
| 882 | tools, |
| 883 | "_descriptor_inventory", |
| 884 | lambda _tree: ( |
| 885 | (101,) if blocked_state == "io_uring" else (), |
| 886 | {41: 101}, |
| 887 | ), |
| 888 | ) |
| 889 | tcp_record = _TcpSocketRecord( |
| 890 | family="AF_INET", |
| 891 | local_raw="3C00000A:D6F6", |
| 892 | remote_raw="01010101:01BB" if blocked_state == "tcp" else "00000000:0000", |
| 893 | inode=41, |
| 894 | ) |
| 895 | monkeypatch.setattr(tools, "_tcp_records", lambda: (tcp_record,)) |
| 896 | with pytest.raises(SnapshotCreateError, match=message) as excinfo: |
| 897 | tools.inventory(100, tmp_path) |
| 898 | if blocked_state == "tcp": |
| 899 | assert "pid 101" in str(excinfo.value) |
| 900 | assert "1.1.1.1:443" in str(excinfo.value) |
| 901 | |
| 902 | |
| 903 | def test_decode_endpoint_families(): |
| 904 | assert _decode_endpoint("AF_INET", "3C00000A:D6F6") == "10.0.0.60:55030" |
| 905 | assert ( |
| 906 | _decode_endpoint("AF_INET6", "00000000000000000000000001000000:01BB") |
| 907 | == "[::1]:443" |
| 908 | ) |
| 909 | |
| 910 | |
| 911 | def test_tcp_records_skips_missing_tables(tmp_path, monkeypatch): |
| 912 | tcp, tcp6 = tmp_path / "tcp", tmp_path / "tcp6" |
| 913 | contents = "header\n0: 0100007F:1F90 00000000:0000 01 0 0 0 0 0 41\n" |
| 914 | tcp.write_text(contents) |
| 915 | tables = (("AF_INET", tcp), ("AF_INET6", tcp6)) |
| 916 | monkeypatch.setattr(snapshot_runtime, "_TCP_TABLES", tables) |
| 917 | |
| 918 | tools = LocalSnapshotTools() |
| 919 | assert tuple(record.family for record in tools._tcp_records()) == ("AF_INET",) |
| 920 | |
| 921 | tcp6.write_text(contents) |
| 922 | families = tuple(record.family for record in tools._tcp_records()) |
| 923 | assert families == ("AF_INET", "AF_INET6") |
| 924 | |
| 925 | |
| 926 | def test_snapshot_manifest_records_external_cache_files(tmp_path: Path): |
| 927 | artifact = tmp_path / "snapshot" |
| 928 | artifact.mkdir(mode=0o700) |
| 929 | recorded = ( |
| 930 | (str(tmp_path / "kernel.so"), "b" * 64), |
| 931 | (str(tmp_path / "launcher.lock"), "locked"), |
| 932 | ) |
| 933 | write_manifest_atomic(artifact, _manifest(external_cache_files=recorded)) |
| 934 | |
| 935 | assert read_manifest(artifact).external_cache_files == recorded |
| 936 | |
| 937 | # A manifest written before this field existed still loads. |
| 938 | legacy = _manifest().model_dump(mode="json") |
| 939 | del legacy["external_cache_files"] |
| 940 | (artifact / "manifest.json").write_text(json.dumps(legacy)) |
| 941 | |
| 942 | assert read_manifest(artifact).external_cache_files == () |
| 943 | |
| 944 | |
| 945 | def _pin_generated_cache_roots(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: |
| 946 | """Keep the resolver off whatever cache roots the host running this has.""" |
| 947 | monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path / "vllm_cache")) |
| 948 | monkeypatch.setenv("TRITON_CACHE_DIR", str(tmp_path / "triton_cache_default")) |
| 949 | monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", str(tmp_path / "inductor_cache")) |
| 950 | |
| 951 | |
| 952 | def _unreadable(_path: Path) -> str: |
| 953 | raise OSError("cache file is unreadable") |
| 954 | |
| 955 | |
| 956 | def test_snapshot_inventory_records_open_generated_cache_files( |
| 957 | tmp_path: Path, monkeypatch: pytest.MonkeyPatch |
| 958 | ): |
| 959 | cache_root = tmp_path / "triton_cache" |
| 960 | cache_root.mkdir() |
| 961 | kernel = cache_root / "kernel.so" |
| 962 | kernel.write_bytes(b"compiled kernel") |
| 963 | lock = cache_root / "launcher.lock" |
| 964 | lock.write_bytes(b"held by pid 100") |
| 965 | weights = tmp_path / "model.safetensors" |
| 966 | weights.write_bytes(b"weights") |
| 967 | artifact = tmp_path / "snapshot" |
| 968 | artifact.mkdir(mode=0o700) |
| 969 | (artifact / "child.log.lock").write_bytes(b"") |
| 970 | _pin_generated_cache_roots(monkeypatch, tmp_path) |
| 971 | monkeypatch.setenv("TRITON_CACHE_DIR", str(cache_root)) |
| 972 | |
| 973 | targets = { |
| 974 | 100: (str(kernel), str(weights), "socket:[41]"), |
| 975 | 101: (str(lock), str(cache_root), str(artifact / "child.log.lock")), |
| 976 | 102: (f"{cache_root / 'evicted.so'} (deleted)",), |
| 977 | } |
| 978 | tools = LocalSnapshotTools() |
| 979 | monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101, 102)) |
| 980 | monkeypatch.setattr(tools, "_descriptor_targets", lambda pid: targets[pid]) |
| 981 | monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",)) |
| 982 | monkeypatch.setattr(tools, "_tcp_records", lambda: ()) |
| 983 | monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity()) |
| 984 | |
| 985 | manifest = tools.make_manifest( |
| 986 | argparse.Namespace( |
| 987 | model_tag="Qwen/Qwen3-0.6B", |
| 988 | revision=_MODEL_REVISION, |
| 989 | served_model_name=None, |
| 990 | tokenizer_revision=None, |
| 991 | ), |
| 992 | ("Qwen/Qwen3-0.6B",), |
| 993 | tools.inventory(100, artifact), |
| 994 | _oracle(), |
| 995 | artifact, |
| 996 | ) |
| 997 | |
| 998 | # Weights, sockets, the cache directory itself, files inside the artifact, |
| 999 | # and deleted targets that CRIU remaps are all left out. |
| 1000 | assert manifest.external_cache_files == ( |
| 1001 | (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()), |
| 1002 | (str(lock), "locked"), |
| 1003 | ) |
| 1004 | |
| 1005 | |
| 1006 | def test_snapshot_inventory_marks_a_cache_file_it_cannot_read( |
| 1007 | tmp_path: Path, monkeypatch: pytest.MonkeyPatch |
| 1008 | ): |
| 1009 | cache_root = tmp_path / "triton_cache" |
| 1010 | cache_root.mkdir() |
| 1011 | kernel = cache_root / "kernel.so" |
| 1012 | kernel.write_bytes(b"compiled kernel") |
| 1013 | lock = cache_root / "launcher.lock" |
| 1014 | lock.write_bytes(b"held by pid 100") |
| 1015 | artifact = tmp_path / "snapshot" |
| 1016 | artifact.mkdir(mode=0o700) |
| 1017 | _pin_generated_cache_roots(monkeypatch, tmp_path) |
| 1018 | monkeypatch.setenv("TRITON_CACHE_DIR", str(cache_root)) |
| 1019 | |
| 1020 | tools = LocalSnapshotTools() |
| 1021 | monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100,)) |
| 1022 | monkeypatch.setattr( |
| 1023 | tools, "_descriptor_targets", lambda _pid: (str(kernel), str(lock)) |
| 1024 | ) |
| 1025 | monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("100, GPU-abc",)) |
| 1026 | monkeypatch.setattr(tools, "_tcp_records", lambda: ()) |
| 1027 | monkeypatch.setattr(tools, "_sha256", _unreadable) |
| 1028 | |
| 1029 | # A file that cannot be digested keeps its own marker, so the manifest says |
| 1030 | # which entries restore can check for existence only. |
| 1031 | assert tools.inventory(100, artifact).external_cache_files == ( |
| 1032 | (str(kernel), "unreadable"), |
| 1033 | (str(lock), "locked"), |
| 1034 | ) |
| 1035 | |
| 1036 | |
| 1037 | def test_snapshot_inventory_matches_resolved_cache_roots_not_siblings( |
| 1038 | tmp_path: Path, monkeypatch: pytest.MonkeyPatch |
| 1039 | ): |
| 1040 | # Descriptor targets come back kernel-canonical, so the fixture must be too. |
| 1041 | base = Path(os.path.realpath(tmp_path)) |
| 1042 | real_cache = base / "mnt_cache" |
| 1043 | real_cache.mkdir() |
| 1044 | kernel = real_cache / "kernel.so" |
| 1045 | kernel.write_bytes(b"compiled kernel") |
| 1046 | linked_cache = base / "linked_cache" |
| 1047 | linked_cache.symlink_to(real_cache, target_is_directory=True) |
| 1048 | sibling = base / "vllm_cache-backup" |
| 1049 | sibling.mkdir() |
| 1050 | stale = sibling / "stale.so" |
| 1051 | stale.write_bytes(b"stale kernel") |
| 1052 | artifact = base / "snapshot" |
| 1053 | artifact.mkdir(mode=0o700) |
| 1054 | |
| 1055 | _pin_generated_cache_roots(monkeypatch, base) |
| 1056 | monkeypatch.setenv("TRITON_CACHE_DIR", str(linked_cache)) |
| 1057 | |
| 1058 | tools = LocalSnapshotTools() |
| 1059 | monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100,)) |
| 1060 | monkeypatch.setattr( |
| 1061 | tools, "_descriptor_targets", lambda _pid: (str(kernel), str(stale)) |
| 1062 | ) |
| 1063 | monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("100, GPU-abc",)) |
| 1064 | monkeypatch.setattr(tools, "_tcp_records", lambda: ()) |
| 1065 | |
| 1066 | # A cache root reached through a symlink still covers the file the kernel |
| 1067 | # reports by its real path, and a directory that merely shares the root's |
| 1068 | # name prefix is not part of the cache at all. |
| 1069 | assert tools.inventory(100, artifact).external_cache_files == ( |
| 1070 | (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()), |
| 1071 | ) |
| 1072 | |
| 1073 | |
| 1074 | def test_snapshot_restore_rejects_rotated_cache_file_before_mutation( |
| 1075 | tmp_path: Path, monkeypatch: pytest.MonkeyPatch |
| 1076 | ): |
| 1077 | _mark_snapshot_pids_free(monkeypatch) |
| 1078 | kernel = tmp_path / "kernel.so" |
| 1079 | kernel.write_bytes(b"compiled kernel") |
| 1080 | lock = tmp_path / "launcher.lock" |
| 1081 | lock.write_bytes(b"held by pid 100") |
| 1082 | manifest = _manifest( |
| 1083 | external_cache_files=( |
| 1084 | (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()), |
| 1085 | (str(lock), "locked"), |
| 1086 | ) |
| 1087 | ) |
| 1088 | artifact = _local_restore_artifact(tmp_path) |
| 1089 | before = {path.name: path.read_bytes() for path in artifact.iterdir()} |
| 1090 | tools = LocalSnapshotTools() |
| 1091 | monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) |
| 1092 | |
| 1093 | kernel.write_bytes(b"recompiled kernel") |
| 1094 | with pytest.raises(SnapshotRestoreError, match="changed since capture") as changed: |
| 1095 | tools.restore(artifact, manifest) |
| 1096 | assert str(kernel) in str(changed.value) |
| 1097 | |
| 1098 | # Restoring the digest gets past the first entry, so the lock file is only |
| 1099 | # checked for existence: the second failure proves the first one passed. |
| 1100 | kernel.write_bytes(b"compiled kernel") |
| 1101 | lock.unlink() |
| 1102 | with pytest.raises(SnapshotRestoreError, match="vanished since capture") as gone: |
| 1103 | tools.restore(artifact, manifest) |
| 1104 | assert str(lock) in str(gone.value) |
| 1105 | |
| 1106 | assert {path.name: path.read_bytes() for path in artifact.iterdir()} == before |
| 1107 | |
| 1108 | |
| 1109 | def test_snapshot_restore_rejects_a_cache_file_it_cannot_read( |
| 1110 | tmp_path: Path, monkeypatch: pytest.MonkeyPatch |
| 1111 | ): |
| 1112 | _mark_snapshot_pids_free(monkeypatch) |
| 1113 | kernel = tmp_path / "kernel.so" |
| 1114 | kernel.write_bytes(b"compiled kernel") |
| 1115 | manifest = _manifest( |
| 1116 | external_cache_files=( |
| 1117 | (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()), |
| 1118 | ) |
| 1119 | ) |
| 1120 | artifact = _local_restore_artifact(tmp_path) |
| 1121 | tools = LocalSnapshotTools() |
| 1122 | monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) |
| 1123 | monkeypatch.setattr(tools, "_sha256", _unreadable) |
| 1124 | |
| 1125 | # A recorded digest that cannot be read again is a failed check, not a file |
| 1126 | # to wave through on existence. |
| 1127 | with pytest.raises( |
| 1128 | SnapshotRestoreError, match="unreadable since capture" |
| 1129 | ) as blocked: |
| 1130 | tools.restore(artifact, manifest) |
| 1131 | assert str(kernel) in str(blocked.value) |
| 1132 | |
| 1133 | |
| 1134 | def test_snapshot_runtime_installer_logs_the_failing_exit_code(): |
| 1135 | installer = Path(__file__).parents[3] / "tools" / "install_snapshot_runtime.sh" |
| 1136 | if not installer.is_file(): |
| 1137 | pytest.skip("installer script is not shipped in this test image") |
| 1138 | timing = [ |
| 1139 | line |
| 1140 | for line in installer.read_text().splitlines() |
| 1141 | if line.startswith(("_snapshot_install_started=", "trap ")) |
| 1142 | ] |
| 1143 | assert len(timing) == 2 |
| 1144 | |
| 1145 | result = subprocess.run( |
| 1146 | ["bash", "-c", "\n".join([*timing, "exit 7"])], |
| 1147 | capture_output=True, |
| 1148 | text=True, |
| 1149 | ) |
| 1150 | |
| 1151 | assert result.returncode == 7 |
| 1152 | assert "(exit 7)" in result.stdout |