1437 lines
50 KiB
Python
1437 lines
50 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Unit tests for the `vllm launch` CLI subcommand."""
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from unittest.mock import create_autospec, patch
|
|
|
|
import pytest
|
|
|
|
import vllm.snapshot.runtime as snapshot_runtime
|
|
import vllm.snapshot.server as snapshot_server
|
|
from vllm.entrypoints.cli import snapshot as snapshot_cli
|
|
from vllm.entrypoints.cli.launch import (
|
|
LaunchSubcommand,
|
|
RenderSubcommand,
|
|
cmd_init,
|
|
)
|
|
from vllm.snapshot.controller import (
|
|
LocalSnapshotTools,
|
|
SnapshotCreateError,
|
|
SnapshotRestoreError,
|
|
_child_engine_argv,
|
|
create_snapshot,
|
|
restore_snapshot,
|
|
)
|
|
from vllm.snapshot.manifest import (
|
|
SnapshotCompatibilityError,
|
|
SnapshotManifest,
|
|
SnapshotRuntimeIdentity,
|
|
SnapshotSecurityError,
|
|
read_manifest,
|
|
validate_artifact_root,
|
|
validate_identity,
|
|
write_manifest_atomic,
|
|
)
|
|
from vllm.snapshot.runtime import ProcessInventory, _decode_endpoint, _TcpSocketRecord
|
|
from vllm.snapshot.server import SnapshotCanaryError
|
|
from vllm.snapshot.types import Oracle, oracles_match
|
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
|
|
|
|
|
@pytest.fixture
|
|
def launch_parser():
|
|
parser = FlexibleArgumentParser(description="test")
|
|
subparsers = parser.add_subparsers(required=False, dest="subparser")
|
|
LaunchSubcommand().subparser_init(subparsers)
|
|
return parser
|
|
|
|
|
|
def test_subcommand_name():
|
|
assert LaunchSubcommand().name == "launch"
|
|
|
|
|
|
def test_cmd_init_returns_subcommand():
|
|
result = cmd_init()
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], LaunchSubcommand)
|
|
|
|
|
|
# -- Parsing: `vllm launch render` --
|
|
|
|
|
|
def test_parse_launch_render(launch_parser):
|
|
args = launch_parser.parse_args(["launch", "render", "--model", "test-model"])
|
|
assert args.launch_component == "render"
|
|
|
|
|
|
def test_parse_launch_requires_component(launch_parser):
|
|
with pytest.raises(SystemExit):
|
|
launch_parser.parse_args(["launch", "--model", "test-model"])
|
|
|
|
|
|
def test_parse_launch_invalid_component(launch_parser):
|
|
with pytest.raises(SystemExit):
|
|
launch_parser.parse_args(["launch", "unknown", "--model", "test-model"])
|
|
|
|
|
|
# -- Dispatch --
|
|
|
|
|
|
def test_cmd_launch_render_calls_run():
|
|
args = argparse.Namespace(model_tag=None, model="test-model")
|
|
with patch("vllm.entrypoints.cli.launch.uvloop.run") as mock_uvloop_run:
|
|
RenderSubcommand.cmd(args)
|
|
mock_uvloop_run.assert_called_once()
|
|
|
|
|
|
def test_cmd_launch_model_tag_overrides():
|
|
args = argparse.Namespace(
|
|
model_tag="tag-model",
|
|
model="original-model",
|
|
launch_command=lambda a: None,
|
|
)
|
|
LaunchSubcommand.cmd(args)
|
|
assert args.model == "tag-model"
|
|
|
|
|
|
def test_cmd_launch_model_tag_none():
|
|
args = argparse.Namespace(
|
|
model_tag=None,
|
|
model="original-model",
|
|
launch_command=lambda a: None,
|
|
)
|
|
LaunchSubcommand.cmd(args)
|
|
assert args.model == "original-model"
|
|
|
|
|
|
def test_cmd_dispatches():
|
|
called = {}
|
|
|
|
def fake_dispatch(args):
|
|
called["args"] = args
|
|
|
|
args = argparse.Namespace(launch_command=fake_dispatch)
|
|
LaunchSubcommand.cmd(args)
|
|
assert "args" in called
|
|
|
|
|
|
# -- Module registration --
|
|
|
|
|
|
def test_subparser_init_returns_parser():
|
|
parser = FlexibleArgumentParser(description="test")
|
|
subparsers = parser.add_subparsers(required=False, dest="subparser")
|
|
result = LaunchSubcommand().subparser_init(subparsers)
|
|
assert isinstance(result, FlexibleArgumentParser)
|
|
|
|
|
|
def test_launch_registered_in_main():
|
|
"""Verify that launch module is importable as a CLI module."""
|
|
import vllm.entrypoints.cli.launch as launch_module
|
|
|
|
assert hasattr(launch_module, "cmd_init")
|
|
subcmds = launch_module.cmd_init()
|
|
assert any(s.name == "launch" for s in subcmds)
|
|
|
|
|
|
_MODEL_REVISION = "c1899de289a04d12100db370d81485cdf75e47ca"
|
|
|
|
|
|
def _run_cli(argv: list[str]) -> None:
|
|
from vllm.entrypoints.cli import main as cli_main
|
|
|
|
with patch.object(sys, "argv", ["vllm", *argv]):
|
|
cli_main.main()
|
|
|
|
|
|
def test_snapshot_environment_contract(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
capsys: pytest.CaptureFixture[str],
|
|
):
|
|
from vllm.entrypoints.serve.utils import api_utils
|
|
|
|
monkeypatch.setattr(api_utils, "cli_env_setup", pytest.fail)
|
|
secret = "snapshot-secret"
|
|
monkeypatch.setenv("VLLM_API_KEY", secret)
|
|
monkeypatch.setenv("VLLM_USER_SETTING", secret)
|
|
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", secret)
|
|
environment = dict(LocalSnapshotTools()._environment_identity())
|
|
assert (
|
|
environment["VLLM_USER_SETTING"]
|
|
== hashlib.sha256(b"VLLM_USER_SETTING\0snapshot-secret").hexdigest()
|
|
)
|
|
assert (
|
|
environment["VLLM_USER_SETTING"] != environment["VLLM_WORKER_MULTIPROC_METHOD"]
|
|
)
|
|
assert secret not in str(environment)
|
|
assert "VLLM_API_KEY" not in environment
|
|
|
|
monkeypatch.delenv("VLLM_WORKER_MULTIPROC_METHOD")
|
|
implicit_worker_default = dict(LocalSnapshotTools()._environment_identity())
|
|
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
|
assert dict(LocalSnapshotTools()._environment_identity()) == implicit_worker_default
|
|
|
|
monkeypatch.setenv("VLLM_USER_SETTING", f"{secret}-changed")
|
|
changed_environment = dict(LocalSnapshotTools()._environment_identity())
|
|
assert changed_environment["VLLM_USER_SETTING"] != environment["VLLM_USER_SETTING"]
|
|
|
|
artifact = tmp_path / "snapshot"
|
|
artifact.mkdir(mode=0o700)
|
|
write_manifest_atomic(artifact, _manifest(environment=tuple(environment.items())))
|
|
before = os.environ.copy()
|
|
|
|
_run_cli(["snapshot", "inspect", str(artifact)])
|
|
|
|
assert os.environ == before
|
|
output = capsys.readouterr().out
|
|
assert secret not in output
|
|
inspected = json.loads(output)
|
|
inspected_environment = dict(inspected["environment"])
|
|
assert inspected_environment == environment
|
|
|
|
tools = LocalSnapshotTools()
|
|
monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity())
|
|
manifest = tools.make_manifest(
|
|
argparse.Namespace(
|
|
hf_token=secret,
|
|
model_tag="Qwen/Qwen3-0.6B",
|
|
revision=_MODEL_REVISION,
|
|
served_model_name=None,
|
|
tokenizer_revision=None,
|
|
),
|
|
(
|
|
"Qwen/Qwen3-0.6B",
|
|
"--api-key",
|
|
secret,
|
|
f"{secret}-alternate",
|
|
f"--api-key={secret}",
|
|
"--hf-token",
|
|
secret,
|
|
f"--hf-token={secret}",
|
|
),
|
|
ProcessInventory(100, (100, 101), (101,), "GPU-abc"),
|
|
_oracle(),
|
|
tmp_path,
|
|
)
|
|
|
|
assert manifest.engine_argv == (
|
|
"Qwen/Qwen3-0.6B",
|
|
"--api-key",
|
|
"***",
|
|
"***",
|
|
"--api-key=***",
|
|
"--hf-token",
|
|
"***",
|
|
"--hf-token=***",
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("configured", "expected"),
|
|
[(None, "1"), ("0", "0"), ("1", "1")],
|
|
)
|
|
def test_snapshot_child_defaults_nccl_ib_without_mutating_parent(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
configured: str | None,
|
|
expected: str,
|
|
):
|
|
if configured is None:
|
|
monkeypatch.delenv("NCCL_IB_DISABLE", raising=False)
|
|
else:
|
|
monkeypatch.setenv("NCCL_IB_DISABLE", configured)
|
|
parent_environment = os.environ.copy()
|
|
process = create_autospec(subprocess.Popen, instance=True)
|
|
process.pid = 100
|
|
popen = create_autospec(subprocess.Popen, return_value=process)
|
|
monkeypatch.setattr(snapshot_runtime.subprocess, "Popen", popen)
|
|
|
|
LocalSnapshotTools().launch_child(tmp_path, ("model",))
|
|
|
|
assert popen.call_args.kwargs["env"]["NCCL_IB_DISABLE"] == expected
|
|
assert os.environ == parent_environment
|
|
|
|
|
|
def test_snapshot_environment_identity_applies_effective_nccl_ib_default(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
):
|
|
monkeypatch.delenv("NCCL_IB_DISABLE", raising=False)
|
|
implicit = LocalSnapshotTools()._environment_identity()
|
|
monkeypatch.setenv("NCCL_IB_DISABLE", "1")
|
|
explicit_disabled = LocalSnapshotTools()._environment_identity()
|
|
monkeypatch.setenv("NCCL_IB_DISABLE", "0")
|
|
explicit_enabled = LocalSnapshotTools()._environment_identity()
|
|
|
|
assert implicit == explicit_disabled
|
|
assert explicit_enabled != implicit
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"case",
|
|
[
|
|
(("--hf_token", "SECRET"), ("--hf_token", "***")),
|
|
(("--hf_token=SECRET",), ("--hf_token=***",)),
|
|
(("--api_key", "SECRET"), ("--api_key", "***")),
|
|
(("--api_key=SECRET",), ("--api_key=***",)),
|
|
(("--hf-tok", "SECRET"), ("--hf-tok", "***")),
|
|
(("--hf-tok=SECRET",), ("--hf-tok=***",)),
|
|
(("--hf_tok", "SECRET"), ("--hf_tok", "***")),
|
|
(("--api-k", "SECRET", "SECOND"), ("--api-k", "***", "***")),
|
|
(("--api-k=SECRET",), ("--api-k=***",)),
|
|
(("--hf-overrides", "{}"), ("--hf-overrides", "{}")),
|
|
],
|
|
)
|
|
def test_snapshot_manifest_redacts_accepted_secret_options(case):
|
|
engine_argv, expected = case
|
|
tools = create_autospec(LocalSnapshotTools, instance=True)
|
|
tools.current_identity.return_value = _runtime_identity()
|
|
tools._artifact_bytes.return_value = 0
|
|
manifest = LocalSnapshotTools.make_manifest(
|
|
tools,
|
|
snapshot_server.parse_vllm_args(["model", *engine_argv]),
|
|
engine_argv,
|
|
ProcessInventory(100, (100, 101), (101,), "GPU-abc"),
|
|
_oracle(),
|
|
Path(),
|
|
)
|
|
|
|
assert manifest.engine_argv == expected
|
|
if "SECRET" in " ".join(engine_argv):
|
|
assert "***" in " ".join(manifest.engine_argv)
|
|
assert "SECRET" not in " ".join(manifest.engine_argv)
|
|
|
|
|
|
def parse_snapshot(*argv: str):
|
|
parser = FlexibleArgumentParser()
|
|
subparsers = parser.add_subparsers(dest="subparser", required=True)
|
|
snapshot_cli.SnapshotSubcommand(
|
|
create_requested=argv[:1] == ("create",)
|
|
).subparser_init(subparsers)
|
|
return parser.parse_args(["snapshot", *argv])
|
|
|
|
|
|
@pytest.mark.parametrize("option", ["--snapshot-dir", "--snapshot_dir", "--snapshot-d"])
|
|
@pytest.mark.parametrize("joined", [False, True])
|
|
def test_snapshot_create_forwards_only_engine_options(option, joined, monkeypatch):
|
|
directory = [f"{option}=/tmp/snapshot"] if joined else [option, "/tmp/snapshot"]
|
|
argv = ["create", "Qwen/Qwen3-0.6B", "--revision", _MODEL_REVISION, *directory]
|
|
args = parse_snapshot(*argv)
|
|
monkeypatch.setattr(sys, "argv", ["vllm", "snapshot", *argv])
|
|
|
|
child = snapshot_server.parse_vllm_args(list(_child_engine_argv(args, None)))
|
|
|
|
assert child.model == "Qwen/Qwen3-0.6B"
|
|
assert child.revision == _MODEL_REVISION
|
|
assert child.enable_sleep_mode
|
|
|
|
|
|
@pytest.mark.parametrize("model_in_cli", [False, True])
|
|
def test_snapshot_create_config_preserves_cli_precedence(
|
|
tmp_path, monkeypatch, model_in_cli
|
|
):
|
|
config = tmp_path / "engine.yaml"
|
|
config.write_text("model: config-model\nmax-model-len: 512\ndtype: float16\n")
|
|
model = ["Qwen/Qwen3-0.6B"] if model_in_cli else []
|
|
argv = [
|
|
"create",
|
|
*model,
|
|
"--snapshot-dir",
|
|
str(tmp_path / "snapshot"),
|
|
"--config",
|
|
str(config),
|
|
"--revision",
|
|
_MODEL_REVISION,
|
|
"--max-model-len",
|
|
"1024",
|
|
]
|
|
args = parse_snapshot(*argv)
|
|
monkeypatch.setattr(sys, "argv", ["vllm", "snapshot", *argv])
|
|
child = snapshot_server.parse_vllm_args(list(_child_engine_argv(args, None)))
|
|
|
|
assert (
|
|
(args.model_tag or args.model)
|
|
== child.model
|
|
== ("Qwen/Qwen3-0.6B" if model_in_cli else "config-model")
|
|
)
|
|
assert args.max_model_len == child.max_model_len == 1024
|
|
assert args.dtype == child.dtype == "float16"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("probe", "status", "output"),
|
|
[
|
|
(None, 0, ""),
|
|
("nvidia-smi", 1, ""),
|
|
("find", 1, ""),
|
|
("ss", 1, ""),
|
|
("nvidia-smi", 0, "12345"),
|
|
("find", 0, "link_remap.fixture"),
|
|
("ss", 0, "LISTEN"),
|
|
],
|
|
)
|
|
def test_snapshot_e2e_cleanup_requires_successful_empty_probes(probe, status, output):
|
|
script = (
|
|
Path(__file__).resolve().parents[3]
|
|
/ ".buildkite/scripts/initialized-snapshot-e2e.sh"
|
|
)
|
|
if not script.exists():
|
|
pytest.skip("snapshot E2E script is not packaged in this test image")
|
|
source = script.read_text()
|
|
# Exercise the real probes without entering Docker/CRIU setup.
|
|
functions = source[source.index("link_remaps() {") : source.index("wait_clean() {")]
|
|
shell = (
|
|
"""set -Eeuo pipefail
|
|
GPU_UUID=fixture PORT_ONE=18001 PORT_TWO=18002 LINK_REMAP_BASELINE=''
|
|
SNAPSHOT_PIDS=(999999999)
|
|
nvidia-smi() { :; }
|
|
find() { :; }
|
|
ss() { :; }
|
|
"""
|
|
+ functions
|
|
)
|
|
if probe:
|
|
shell += f'\n{probe}() {{ printf %s "{output}"; return {status}; }}\n'
|
|
shell += "\nif state_clean; then exit 0; else exit 1; fi\n"
|
|
|
|
result = subprocess.run(["bash", "-c", shell], capture_output=True, timeout=10)
|
|
|
|
assert result.returncode == (0 if probe is None else 1), result.stderr
|
|
|
|
|
|
@pytest.mark.parametrize("container_attempted", [False, True])
|
|
def test_snapshot_e2e_cleanup_checks_state_only_after_container_attempt(
|
|
container_attempted: bool,
|
|
):
|
|
script = (
|
|
Path(__file__).resolve().parents[3]
|
|
/ ".buildkite/scripts/initialized-snapshot-e2e.sh"
|
|
)
|
|
if not script.exists():
|
|
pytest.skip("snapshot E2E script is not packaged in this test image")
|
|
source = script.read_text()
|
|
cleanup = source[source.index("cleanup() {") : source.index("on_exit() {")]
|
|
shell = (
|
|
"set -Eeuo pipefail\n"
|
|
"CONTAINER_NAME='' PORT_TWO=18002\n"
|
|
f"CHECK_CLEANUP_STATE={int(container_attempted)}\n"
|
|
"wait_clean() { echo checked-state; return 1; }\n"
|
|
+ cleanup
|
|
+ "\nif cleanup; then exit 0; else exit 1; fi\n"
|
|
)
|
|
|
|
result = subprocess.run(["bash", "-c", shell], capture_output=True, timeout=10)
|
|
|
|
assert result.returncode == int(container_attempted), result.stderr
|
|
assert (b"checked-state" in result.stdout) == container_attempted
|
|
|
|
|
|
def test_snapshot_create_cli_accepts_only_pinned_compact_mode(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
):
|
|
monkeypatch.setattr(snapshot_cli.platform, "system", lambda: "Linux")
|
|
monkeypatch.setattr(snapshot_cli.platform, "machine", lambda: "x86_64")
|
|
base = ("create", "Qwen/Qwen3-0.6B", "--snapshot-dir=/tmp/snapshot")
|
|
compact = parse_snapshot(*base, "--revision", _MODEL_REVISION)
|
|
|
|
snapshot_cli.validate_create_args(compact)
|
|
assert compact.model_tag == "Qwen/Qwen3-0.6B"
|
|
|
|
dispatched: list[argparse.Namespace] = []
|
|
monkeypatch.setattr(snapshot_cli, "run_create", dispatched.append)
|
|
_run_cli(["snapshot", *base, "--revision", _MODEL_REVISION])
|
|
assert dispatched[0].revision == _MODEL_REVISION
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("attribute", "value", "message"),
|
|
[
|
|
("prefill_context_parallel_size", 2, "prefill context parallel size 1"),
|
|
("model_tag", None, "immutable remote model"),
|
|
("logprobs_mode", "raw_logits", "log-probability mode"),
|
|
("speculative_config", {}, "does not support speculative decoding"),
|
|
("spec_method", "ngram", "does not support speculative decoding"),
|
|
("spec_model", "draft-model", "does not support speculative decoding"),
|
|
("spec_tokens", 3, "does not support speculative decoding"),
|
|
],
|
|
)
|
|
def test_snapshot_create_rejects_invalid_config(
|
|
attribute: str,
|
|
value: object,
|
|
message: str,
|
|
tmp_path: Path,
|
|
):
|
|
args = parse_snapshot(
|
|
"create",
|
|
"Qwen/Qwen3-0.6B",
|
|
"--snapshot-dir=/tmp/snapshot",
|
|
"--revision",
|
|
_MODEL_REVISION,
|
|
)
|
|
if attribute == "model_tag":
|
|
local_model = tmp_path / "model"
|
|
local_model.mkdir()
|
|
value = str(local_model)
|
|
setattr(args, attribute, value)
|
|
with pytest.raises(ValueError, match=message):
|
|
snapshot_cli.validate_create_args(args)
|
|
|
|
|
|
def _oracle(
|
|
token_ids: tuple[int, ...] = (12095,),
|
|
text: str = " Paris",
|
|
logprob: float = -0.125,
|
|
) -> Oracle:
|
|
return Oracle(token_ids, text, logprob)
|
|
|
|
|
|
def _manifest(**changes: object) -> SnapshotManifest:
|
|
manifest = SnapshotManifest(
|
|
schema_version=1,
|
|
boundary="post-engine-init-reloadable-state-released",
|
|
created_at="2026-08-06T00:00:00Z",
|
|
artifact_bytes=8,
|
|
vllm_version="0.1.dev1+snapshot",
|
|
python_version="3.12.3",
|
|
torch_version="2.9.0",
|
|
cuda_runtime="12.9",
|
|
driver_version="575.57.08",
|
|
criu_version="4.1",
|
|
cuda_checkpoint_sha256="a" * 64,
|
|
kernel_release="6.8.0",
|
|
host_id="host-a",
|
|
gpu_name="NVIDIA A10",
|
|
gpu_uuid="GPU-abc",
|
|
model="Qwen/Qwen3-0.6B",
|
|
served_model_name="Qwen/Qwen3-0.6B",
|
|
model_revision="model-sha",
|
|
tokenizer_revision="tokenizer-sha",
|
|
engine_argv=("Qwen/Qwen3-0.6B",),
|
|
environment=(("VLLM_USE_V1", "1"),),
|
|
process_tree=(100, 101),
|
|
cuda_holders=(101,),
|
|
oracle_token_ids=(12095,),
|
|
oracle_text=" Paris",
|
|
oracle_sampled_token_logprob=-0.125,
|
|
)
|
|
return SnapshotManifest.model_validate(
|
|
{**manifest.model_dump(mode="python"), **changes}, strict=True
|
|
)
|
|
|
|
|
|
def _runtime_identity(**changes: object) -> Any:
|
|
fields = SnapshotRuntimeIdentity.model_fields.keys()
|
|
payload = _manifest().model_dump(mode="python", include=fields)
|
|
return SnapshotRuntimeIdentity.model_validate({**payload, **changes}, strict=True)
|
|
|
|
|
|
def _fake_snapshot_tools() -> Any:
|
|
tools = create_autospec(LocalSnapshotTools, instance=True)
|
|
tools.launch_child.return_value = 100
|
|
tools.wait_ready.return_value = _oracle()
|
|
tools.inventory.return_value = ProcessInventory(100, (100, 101), (101,), "GPU-abc")
|
|
|
|
tools.current_identity.side_effect = lambda _gpu_uuid: _runtime_identity()
|
|
tools.restore.return_value = 100
|
|
tools.request_oracle.return_value = _oracle()
|
|
return tools
|
|
|
|
|
|
@pytest.mark.parametrize("value", ["0", "-1", "nan", "inf"])
|
|
def test_snapshot_timeout_must_be_positive_and_finite(
|
|
value: str, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
monkeypatch.setenv("VLLM_SNAPSHOT_TIMEOUT_S", value)
|
|
with pytest.raises(ValueError, match="positive finite number"):
|
|
LocalSnapshotTools()
|
|
|
|
|
|
def test_snapshot_create_abort_reaps_after_process_group_cleanup(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
tools = LocalSnapshotTools()
|
|
process = subprocess.Popen(
|
|
[sys.executable, "-c", "raise SystemExit(1)"],
|
|
stdin=subprocess.DEVNULL,
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
start_new_session=True,
|
|
)
|
|
tools._children[process.pid] = process
|
|
(tmp_path / "child.log").write_text("root failed")
|
|
killed: list[tuple[int, signal.Signals]] = []
|
|
monkeypatch.setattr(
|
|
os,
|
|
"killpg",
|
|
lambda pid, sig: killed.append((pid, sig)),
|
|
)
|
|
monkeypatch.setattr(os, "waitid", lambda *_args: object(), raising=False)
|
|
|
|
try:
|
|
with pytest.raises(SnapshotCreateError, match="root failed"):
|
|
tools.wait_ready(tmp_path, process.pid)
|
|
assert process.returncode is None
|
|
tools.abort_create(process.pid)
|
|
finally:
|
|
if process.pid in tools._children:
|
|
tools.abort_create(process.pid)
|
|
|
|
assert killed == [(process.pid, signal.SIGKILL)]
|
|
assert process.pid not in tools._children
|
|
|
|
|
|
def _supply_waitid(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""`wait_ready` binds Linux-only `os.waitid` before it reads the marker."""
|
|
monkeypatch.setattr(os, "waitid", lambda *_args: None, raising=False)
|
|
|
|
|
|
def test_snapshot_wait_ready_reads_back_the_oracle_the_child_wrote(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
_supply_waitid(monkeypatch)
|
|
snapshot_server.write_ready_atomic(tmp_path / "ready.json", _oracle())
|
|
|
|
assert LocalSnapshotTools().wait_ready(tmp_path, os.getpid()) == _oracle()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"payload",
|
|
[
|
|
{"token_ids": [12095], "sampled_token_logprob": -0.125},
|
|
{"token_ids": ["12095"], "text": " Paris", "sampled_token_logprob": -0.125},
|
|
],
|
|
)
|
|
def test_snapshot_wait_ready_rejects_a_malformed_marker(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, payload: dict[str, object]
|
|
):
|
|
_supply_waitid(monkeypatch)
|
|
(tmp_path / "ready.json").write_text(json.dumps(payload))
|
|
|
|
with pytest.raises(SnapshotCreateError, match="ready marker invalid"):
|
|
LocalSnapshotTools().wait_ready(tmp_path, os.getpid())
|
|
|
|
|
|
def test_snapshot_create_abort_rejects_surviving_process_group(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
):
|
|
tools = LocalSnapshotTools()
|
|
process = create_autospec(subprocess.Popen, instance=True)
|
|
process.wait.side_effect = subprocess.TimeoutExpired("snapshot child", 10)
|
|
tools._children[100] = process
|
|
monkeypatch.setattr(os, "killpg", lambda _pid, _sig: None)
|
|
|
|
with pytest.raises(SnapshotCreateError, match="survived SIGKILL"):
|
|
tools.abort_create(100)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_snapshot_child_rejects_mismatch_before_publication(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
from vllm.entrypoints.launchers.api_server import entry as api_server
|
|
|
|
canaries = iter((_oracle(), _oracle((42,), " Lyon")))
|
|
|
|
@asynccontextmanager
|
|
async def engine_context():
|
|
yield object()
|
|
|
|
async def next_canary(_engine: object) -> Oracle:
|
|
return next(canaries)
|
|
|
|
async def return_none(_engine: object) -> None:
|
|
return None
|
|
|
|
monkeypatch.setattr(
|
|
api_server, "build_async_engine_client", lambda _args: engine_context()
|
|
)
|
|
monkeypatch.setattr(snapshot_server, "run_engine_canary", next_canary)
|
|
monkeypatch.setattr(snapshot_server, "_release_reloadable_state", return_none)
|
|
monkeypatch.setattr(snapshot_server, "_restore_reloadable_state", return_none)
|
|
monkeypatch.setattr(snapshot_server, "detach_snapshot_streams", lambda: None)
|
|
monkeypatch.setattr(
|
|
snapshot_server,
|
|
"write_ready_atomic",
|
|
lambda *_args: pytest.fail("mismatched snapshot was published"),
|
|
)
|
|
control = snapshot_server.ControlArgs(
|
|
ready_file=tmp_path / "ready.json",
|
|
release_file=tmp_path / "release.json",
|
|
release_timeout_s=1,
|
|
)
|
|
with pytest.raises(SnapshotCanaryError, match="rehearsal changed"):
|
|
await snapshot_server.run_vllm_snapshot_child(control, argparse.Namespace())
|
|
|
|
|
|
def test_snapshot_release_marker_ignores_unknown_keys(tmp_path: Path):
|
|
marker = tmp_path / "release.json"
|
|
marker.write_text(json.dumps({"release": True, "port": 8000, "unknown": 1}))
|
|
|
|
listener = snapshot_server.read_release_marker(marker)
|
|
|
|
assert listener == snapshot_server.ListenerConfig(host=None, port=8000)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"payload",
|
|
[
|
|
["release"],
|
|
{"port": 8000},
|
|
{"release": "true", "port": 8000},
|
|
{"release": 1, "port": 8000},
|
|
{"release": True, "port": 0},
|
|
{"release": True, "port": 70000},
|
|
{"release": True, "port": True},
|
|
{"release": True, "port": "8000"},
|
|
{"release": True, "port": 8000, "host": 5},
|
|
],
|
|
)
|
|
def test_snapshot_release_marker_rejects_an_invalid_payload(
|
|
tmp_path: Path, payload: object
|
|
):
|
|
marker = tmp_path / "release.json"
|
|
marker.write_text(json.dumps(payload))
|
|
|
|
with pytest.raises(snapshot_server.SnapshotBarrierError, match="release marker"):
|
|
snapshot_server.read_release_marker(marker)
|
|
|
|
|
|
def test_snapshot_create_rolls_back_failed_dump(tmp_path: Path):
|
|
target = tmp_path / "snapshot"
|
|
tools = _fake_snapshot_tools()
|
|
tools.dump.side_effect = RuntimeError("dump failed")
|
|
with pytest.raises(RuntimeError, match="dump failed"):
|
|
create_snapshot(
|
|
argparse.Namespace(snapshot_dir=str(target), model_tag="Qwen/Qwen3-0.6B"),
|
|
tools=tools,
|
|
)
|
|
tools.abort_create.assert_called_once_with(100)
|
|
assert not target.exists()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"failure_kind", ["called-process", "timeout", "missing", "unreadable"]
|
|
)
|
|
def test_snapshot_create_preserves_bounded_criu_diagnostics_after_cleanup(
|
|
tmp_path: Path, failure_kind: str, caplog: pytest.LogCaptureFixture
|
|
):
|
|
target = tmp_path / "snapshot"
|
|
tools = _fake_snapshot_tools()
|
|
criu_tools = LocalSnapshotTools()
|
|
criu_tools.plugin_dir = tmp_path
|
|
original_run = criu_tools._run
|
|
criu_tools._privileged = lambda: [] # type: ignore[method-assign]
|
|
captured = "captured stderr " + "x" * 9000
|
|
if failure_kind == "timeout":
|
|
primary: subprocess.SubprocessError = subprocess.TimeoutExpired(
|
|
["criu", "dump"], 1, output=b"captured stdout\xff", stderr=captured.encode()
|
|
)
|
|
else:
|
|
primary = subprocess.CalledProcessError(
|
|
1, ["criu", "dump"], output="captured stdout", stderr=captured
|
|
)
|
|
|
|
def run(command: list[str], **_kwargs: object):
|
|
if "tail" not in command:
|
|
raise primary
|
|
if failure_kind == "unreadable":
|
|
raise PermissionError("log is root-owned")
|
|
if failure_kind == "timeout":
|
|
return subprocess.CompletedProcess(command, 0, "timeout log tail", "")
|
|
return original_run(command, **_kwargs)
|
|
|
|
criu_tools._run = run # type: ignore[method-assign]
|
|
|
|
def fail_dump(workdir: Path, _inventory: ProcessInventory) -> None:
|
|
images = workdir / "images"
|
|
images.mkdir()
|
|
if failure_kind != "missing":
|
|
(images / "dump.log").write_text(
|
|
"omitted dump prefix" + "d" * 9000 + "distinctive dump diagnostic"
|
|
)
|
|
criu_tools._criu("dump", workdir, [])
|
|
|
|
tools.dump.side_effect = fail_dump
|
|
with pytest.raises(type(primary)) as excinfo:
|
|
create_snapshot(
|
|
argparse.Namespace(snapshot_dir=str(target), model_tag="Qwen/Qwen3-0.6B"),
|
|
tools=tools,
|
|
)
|
|
|
|
assert excinfo.value is primary
|
|
assert not target.exists()
|
|
diagnostics = "\n".join(caplog.messages)
|
|
assert "captured stdout" in diagnostics
|
|
assert "captured stderr" not in diagnostics
|
|
assert "x" * 100 in diagnostics
|
|
assert len(diagnostics) < 17000
|
|
if failure_kind in {"missing", "unreadable"}:
|
|
assert "dump.log unavailable" in diagnostics
|
|
expected_error = (
|
|
"CalledProcessError" if failure_kind == "missing" else "PermissionError"
|
|
)
|
|
assert expected_error in diagnostics
|
|
elif failure_kind == "called-process":
|
|
assert "distinctive dump diagnostic" in diagnostics
|
|
assert "omitted dump prefix" not in diagnostics
|
|
else:
|
|
assert "timeout log tail" in diagnostics
|
|
if failure_kind == "timeout":
|
|
assert "captured stdout�" in diagnostics
|
|
|
|
|
|
def test_snapshot_criu_failure_selects_restore_log(
|
|
tmp_path: Path, caplog: pytest.LogCaptureFixture
|
|
):
|
|
tools = LocalSnapshotTools()
|
|
tools.plugin_dir = tmp_path
|
|
tools._privileged = lambda: [] # type: ignore[method-assign]
|
|
artifact = tmp_path / "snapshot"
|
|
images = artifact / "images"
|
|
images.mkdir(parents=True)
|
|
(images / "restore.log").write_text(
|
|
"omitted restore prefix" + "r" * 9000 + "distinctive restore diagnostic"
|
|
)
|
|
primary = subprocess.CalledProcessError(1, ["criu", "restore"])
|
|
original_run = tools._run
|
|
tailed: list[str] = []
|
|
|
|
def run(command: list[str], **_kwargs: object):
|
|
if "tail" not in command:
|
|
raise primary
|
|
tailed.append(command[-1])
|
|
return original_run(command, **_kwargs)
|
|
|
|
tools._run = run # type: ignore[method-assign]
|
|
|
|
with pytest.raises(subprocess.CalledProcessError) as excinfo:
|
|
tools._criu("restore", artifact, [])
|
|
|
|
assert excinfo.value is primary
|
|
assert tailed == [str(images / "restore.log")]
|
|
diagnostics = "\n".join(caplog.messages)
|
|
assert "distinctive restore diagnostic" in diagnostics
|
|
assert "omitted restore prefix" not in diagnostics
|
|
|
|
|
|
def _restore_fixture(tmp_path: Path) -> argparse.Namespace:
|
|
artifact = tmp_path / "snapshot"
|
|
artifact.mkdir(mode=0o700)
|
|
write_manifest_atomic(artifact, _manifest())
|
|
return argparse.Namespace(snapshot_dir=str(artifact), host="127.0.0.1", port=9000)
|
|
|
|
|
|
def test_snapshot_restore_rejects_identity_mismatch(tmp_path: Path):
|
|
args = _restore_fixture(tmp_path)
|
|
|
|
identity_mismatch = _fake_snapshot_tools()
|
|
identity_mismatch.current_identity.side_effect = lambda _gpu_uuid: (
|
|
_runtime_identity(vllm_version="different")
|
|
)
|
|
with pytest.raises(SnapshotCompatibilityError, match="vllm_version"):
|
|
restore_snapshot(args, tools=identity_mismatch)
|
|
identity_mismatch.restore.assert_not_called()
|
|
|
|
|
|
def test_snapshot_identity_mismatch_names_every_differing_field():
|
|
with pytest.raises(SnapshotCompatibilityError) as mismatch:
|
|
validate_identity(
|
|
_manifest(),
|
|
_runtime_identity(vllm_version="different", gpu_uuid="GPU-other"),
|
|
)
|
|
|
|
assert str(mismatch.value) == "snapshot mismatch: vllm_version, gpu_uuid"
|
|
|
|
|
|
def test_snapshot_restore_cleans_oracle_mismatch(tmp_path: Path):
|
|
args = _restore_fixture(tmp_path)
|
|
oracle_mismatch = _fake_snapshot_tools()
|
|
oracle_mismatch.request_oracle.return_value = _oracle(logprob=-0.5)
|
|
with pytest.raises(SnapshotRestoreError, match="oracle mismatch"):
|
|
restore_snapshot(args, tools=oracle_mismatch)
|
|
oracle_mismatch.cleanup.assert_called_once()
|
|
oracle_mismatch.complete_restore.assert_not_called()
|
|
assert oracles_match(_oracle(logprob=-100.0), _oracle(logprob=-100.0005))
|
|
assert not oracles_match(_oracle(logprob=-100.0), _oracle(logprob=-100.002))
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("invalid_update", "diagnostic"),
|
|
[
|
|
({"unexpected": True}, "unexpected"),
|
|
({"schema_version": True}, "schema_version"),
|
|
({"boundary": "unknown"}, "boundary"),
|
|
({"process_tree": [100, 100]}, "root"),
|
|
({"cuda_holders": [999]}, "root"),
|
|
({"oracle_token_ids": [1, 2]}, "oracle_token_ids"),
|
|
({"oracle_sampled_token_logprob": 0}, "oracle_sampled_token_logprob"),
|
|
(
|
|
{"oracle_sampled_token_logprob": float("nan")},
|
|
"oracle_sampled_token_logprob",
|
|
),
|
|
],
|
|
)
|
|
def test_snapshot_manifest_validation(
|
|
tmp_path: Path,
|
|
invalid_update: dict[str, object],
|
|
diagnostic: str,
|
|
):
|
|
artifact = tmp_path / "snapshot"
|
|
artifact.mkdir(mode=0o700)
|
|
write_manifest_atomic(artifact, _manifest())
|
|
invalid = _manifest().model_dump(mode="json")
|
|
invalid.update(invalid_update)
|
|
(artifact / "manifest.json").write_text(json.dumps(invalid))
|
|
with pytest.raises(
|
|
SnapshotCompatibilityError,
|
|
match=f"invalid snapshot manifest: {diagnostic}",
|
|
):
|
|
read_manifest(artifact)
|
|
|
|
|
|
def _local_restore_artifact(tmp_path: Path, name: str = "snapshot") -> Path:
|
|
artifact = tmp_path / name
|
|
artifact.mkdir(mode=0o700)
|
|
(artifact / "child.log").write_bytes(b"startup log")
|
|
(artifact / "child.log.snapshot-size").write_text("11\n")
|
|
return artifact
|
|
|
|
|
|
def _mark_snapshot_pids_free(monkeypatch: pytest.MonkeyPatch) -> list[int]:
|
|
probed_pids: list[int] = []
|
|
|
|
def missing(pid: int, _signal: int) -> None:
|
|
probed_pids.append(pid)
|
|
raise ProcessLookupError
|
|
|
|
monkeypatch.setattr(snapshot_runtime.os, "kill", missing)
|
|
return probed_pids
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("probe_error", "diagnostic"),
|
|
[
|
|
(None, "already occupied: 100"),
|
|
(PermissionError(), "already occupied: 100"),
|
|
(OSError("probe failed"), "availability probe failed: 100"),
|
|
],
|
|
)
|
|
def test_snapshot_restore_rejects_unavailable_pid_before_mutation(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
probe_error: OSError | None,
|
|
diagnostic: str,
|
|
):
|
|
args = _restore_fixture(tmp_path)
|
|
artifact = Path(args.snapshot_dir)
|
|
(artifact / "child.log").write_bytes(b"startup log")
|
|
(artifact / "release.json").write_bytes(b"release sentinel")
|
|
before = {path.name: path.read_bytes() for path in artifact.iterdir()}
|
|
tools = LocalSnapshotTools()
|
|
|
|
def fail_probe(_pid: int, _signal: int) -> None:
|
|
if probe_error is not None:
|
|
raise probe_error
|
|
|
|
monkeypatch.setattr(snapshot_runtime.os, "kill", fail_probe)
|
|
monkeypatch.setattr(tools, "preflight", lambda *_args: None)
|
|
monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity())
|
|
monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
|
|
|
|
with pytest.raises(SnapshotRestoreError, match=diagnostic):
|
|
restore_snapshot(args, tools=tools)
|
|
assert {path.name: path.read_bytes() for path in artifact.iterdir()} == before
|
|
|
|
|
|
@pytest.mark.parametrize("mode", ["pinned", "unpinned", "cleanup-fails"])
|
|
def test_snapshot_restore_terminates_only_a_pinned_tree(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
mode: str,
|
|
):
|
|
_mark_snapshot_pids_free(monkeypatch)
|
|
artifact = _local_restore_artifact(tmp_path)
|
|
manifest = _manifest()
|
|
root_pid = manifest.process_tree[0]
|
|
tools = LocalSnapshotTools()
|
|
terminated: list[int] = []
|
|
|
|
def terminate_and_wait(pid: int) -> None:
|
|
terminated.append(pid)
|
|
if mode == "cleanup-fails":
|
|
raise SnapshotRestoreError("restored process cleanup is incomplete")
|
|
tools._restored_processes.pop(pid)
|
|
|
|
def pin(
|
|
_artifact: Path, process_tree: tuple[int, ...], _holders: tuple[int, ...]
|
|
) -> None:
|
|
if mode == "unpinned":
|
|
raise SnapshotRestoreError(
|
|
"restored session does not match the captured process tree"
|
|
)
|
|
tools._restored_processes[process_tree[0]] = ()
|
|
|
|
monkeypatch.setattr(tools, "cleanup", terminate_and_wait)
|
|
monkeypatch.setattr(tools, "_pin_restored_tree", pin)
|
|
monkeypatch.setattr(tools, "_criu", lambda *_args: None)
|
|
monkeypatch.setattr(
|
|
tools,
|
|
"_run",
|
|
lambda command, **_kwargs: subprocess.CompletedProcess(command, 0, "333\n", ""),
|
|
)
|
|
|
|
with pytest.raises(SnapshotRestoreError, match="does not match") as failure:
|
|
tools.restore(artifact, manifest)
|
|
|
|
# A pinned tree is ours: it must be terminated and waited for, not left to
|
|
# honour the abort marker. An unpinned PID is never signalled.
|
|
assert terminated == ([] if mode == "unpinned" else [root_pid])
|
|
if mode == "cleanup-fails":
|
|
# The failure to terminate is reported beside the primary error, and a
|
|
# tree that may still be running stays pinned.
|
|
assert "restore cleanup failed" in str(failure.value)
|
|
assert tools._restored_processes == {root_pid: ()}
|
|
else:
|
|
assert tools._restored_processes == {}
|
|
assert json.loads((artifact / "release.json").read_text())["release"] is False
|
|
|
|
|
|
def test_snapshot_private_path_and_link_remap_security(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
_mark_snapshot_pids_free(monkeypatch)
|
|
real = tmp_path / "real"
|
|
real.mkdir(mode=0o700)
|
|
linked = tmp_path / "linked"
|
|
linked.symlink_to(real, target_is_directory=True)
|
|
with pytest.raises(SnapshotSecurityError, match="symlink"):
|
|
validate_artifact_root(linked, creating=False)
|
|
|
|
unsafe = tmp_path / "unsafe"
|
|
unsafe.mkdir()
|
|
unsafe.chmod(0o777)
|
|
with pytest.raises(SnapshotSecurityError, match="world-writable"):
|
|
validate_artifact_root(unsafe / "snapshot", creating=True)
|
|
|
|
artifact = _local_restore_artifact(tmp_path)
|
|
saved_remaps = artifact / "link-remaps"
|
|
saved_remaps.mkdir()
|
|
(saved_remaps / "link_remap.270").write_bytes(b"semaphore state")
|
|
shm_dir = tmp_path / "shm"
|
|
shm_dir.mkdir()
|
|
target = shm_dir / "link_remap.270"
|
|
replacement = tmp_path / "replacement"
|
|
replacement.write_bytes(b"semaphore state")
|
|
target.symlink_to(replacement)
|
|
tools = LocalSnapshotTools()
|
|
tools.shm_dir = shm_dir
|
|
monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
|
|
with pytest.raises(SnapshotRestoreError, match="conflicting CRIU link remap"):
|
|
tools.restore(artifact, _manifest())
|
|
|
|
|
|
@pytest.mark.parametrize("inventory_error", [None, PermissionError("denied")])
|
|
def test_snapshot_dump_rejects_unavailable_link_remap_inventory_before_criu(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
inventory_error: PermissionError | None,
|
|
):
|
|
tools = LocalSnapshotTools()
|
|
workdir = tmp_path / "snapshot"
|
|
workdir.mkdir()
|
|
tools.shm_dir = tmp_path / "missing-shm"
|
|
if inventory_error is not None:
|
|
tools.shm_dir.mkdir()
|
|
|
|
def fail_inventory(_path: Path):
|
|
raise inventory_error
|
|
|
|
monkeypatch.setattr(Path, "iterdir", fail_inventory)
|
|
monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
|
|
|
|
with pytest.raises(SnapshotCreateError, match="inventory CRIU link remaps"):
|
|
tools.dump(workdir, ProcessInventory(100, (100,), (100,), "GPU"))
|
|
|
|
|
|
def test_snapshot_restore_rejects_saved_remap_without_prefix(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
_mark_snapshot_pids_free(monkeypatch)
|
|
artifact = _local_restore_artifact(tmp_path)
|
|
saved = artifact / "link-remaps"
|
|
saved.mkdir()
|
|
(saved / "270").write_bytes(b"invalid remap name")
|
|
tools = LocalSnapshotTools()
|
|
tools.shm_dir = tmp_path / "shm"
|
|
tools.shm_dir.mkdir()
|
|
monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
|
|
|
|
with pytest.raises(SnapshotRestoreError, match="invalid saved link remap"):
|
|
tools.restore(artifact, _manifest())
|
|
|
|
assert not tuple(tools.shm_dir.iterdir())
|
|
|
|
|
|
def test_snapshot_link_remap_inventory_matches_only_numeric_suffixes(tmp_path: Path):
|
|
for name in (
|
|
"link_remap.0",
|
|
"link_remap.270",
|
|
"link_remap.",
|
|
"link_remap.12a",
|
|
"link-remap.12",
|
|
"prefix.link_remap.12",
|
|
"270",
|
|
):
|
|
(tmp_path / name).touch()
|
|
tools = LocalSnapshotTools()
|
|
tools.shm_dir = tmp_path
|
|
|
|
assert tools._link_remap_names() == {"link_remap.0", "link_remap.270"}
|
|
|
|
|
|
def test_snapshot_dump_preserves_primary_failure_when_cleanup_inventory_fails(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
tools = LocalSnapshotTools()
|
|
workdir = tmp_path / "snapshot"
|
|
workdir.mkdir()
|
|
inventory_calls = 0
|
|
|
|
def inventory_names() -> set[str]:
|
|
nonlocal inventory_calls
|
|
inventory_calls += 1
|
|
if inventory_calls > 1:
|
|
raise SnapshotCreateError(
|
|
"could not inventory CRIU link remaps: PermissionError: denied"
|
|
)
|
|
return set()
|
|
|
|
monkeypatch.setattr(tools, "_link_remap_names", inventory_names)
|
|
|
|
def fail_dump(*_args):
|
|
raise RuntimeError("dump failed")
|
|
|
|
monkeypatch.setattr(tools, "_criu", fail_dump)
|
|
|
|
with pytest.raises(SnapshotCreateError) as failure:
|
|
tools.dump(workdir, ProcessInventory(100, (100,), (100,), "GPU"))
|
|
|
|
assert "RuntimeError: dump failed" in str(failure.value)
|
|
assert "cleanup inventory failed" in str(failure.value)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("blocked_state", "message"),
|
|
[("tcp", "external established TCP"), ("io_uring", "kernel.io_uring_disabled=1")],
|
|
)
|
|
def test_snapshot_rejects_unsafe_process_state_before_criu(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, blocked_state: str, message: str
|
|
):
|
|
tools = LocalSnapshotTools()
|
|
monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101))
|
|
monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",))
|
|
monkeypatch.setattr(
|
|
tools,
|
|
"_descriptor_inventory",
|
|
lambda _tree: (
|
|
(101,) if blocked_state == "io_uring" else (),
|
|
{41: 101},
|
|
),
|
|
)
|
|
tcp_record = _TcpSocketRecord(
|
|
family="AF_INET",
|
|
local_raw="3C00000A:D6F6",
|
|
remote_raw="01010101:01BB" if blocked_state == "tcp" else "00000000:0000",
|
|
inode=41,
|
|
)
|
|
monkeypatch.setattr(tools, "_tcp_records", lambda: (tcp_record,))
|
|
with pytest.raises(SnapshotCreateError, match=message) as excinfo:
|
|
tools.inventory(100, tmp_path)
|
|
if blocked_state == "tcp":
|
|
assert "pid 101" in str(excinfo.value)
|
|
assert "1.1.1.1:443" in str(excinfo.value)
|
|
|
|
|
|
def test_snapshot_rejects_open_infiniband_device_before_criu(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
tools = LocalSnapshotTools()
|
|
monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101))
|
|
monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",))
|
|
monkeypatch.setattr(
|
|
tools,
|
|
"_descriptor_targets",
|
|
lambda pid: ("/dev/infiniband/uverbs7",) if pid == 101 else (),
|
|
)
|
|
monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
|
|
|
|
with pytest.raises(SnapshotCreateError) as failure:
|
|
tools.inventory(100, tmp_path)
|
|
|
|
message = str(failure.value)
|
|
assert "pid 101" in message
|
|
assert "/dev/infiniband/uverbs7" in message
|
|
assert "NCCL_IB_DISABLE=1" in message
|
|
|
|
|
|
def test_decode_endpoint_families():
|
|
assert _decode_endpoint("AF_INET", "3C00000A:D6F6") == "10.0.0.60:55030"
|
|
assert (
|
|
_decode_endpoint("AF_INET6", "00000000000000000000000001000000:01BB")
|
|
== "[::1]:443"
|
|
)
|
|
|
|
|
|
def test_tcp_records_skips_missing_tables(tmp_path, monkeypatch):
|
|
tcp, tcp6 = tmp_path / "tcp", tmp_path / "tcp6"
|
|
contents = "header\n0: 0100007F:1F90 00000000:0000 01 0 0 0 0 0 41\n"
|
|
tcp.write_text(contents)
|
|
tables = (("AF_INET", tcp), ("AF_INET6", tcp6))
|
|
monkeypatch.setattr(snapshot_runtime, "_TCP_TABLES", tables)
|
|
|
|
tools = LocalSnapshotTools()
|
|
assert tuple(record.family for record in tools._tcp_records()) == ("AF_INET",)
|
|
|
|
tcp6.write_text(contents)
|
|
families = tuple(record.family for record in tools._tcp_records())
|
|
assert families == ("AF_INET", "AF_INET6")
|
|
|
|
|
|
def test_snapshot_manifest_records_external_cache_files(tmp_path: Path):
|
|
artifact = tmp_path / "snapshot"
|
|
artifact.mkdir(mode=0o700)
|
|
recorded = (
|
|
(str(tmp_path / "kernel.so"), "b" * 64),
|
|
(str(tmp_path / "launcher.lock"), "locked"),
|
|
)
|
|
write_manifest_atomic(artifact, _manifest(external_cache_files=recorded))
|
|
|
|
assert read_manifest(artifact).external_cache_files == recorded
|
|
|
|
# A manifest written before this field existed still loads.
|
|
legacy = _manifest().model_dump(mode="json")
|
|
del legacy["external_cache_files"]
|
|
(artifact / "manifest.json").write_text(json.dumps(legacy))
|
|
|
|
assert read_manifest(artifact).external_cache_files == ()
|
|
|
|
|
|
def _pin_generated_cache_roots(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
|
"""Keep the resolver off whatever cache roots the host running this has."""
|
|
monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path / "vllm_cache"))
|
|
monkeypatch.setenv("TRITON_CACHE_DIR", str(tmp_path / "triton_cache_default"))
|
|
monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", str(tmp_path / "inductor_cache"))
|
|
|
|
|
|
def _unreadable(_path: Path) -> str:
|
|
raise OSError("cache file is unreadable")
|
|
|
|
|
|
def test_snapshot_inventory_records_open_generated_cache_files(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
cache_root = tmp_path / "triton_cache"
|
|
cache_root.mkdir()
|
|
kernel = cache_root / "kernel.so"
|
|
kernel.write_bytes(b"compiled kernel")
|
|
lock = cache_root / "launcher.lock"
|
|
lock.write_bytes(b"held by pid 100")
|
|
weights = tmp_path / "model.safetensors"
|
|
weights.write_bytes(b"weights")
|
|
artifact = tmp_path / "snapshot"
|
|
artifact.mkdir(mode=0o700)
|
|
(artifact / "child.log.lock").write_bytes(b"")
|
|
_pin_generated_cache_roots(monkeypatch, tmp_path)
|
|
monkeypatch.setenv("TRITON_CACHE_DIR", str(cache_root))
|
|
|
|
targets = {
|
|
100: (str(kernel), str(weights), "socket:[41]"),
|
|
101: (str(lock), str(cache_root), str(artifact / "child.log.lock")),
|
|
102: (f"{cache_root / 'evicted.so'} (deleted)",),
|
|
}
|
|
tools = LocalSnapshotTools()
|
|
monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101, 102))
|
|
monkeypatch.setattr(tools, "_descriptor_targets", lambda pid: targets[pid])
|
|
monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",))
|
|
monkeypatch.setattr(tools, "_tcp_records", lambda: ())
|
|
monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity())
|
|
|
|
manifest = tools.make_manifest(
|
|
argparse.Namespace(
|
|
model_tag="Qwen/Qwen3-0.6B",
|
|
revision=_MODEL_REVISION,
|
|
served_model_name=None,
|
|
tokenizer_revision=None,
|
|
),
|
|
("Qwen/Qwen3-0.6B",),
|
|
tools.inventory(100, artifact),
|
|
_oracle(),
|
|
artifact,
|
|
)
|
|
|
|
# Weights, sockets, the cache directory itself, files inside the artifact,
|
|
# and deleted targets that CRIU remaps are all left out.
|
|
assert manifest.external_cache_files == (
|
|
(str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()),
|
|
(str(lock), "locked"),
|
|
)
|
|
|
|
|
|
def test_snapshot_inventory_marks_a_cache_file_it_cannot_read(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
cache_root = tmp_path / "triton_cache"
|
|
cache_root.mkdir()
|
|
kernel = cache_root / "kernel.so"
|
|
kernel.write_bytes(b"compiled kernel")
|
|
lock = cache_root / "launcher.lock"
|
|
lock.write_bytes(b"held by pid 100")
|
|
artifact = tmp_path / "snapshot"
|
|
artifact.mkdir(mode=0o700)
|
|
_pin_generated_cache_roots(monkeypatch, tmp_path)
|
|
monkeypatch.setenv("TRITON_CACHE_DIR", str(cache_root))
|
|
|
|
tools = LocalSnapshotTools()
|
|
monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100,))
|
|
monkeypatch.setattr(
|
|
tools, "_descriptor_targets", lambda _pid: (str(kernel), str(lock))
|
|
)
|
|
monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("100, GPU-abc",))
|
|
monkeypatch.setattr(tools, "_tcp_records", lambda: ())
|
|
monkeypatch.setattr(tools, "_sha256", _unreadable)
|
|
|
|
# A file that cannot be digested keeps its own marker, so the manifest says
|
|
# which entries restore can check for existence only.
|
|
assert tools.inventory(100, artifact).external_cache_files == (
|
|
(str(kernel), "unreadable"),
|
|
(str(lock), "locked"),
|
|
)
|
|
|
|
|
|
def test_snapshot_inventory_matches_resolved_cache_roots_not_siblings(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
# Descriptor targets come back kernel-canonical, so the fixture must be too.
|
|
base = Path(os.path.realpath(tmp_path))
|
|
real_cache = base / "mnt_cache"
|
|
real_cache.mkdir()
|
|
kernel = real_cache / "kernel.so"
|
|
kernel.write_bytes(b"compiled kernel")
|
|
linked_cache = base / "linked_cache"
|
|
linked_cache.symlink_to(real_cache, target_is_directory=True)
|
|
sibling = base / "vllm_cache-backup"
|
|
sibling.mkdir()
|
|
stale = sibling / "stale.so"
|
|
stale.write_bytes(b"stale kernel")
|
|
artifact = base / "snapshot"
|
|
artifact.mkdir(mode=0o700)
|
|
|
|
_pin_generated_cache_roots(monkeypatch, base)
|
|
monkeypatch.setenv("TRITON_CACHE_DIR", str(linked_cache))
|
|
|
|
tools = LocalSnapshotTools()
|
|
monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100,))
|
|
monkeypatch.setattr(
|
|
tools, "_descriptor_targets", lambda _pid: (str(kernel), str(stale))
|
|
)
|
|
monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("100, GPU-abc",))
|
|
monkeypatch.setattr(tools, "_tcp_records", lambda: ())
|
|
|
|
# A cache root reached through a symlink still covers the file the kernel
|
|
# reports by its real path, and a directory that merely shares the root's
|
|
# name prefix is not part of the cache at all.
|
|
assert tools.inventory(100, artifact).external_cache_files == (
|
|
(str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()),
|
|
)
|
|
|
|
|
|
def test_snapshot_restore_rejects_rotated_cache_file_before_mutation(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
_mark_snapshot_pids_free(monkeypatch)
|
|
kernel = tmp_path / "kernel.so"
|
|
kernel.write_bytes(b"compiled kernel")
|
|
lock = tmp_path / "launcher.lock"
|
|
lock.write_bytes(b"held by pid 100")
|
|
manifest = _manifest(
|
|
external_cache_files=(
|
|
(str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()),
|
|
(str(lock), "locked"),
|
|
)
|
|
)
|
|
artifact = _local_restore_artifact(tmp_path)
|
|
before = {path.name: path.read_bytes() for path in artifact.iterdir()}
|
|
tools = LocalSnapshotTools()
|
|
monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
|
|
|
|
kernel.write_bytes(b"recompiled kernel")
|
|
with pytest.raises(SnapshotRestoreError, match="changed since capture") as changed:
|
|
tools.restore(artifact, manifest)
|
|
assert str(kernel) in str(changed.value)
|
|
|
|
# Restoring the digest gets past the first entry, so the lock file is only
|
|
# checked for existence: the second failure proves the first one passed.
|
|
kernel.write_bytes(b"compiled kernel")
|
|
lock.unlink()
|
|
with pytest.raises(SnapshotRestoreError, match="vanished since capture") as gone:
|
|
tools.restore(artifact, manifest)
|
|
assert str(lock) in str(gone.value)
|
|
|
|
assert {path.name: path.read_bytes() for path in artifact.iterdir()} == before
|
|
|
|
|
|
def test_snapshot_restore_rejects_a_cache_file_it_cannot_read(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
_mark_snapshot_pids_free(monkeypatch)
|
|
kernel = tmp_path / "kernel.so"
|
|
kernel.write_bytes(b"compiled kernel")
|
|
manifest = _manifest(
|
|
external_cache_files=(
|
|
(str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()),
|
|
)
|
|
)
|
|
artifact = _local_restore_artifact(tmp_path)
|
|
tools = LocalSnapshotTools()
|
|
monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
|
|
monkeypatch.setattr(tools, "_sha256", _unreadable)
|
|
|
|
# A recorded digest that cannot be read again is a failed check, not a file
|
|
# to wave through on existence.
|
|
with pytest.raises(
|
|
SnapshotRestoreError, match="unreadable since capture"
|
|
) as blocked:
|
|
tools.restore(artifact, manifest)
|
|
assert str(kernel) in str(blocked.value)
|
|
|
|
|
|
def test_snapshot_runtime_installer_logs_the_failing_exit_code():
|
|
installer = Path(__file__).parents[3] / "tools" / "install_snapshot_runtime.sh"
|
|
if not installer.is_file():
|
|
pytest.skip("installer script is not shipped in this test image")
|
|
timing = [
|
|
line
|
|
for line in installer.read_text().splitlines()
|
|
if line.startswith(("_snapshot_install_started=", "trap "))
|
|
]
|
|
assert len(timing) == 2
|
|
|
|
result = subprocess.run(
|
|
["bash", "-c", "\n".join([*timing, "exit 7"])],
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
|
|
assert result.returncode == 7
|
|
assert "(exit 7)" in result.stdout
|