From 13058faa49f5e4a4d50e1ca562ba0c2fdd73c7bc Mon Sep 17 00:00:00 2001 From: Nils Matteson Date: Mon, 21 Sep 2026 16:18:49 -0700 Subject: [PATCH] [Frontend] Add reusable TP1 initialized-engine snapshots (#51360) Signed-off-by: Nils Matteson Signed-off-by: Taneem Ibrahim Signed-off-by: Andreas Karatzas Co-authored-by: Taneem Ibrahim Co-authored-by: elehayym <52448798+Yuzu23@users.noreply.github.com> Co-authored-by: Andreas Karatzas Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../scripts/initialized-snapshot-e2e.sh | 311 ++++ .buildkite/test_areas/entrypoints.yaml | 20 + docker/Dockerfile | 16 + docs/features/initialized_snapshots.md | 152 ++ .../entrypoints/launchers/test_launch_cli.py | 1043 +++++++++++++- .../quantization/test_mimo_v2_qkv_shard.py | 8 + tools/install_snapshot_runtime.sh | 109 ++ vllm/entrypoints/cli/main.py | 12 +- vllm/entrypoints/cli/snapshot.py | 143 ++ vllm/snapshot/__init__.py | 4 + vllm/snapshot/controller.py | 132 ++ vllm/snapshot/manifest.py | 277 ++++ vllm/snapshot/runtime.py | 1246 +++++++++++++++++ vllm/snapshot/server.py | 252 ++++ vllm/snapshot/types.py | 37 + vllm/utils/argparse_utils.py | 21 +- 16 files changed, 3773 insertions(+), 10 deletions(-) create mode 100755 .buildkite/scripts/initialized-snapshot-e2e.sh create mode 100644 docs/features/initialized_snapshots.md create mode 100755 tools/install_snapshot_runtime.sh create mode 100644 vllm/entrypoints/cli/snapshot.py create mode 100644 vllm/snapshot/__init__.py create mode 100644 vllm/snapshot/controller.py create mode 100644 vllm/snapshot/manifest.py create mode 100644 vllm/snapshot/runtime.py create mode 100644 vllm/snapshot/server.py create mode 100644 vllm/snapshot/types.py diff --git a/.buildkite/scripts/initialized-snapshot-e2e.sh b/.buildkite/scripts/initialized-snapshot-e2e.sh new file mode 100755 index 0000000000..cd0511ffdd --- /dev/null +++ b/.buildkite/scripts/initialized-snapshot-e2e.sh @@ -0,0 +1,311 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +set -Eeuo pipefail +readonly MODEL="Qwen/Qwen3-0.6B" +readonly MODEL_REVISION="c1899de289a04d12100db370d81485cdf75e47ca" +readonly EXPECTED_TOKEN_ID="12095" +readonly EXPECTED_TEXT=" Paris" +readonly CLEANUP_RESERVE_S=300 +RUN_LIMIT_S="${SNAPSHOT_E2E_RUN_LIMIT_S:-3300}" +[[ "$RUN_LIMIT_S" =~ ^[1-9][0-9]{2,3}$ ]] \ + || { echo "SNAPSHOT_E2E_RUN_LIMIT_S must be a bounded decimal integer" >&2; exit 1; } +(( RUN_LIMIT_S >= CLEANUP_RESERVE_S + 60 && RUN_LIMIT_S <= 3300 )) \ + || { echo "SNAPSHOT_E2E_RUN_LIMIT_S must be between 360 and 3300" >&2; exit 1; } +readonly RUN_LIMIT_S +readonly HOST_HF_CACHE="${SNAPSHOT_E2E_HOST_HF_CACHE:-}" +HF_CACHE_MOUNT=() +if [[ -n "$HOST_HF_CACHE" ]]; then + [[ "$HOST_HF_CACHE" == /* && -d "$HOST_HF_CACHE" ]] \ + || { echo "SNAPSHOT_E2E_HOST_HF_CACHE must be an existing absolute directory" >&2; exit 1; } + HF_CACHE_MOUNT=(--mount "type=bind,source=$HOST_HF_CACHE,target=/e2e/hf,readonly") + PREPARE_ARTIFACT_ROOT='install -d -m 0700 /e2e' +else + PREPARE_ARTIFACT_ROOT='mkdir -m 0700 /e2e && mkdir -m 0700 /e2e/hf' +fi +readonly PREPARE_ARTIFACT_ROOT +if (( $# != 2 )); then + echo "usage: $0 IMAGE_TAG BUILDKITE_COMMIT" >&2 + exit 2 +fi +readonly IMAGE_TAG="$1" +readonly EXPECTED_COMMIT="$2" +readonly DEADLINE_EPOCH="$(( $(date +%s) + RUN_LIMIT_S ))" +RUN_ID="" CONTAINER_NAME="" CONTAINER_ID="" +GPU_UUID="" PORT_ONE="" PORT_TWO="" +LINK_REMAP_BASELINE="" +SNAPSHOT_PIDS=() +die() { + echo "initialized snapshot E2E: $*" >&2 + exit 1 +} + +run() { + local label="$1" maximum_s="$2" available limit + shift 2 + available="$((DEADLINE_EPOCH - $(date +%s) - CLEANUP_RESERVE_S))" + (( available > 0 )) || die "no time remains before the cleanup reserve" + (( available < maximum_s )) && limit="$available" || limit="$maximum_s" + echo "--- ${label} (timeout ${limit}s)" >&2 + timeout --signal=TERM --kill-after=30s "${limit}s" "$@" +} + +link_remaps() { + find /dev/shm -maxdepth 1 -name 'link_remap.*' \ + -printf '%p|%D|%i|%m|%s\n' | sort +} + +gpu_pids() { + nvidia-smi --id="$GPU_UUID" --query-compute-apps=pid \ + --format=csv,noheader,nounits | sed '/^[[:space:]]*$/d' +} + +port_listeners() { + ss -H -ltn "sport = :$1" +} + +state_clean() { + local pid port output + for pid in "${SNAPSHOT_PIDS[@]}"; do + [[ ! -e "/proc/$pid" ]] || return 1 + done + output="$(gpu_pids)" || return 1 + [[ -z "$output" ]] || return 1 + output="$(link_remaps)" || return 1 + [[ "$output" == "$LINK_REMAP_BASELINE" ]] || return 1 + for port in "$PORT_ONE" "$PORT_TWO"; do + [[ -n "$port" ]] || continue + output="$(port_listeners "$port")" || return 1 + [[ -z "$output" ]] || return 1 + done +} + +wait_clean() { + local _attempt + for ((_attempt = 0; _attempt < 60; _attempt++)); do + state_clean && return 0 + sleep 1 + done + return 1 +} + +container_names() { + timeout 15s docker container ls --all --format '{{.Names}}' +} + +cleanup() { + local failed=0 names="" actual="" expected="" + if [[ -n "$CONTAINER_NAME" ]]; then + names="$(container_names 2>/dev/null)" || return 1 + if grep -Fxq "$CONTAINER_NAME" <<< "$names"; then + actual="$(docker container inspect \ + --format '{{.Id}}|{{index .Config.Labels "ai.vllm.snapshot.e2e.run"}}' \ + "$CONTAINER_NAME" 2>/dev/null)" || return 1 + expected="${CONTAINER_ID:-${actual%%|*}}|$RUN_ID" + if [[ "$actual" != "$expected" ]]; then + echo "refusing to remove a container with mismatched identity" >&2 + failed=1 + else + timeout 60s docker rm --force "$CONTAINER_NAME" >/dev/null \ + || failed=1 + fi + fi + names="$(container_names 2>/dev/null)" || return 1 + grep -Fxq "$CONTAINER_NAME" <<< "$names" && failed=1 + fi + if [[ -n "$PORT_TWO" ]] && ! wait_clean; then + echo "process, GPU, port, or link-remap residue remained" >&2 + failed=1 + fi + return "$failed" +} + +on_exit() { + local status=$? + trap - EXIT INT TERM + set +e + cleanup || status=1 + exit "$status" +} + +trap on_exit EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +parse_inspect() { + timeout 30s python3 -c ' +import json, math, sys +data = json.load(sys.stdin) +model, revision, gpu_uuid, token_id, text = sys.argv[1:] +expected = { + "boundary": "post-engine-init-reloadable-state-released", + "model": model, "model_revision": revision, + "gpu_uuid": gpu_uuid, "oracle_token_ids": [int(token_id)], + "oracle_text": text, +} +bad = next((field for field, value in expected.items() if data.get(field) != value), None) +if bad: + raise SystemExit(f"inspect field {bad} mismatch: {data.get(bad)!r}") +logprob = data.get("oracle_sampled_token_logprob") +if type(logprob) is not float or not math.isfinite(logprob): + raise SystemExit("inspect logprob is not finite") +pids = data.get("process_tree") +if not pids: + raise SystemExit("inspect process_tree is empty") +for pid in pids: + print(f"pid={pid}") +print(f"oracle_logprob={logprob}") +' "$MODEL" "$MODEL_REVISION" "$GPU_UUID" "$EXPECTED_TOKEN_ID" \ + "$EXPECTED_TEXT" +} + +assert_oracle() { + local port="$1" ordinal="$2" response + response="$(run "restore ${ordinal} public completion" 120 curl \ + --fail --silent --show-error --max-time 120 \ + --header 'Content-Type: application/json' \ + --data "{\"model\":\"${MODEL}\",\"prompt\":\"The capital of France is\",\"temperature\":0,\"seed\":0,\"min_tokens\":1,\"max_tokens\":1,\"logprobs\":0,\"return_token_ids\":true}" \ + "http://127.0.0.1:${port}/v1/completions")" + printf '%s' "$response" | timeout 30s python3 -c ' +import json, math, sys +choice = json.load(sys.stdin)["choices"][0] +token_id, text, expected, ordinal = sys.argv[1:] +values = choice["logprobs"]["token_logprobs"] +actual = float(values[0]) if len(values) == 1 else math.nan +if (choice["token_ids"] != [int(token_id)] or choice["text"] != text + or not math.isfinite(actual) or abs(actual - float(expected)) > 1e-3): + raise SystemExit(f"public oracle mismatch: {choice!r}") +print(f"restore={ordinal} token_id={token_id} text={text!r} logprob={actual}") +' "$EXPECTED_TOKEN_ID" "$EXPECTED_TEXT" "$INSPECT_LOGPROB" "$ordinal" +} + +[[ "$(uname -s)" == "Linux" && "$(uname -m)" == "x86_64" ]] \ + || die "host must be Linux x86_64" +[[ "$EXPECTED_COMMIT" =~ ^[0-9a-f]{40}$ ]] \ + || die "BUILDKITE_COMMIT must be a full lowercase commit hash" +[[ "$(< /proc/sys/kernel/io_uring_disabled)" == "2" ]] \ + || die "kernel.io_uring_disabled must already equal 2" +run "Docker preflight" 60 docker info >/dev/null + +VISIBLE_GPU="${NVIDIA_VISIBLE_DEVICES:-}" +if [[ -z "$VISIBLE_GPU" || "$VISIBLE_GPU" == "all" || "$VISIBLE_GPU" == "void" ]]; then + VISIBLE_GPU="${CUDA_VISIBLE_DEVICES:-}" +fi +if [[ -z "$VISIBLE_GPU" || "$VISIBLE_GPU" == "all" || "$VISIBLE_GPU" == "void" ]]; then + CDI_GPU="${BUILDKITE_PLUGIN_DOCKER_DEVICES_0:-}" + if [[ "$CDI_GPU" == nvidia.com/gpu=* ]]; then + VISIBLE_GPU="${CDI_GPU#nvidia.com/gpu=}" + else + VISIBLE_GPU="" + fi +fi +[[ -n "$VISIBLE_GPU" && "$VISIBLE_GPU" != all && "$VISIBLE_GPU" != void \ + && "$VISIBLE_GPU" != *","* ]] \ + || die "the runner must expose exactly one selected GPU" +GPU_ROW="$(nvidia-smi --id="$VISIBLE_GPU" \ + --query-gpu=name,uuid,mig.mode.current,memory.total \ + --format=csv,noheader,nounits)" +IFS=',' read -r GPU_NAME GPU_UUID MIG_MODE GPU_MEMORY_MIB <<< "$GPU_ROW" +for variable in GPU_NAME GPU_UUID MIG_MODE GPU_MEMORY_MIB; do + printf -v "$variable" '%s' "${!variable#"${!variable%%[![:space:]]*}"}" +done +[[ "$GPU_NAME" == *H200* && "$GPU_UUID" == GPU-* && "$MIG_MODE" == Disabled ]] \ + || die "selected device is not a full non-MIG H200: $GPU_ROW" +(( GPU_MEMORY_MIB >= 130000 )) || die "selected H200 does not expose full memory" +GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU" +[[ -z "$GPU_PIDS" ]] || die "selected H200 is not idle" + +RUN_ID="snapshot-e2e-$(date +%s)-$$-$RANDOM" +CONTAINER_NAME="vllm-$RUN_ID" +read -r PORT_ONE PORT_TWO < <(python3 -c 'import socket; s=[socket.socket(),socket.socket()]; [x.bind(("127.0.0.1",0)) for x in s]; print(*(x.getsockname()[1] for x in s))') +LINK_REMAP_BASELINE="$(link_remaps)" + +run "pull exact candidate image" 900 docker pull --platform linux/amd64 \ + "$IMAGE_TAG" >/dev/null +IFS='|' read -r IMAGE_ID IMAGE_PLATFORM OCI_COMMIT VLLM_COMMIT < <( + docker image inspect \ + --format '{{.Id}}|{{.Os}}/{{.Architecture}}|{{index .Config.Labels "org.opencontainers.image.revision"}}|{{index .Config.Labels "ai.vllm.build.commit"}}' \ + "$IMAGE_TAG" +) +[[ "$IMAGE_ID" =~ ^sha256:[0-9a-f]{64}$ && "$IMAGE_PLATFORM" == linux/amd64 ]] \ + || die "candidate image identity is invalid" +[[ "$OCI_COMMIT" == "$EXPECTED_COMMIT" && "$VLLM_COMMIT" == "$EXPECTED_COMMIT" ]] \ + || die "candidate image commit labels do not match" +CONTAINER_ID="$(run "start exact candidate container" 90 docker run --detach \ + --name "$CONTAINER_NAME" --label "ai.vllm.snapshot.e2e.run=$RUN_ID" \ + --gpus "device=$GPU_UUID" --user 0 --privileged --pid=host --ipc=host \ + --network=host --env "CUDA_VISIBLE_DEVICES=$GPU_UUID" --env HF_HUB_DISABLE_TELEMETRY=1 \ + --env VLLM_NO_USAGE_STATS=1 --env VLLM_USE_V2_MODEL_RUNNER=1 \ + ${HF_CACHE_MOUNT[@]+"${HF_CACHE_MOUNT[@]}"} \ + --entrypoint sleep "$IMAGE_ID" infinity)" +[[ "$CONTAINER_ID" =~ ^[0-9a-f]{64}$ ]] || die "invalid container ID" +run "prepare private artifact root" 60 docker exec "$CONTAINER_NAME" sh -c \ + "$PREPARE_ARTIFACT_ROOT && df -Pk / | awk 'NR == 2 && \$4 >= 10485760 {ok=1} END {exit !ok}'" +run "verify snapshot runtime" 60 docker exec "$CONTAINER_NAME" sh -c \ + 'criu --version; test -x /usr/local/sbin/cuda-checkpoint; test -f /usr/local/lib/criu/cuda_plugin.so' +run "verify selected CUDA device" 60 docker exec "$CONTAINER_NAME" python3 -c ' +import sys, torch +if torch.cuda.device_count() != 1: + raise SystemExit("expected exactly one visible CUDA device") +actual = str(torch.cuda.get_device_properties(0).uuid).removeprefix("GPU-") +if actual != sys.argv[1].removeprefix("GPU-"): + raise SystemExit(f"CUDA device {actual} does not match {sys.argv[1]}") +print(f"selected_cuda_uuid=GPU-{actual}") +' "$GPU_UUID" +OFFLINE_EXEC=(docker exec --env HF_HOME=/e2e/hf --env HF_HUB_OFFLINE=1 \ + --env TRANSFORMERS_OFFLINE=1 --env VLLM_SNAPSHOT_TIMEOUT_S=900 \ + "$CONTAINER_NAME") +if [[ -n "$HOST_HF_CACHE" ]]; then + run "validate pinned model in read-only host cache" 120 \ + "${OFFLINE_EXEC[@]}" python3 -c \ + 'from pathlib import Path; from huggingface_hub import snapshot_download; import sys; p=snapshot_download(repo_id=sys.argv[1], revision=sys.argv[2], local_files_only=True); assert Path(p).name == sys.argv[2]' \ + "$MODEL" "$MODEL_REVISION" +else + run "prefetch pinned public model" 900 docker exec --env HF_HOME=/e2e/hf \ + "$CONTAINER_NAME" python3 -c \ + 'from pathlib import Path; from huggingface_hub import snapshot_download; import sys; p=snapshot_download(repo_id=sys.argv[1], revision=sys.argv[2]); assert Path(p).name == sys.argv[2]' \ + "$MODEL" "$MODEL_REVISION" +fi + +run "create compact initialized snapshot" 1200 "${OFFLINE_EXEC[@]}" \ + vllm snapshot create "$MODEL" --revision "$MODEL_REVISION" \ + --tokenizer-revision "$MODEL_REVISION" --snapshot-dir /e2e/artifact \ + --dtype float16 --max-model-len 512 --gpu-memory-utilization 0.50 +wait_clean || die "create left process, GPU, port, or link-remap residue" + +INSPECT_JSON="$(run "inspect compact initialized snapshot" 120 \ + "${OFFLINE_EXEC[@]}" vllm snapshot inspect /e2e/artifact)" +INSPECT_RECEIPT="$(printf '%s' "$INSPECT_JSON" | parse_inspect)" +mapfile -t SNAPSHOT_PIDS < <(sed -n 's/^pid=//p' <<< "$INSPECT_RECEIPT") +INSPECT_LOGPROB="$(sed -n 's/^oracle_logprob=//p' <<< "$INSPECT_RECEIPT")" +(( ${#SNAPSHOT_PIDS[@]} > 0 )) || die "inspect returned no PIDs" +state_clean || die "create left recorded state alive" +echo "$INSPECT_RECEIPT" + +run "restore compact snapshot 1" 900 "${OFFLINE_EXEC[@]}" \ + vllm snapshot restore /e2e/artifact --host 127.0.0.1 --port "$PORT_ONE" +PORT_LISTENERS="$(port_listeners "$PORT_ONE")" || die "failed to inspect restore 1 port" +[[ -n "$PORT_LISTENERS" ]] || die "restore 1 did not bind its port" +GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU" +[[ -n "$GPU_PIDS" ]] || die "restore 1 did not engage the selected GPU" +assert_oracle "$PORT_ONE" 1 +run "stop container after restore 1" 60 docker stop --time 20 "$CONTAINER_NAME" >/dev/null +[[ "$(docker container inspect --format '{{.Id}}' "$CONTAINER_NAME")" == "$CONTAINER_ID" ]] \ + || die "docker stop replaced the container" +wait_clean || die "restore 1 teardown left residue" + +run "start the same container" 60 docker start "$CONTAINER_NAME" >/dev/null +[[ "$(docker container inspect --format '{{.Id}}' "$CONTAINER_NAME")" == "$CONTAINER_ID" ]] \ + || die "docker start replaced the container" +run "restore compact snapshot 2" 900 "${OFFLINE_EXEC[@]}" \ + vllm snapshot restore /e2e/artifact --host 127.0.0.1 --port "$PORT_TWO" +PORT_LISTENERS="$(port_listeners "$PORT_TWO")" || die "failed to inspect restore 2 port" +[[ -n "$PORT_LISTENERS" ]] || die "restore 2 did not bind its port" +GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU" +[[ -n "$GPU_PIDS" ]] || die "restore 2 did not engage the selected GPU" +assert_oracle "$PORT_TWO" 2 +run "stop container after restore 2" 60 docker stop --time 20 "$CONTAINER_NAME" >/dev/null +wait_clean || die "restore 2 teardown left residue" +cleanup || die "exact cleanup failed" +echo "container_id=$CONTAINER_ID restore_1_port=$PORT_ONE restore_2_port=$PORT_TWO" +echo "result=PASS" diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 60ac5d932c..9e0a8399a3 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -2,6 +2,26 @@ group: Entrypoints depends_on: - image-build steps: +- label: ":nvidia: (H200) Initialized Snapshot E2E" + key: initialized-snapshot-e2e + device: h200 + num_devices: 1 + no_plugin: true + optional: true + timeout_in_minutes: 60 + working_dir: "." + env: + SNAPSHOT_E2E_RUN_LIMIT_S: "1200" + source_file_dependencies: + - .buildkite/scripts/initialized-snapshot-e2e.sh + - .buildkite/test_areas/entrypoints.yaml + - docker/Dockerfile + - tools/install_snapshot_runtime.sh + - vllm/entrypoints/cli/ + - vllm/snapshot/ + commands: + - bash .buildkite/scripts/initialized-snapshot-e2e.sh "$IMAGE_TAG" "$BUILDKITE_COMMIT" + - label: ":nvidia: (H200 MIG 35GB) Entrypoints Unit" device: h200_35gb key: entrypoints-unit-tests diff --git a/docker/Dockerfile b/docker/Dockerfile index ad010285b1..5fd25f5388 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -735,6 +735,8 @@ RUN --mount=type=cache,target=/opt/uv/cache \ FROM ${FINAL_BASE_IMAGE} AS vllm-runtime-base ARG INSTALL_RUBIN_PRERELEASE +ARG TARGETPLATFORM +ARG max_jobs=2 ARG CUDA_VERSION ARG PYTHON_VERSION ARG NCCL_VERSION @@ -854,6 +856,20 @@ RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \ && ldconfig; \ fi +# Install the initialized-snapshot runtime on the architecture supported by +# `vllm snapshot`. CRIU and its CUDA plugin are built from the same pinned +# source. The `cuda-checkpoint`-compatible helper is built from NVIDIA Dynamo's +# Apache-2.0 source against the public CUDA Driver checkpoint APIs instead of +# redistributing NVIDIA's proprietary cuda-checkpoint binary. +COPY tools/install_snapshot_runtime.sh /usr/local/bin/install_snapshot_runtime.sh +RUN if [ "${TARGETPLATFORM}" = "linux/amd64" ]; then \ + MAX_JOBS="${max_jobs}" /usr/local/bin/install_snapshot_runtime.sh; \ + else \ + echo "Initialized snapshots are not installed for ${TARGETPLATFORM}"; \ + fi + +ENV CRIU_CUDA_PLUGIN_DIR=/usr/local/lib/criu + # Install uv for faster pip installs RUN python3 -m pip install uv diff --git a/docs/features/initialized_snapshots.md b/docs/features/initialized_snapshots.md new file mode 100644 index 0000000000..538695d0a8 --- /dev/null +++ b/docs/features/initialized_snapshots.md @@ -0,0 +1,152 @@ +# Initialized engine snapshots + +Initialized engine snapshots are an experimental way to trade local disk space +and host privileges for a faster vLLM activation. Snapshot creation initializes +the engine and records deterministic generation output before CRIU captures the +process tree and CUDA state. Restore validates the saved environment, restores +the engine, binds the HTTP server, and reproduces the recorded token and sampled +token log probability before returning. + +This path is intended for repeatedly activating the same model and engine +configuration on the same machine. It is not a portable model artifact. + +## Requirements + +Snapshots currently require: + +- Linux on x86-64 with one NVIDIA GPU. +- TP1 with one unauthenticated plaintext HTTP server. Other parallel sizes, + TLS, middleware, Unix sockets, and speculative decoding are unsupported. +- Because snapshot mode requires an unauthenticated plaintext HTTP server, use + a trusted host or network boundary, or external controls. See the vLLM + [Security guide](../usage/security.md). +- [CRIU](https://github.com/checkpoint-restore/criu), its CUDA plugin, a + `cuda-checkpoint`-compatible helper, and `nvidia-smi` on `PATH`. +- Root or passwordless `sudo` for CRIU. +- `io_uring` disabled before launch because CRIU cannot dump it. Use + `kernel.io_uring_disabled=1` for an unprivileged process. A process running + as root, including the `docker exec` flow below, bypasses `=1`, so use `=2` + host-wide. +- No established TCP connection to a peer outside the captured process tree. + Current Hugging Face hub clients hold their connections for the process + lifetime, so download the model in a separate step and run create with + `HF_HUB_OFFLINE=1`, as the quickstart below does. +- A remote model ID and an immutable 40-character `--revision`. Local model + directories and mutable revisions are not supported. +- Enough disk for the artifact, with the same installed vLLM package, model + files, container filesystem, and generated-cache paths available at restore. + Generated-cache files are not copied into the artifact. The manifest + fingerprints the ones the captured tree holds open, so restore fails early + and names the file when one is removed or replaced. Only open descriptors + are recorded. A library the tree mapped and then closed is not fingerprinted, + and CRIU still reopens it by path, so replacing that file permanently + invalidates the artifact without an early error. + +The official `vllm/vllm-openai` Linux x86-64 image includes the snapshot +runtime. It still requires a compatible host driver, kernel, and privileges. +Arm64 images omit it. Source installs must set `CRIU_CUDA_PLUGIN_DIR` to the +directory containing `cuda_plugin.so`. + +Run snapshot commands with `docker exec` inside a long-lived container. Restore +hands the API server off as a detached process, so a one-shot container would +stop that server when its PID 1 exits. Snapshot preflight requires every +component of the artifact path to be owned by the invoking user or root, with +the directory itself at mode 0700, and the commands below run as root inside +the official image, so the bind-mounted host directory is created with `sudo` +and root ownership. The model downloads in its own step so the captured tree +holds no hub connection, and create runs offline. This example also keeps the +container filesystem and `/dev/shm` namespace stable for the lifetime of the +artifact: + +```bash +sudo sysctl kernel.io_uring_disabled=2 + +snapshot_root="$(pwd)/vllm-snapshots" +sudo install -d -m 0700 -o root -g root "${snapshot_root}" + +docker run --detach --name vllm-snapshot \ + --gpus all \ + --privileged \ + --pid=host \ + --ipc=host \ + --network=host \ + --mount "type=bind,source=${snapshot_root},target=/snapshots" \ + --entrypoint sleep \ + vllm/vllm-openai:latest infinity + +docker exec vllm-snapshot hf download Qwen/Qwen3-0.6B \ + --revision c1899de289a04d12100db370d81485cdf75e47ca + +docker exec -e HF_HUB_OFFLINE=1 vllm-snapshot vllm snapshot create Qwen/Qwen3-0.6B \ + --snapshot-dir /snapshots/qwen3-0.6b \ + --revision c1899de289a04d12100db370d81485cdf75e47ca \ + --dtype float16 \ + --max-model-len 512 + +docker exec vllm-snapshot vllm snapshot inspect /snapshots/qwen3-0.6b + +docker exec vllm-snapshot vllm snapshot restore \ + /snapshots/qwen3-0.6b --host 0.0.0.0 --port 8000 +``` + +Keep that container and its mounts available while the snapshot is in use. +Stop and remove it only after the restored API server is no longer needed. + +Create initializes the engine, records a one-token canary, releases and reloads +weights and KV cache to rehearse restore, then releases them again for capture. +The manifest is published only after CRIU completes and the source tree stops. +Creation is offline preparation and is not part of restore latency. + +The private `0700` artifact contains process memory, CUDA state, engine +arguments, compatibility identity, and canary output; its manifest is `0600`. +Literal API keys and Hugging Face tokens are redacted from the manifest's engine +arguments. The manifest records selected environment names plus deterministic +name-bound fingerprints, not their values. The protected CRIU artifact can +still contain process secrets, so treat it as sensitive data. Restore reloads +model files and KV cache before binding HTTP, then reproduces the canary or +tears down the restored tree. The inspect command prints that identity and +canary without executing the saved process. + +## Restore behavior + +Restore fails before CRIU runs if the saved identity does not match the current +host. It does not silently fall back to ordinary startup. After CRIU restores +the process tree, vLLM releases the saved engine to bind the requested HTTP +address and checks the first generated token and sampled-token log probability +against the snapshot canary. The command returns only after that check passes. +The restored API server continues to run as a detached process. + +The pre-release port probe is best-effort only. It neither reserves the port +nor authenticates the listener that appears afterward. + +Rollback terminates and waits for the restored tree after process identity +verification succeeds. If verification fails, vLLM writes an abort marker for +the snapshot server instead of signaling unverified PIDs. If that server has +already exited, surviving engine processes can require operator cleanup. +Verify that the previous tree has stopped before retrying; do not kill a process +solely because its PID appears in the manifest, since PIDs can be reused. + +The artifact is reusable after its previous restored tree stops. Only one +snapshot or external CRIU operation may use a shared `/dev/shm` mount at a time. + +## Tradeoffs and limitations + +- Creation has its own latency and briefly requires the full engine. Artifact + size can approach the captured process and GPU memory. +- Artifacts are not guaranteed to survive a power loss. The manifest is + fsynced, but publication does not explicitly flush all captured payload files + and directories. Recreate the artifact after an unclean host shutdown. +- Restore currently requires the same host, GPU, driver, kernel, Python, + PyTorch, installed vLLM version, model revision, engine arguments, + selected environment variables, and CRIU plugin binaries. +- Only dense float16 TP1 has been validated. Other model formats depend on their + existing sleep level 2 reload support; NCCL state is not restored. +- CRIU support varies by kernel and driver. Preserve package, library, model, + and generated-cache paths for the artifact lifetime. +- A snapshot can include application secrets or request state present in the + process. Build it before traffic and protect it like process memory. +- Any feature that opens an external connection must close it before capture. + +For a lower-complexity option that retains a live process, see +[Sleep mode](sleep_mode.md). Sleep mode and initialized snapshots retain +different amounts of state and have different idle resource costs. diff --git a/tests/entrypoints/launchers/test_launch_cli.py b/tests/entrypoints/launchers/test_launch_cli.py index 443dd82fdd..fc83010698 100644 --- a/tests/entrypoints/launchers/test_launch_cli.py +++ b/tests/entrypoints/launchers/test_launch_cli.py @@ -3,15 +3,48 @@ """Unit tests for the `vllm launch` CLI subcommand.""" import argparse -from unittest.mock import patch +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 @@ -109,3 +142,1011 @@ def test_launch_registered_in_main(): 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( + "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 + + +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() + + +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( + ("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_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 diff --git a/tests/models/quantization/test_mimo_v2_qkv_shard.py b/tests/models/quantization/test_mimo_v2_qkv_shard.py index 91eb036bc7..68afe7fdec 100644 --- a/tests/models/quantization/test_mimo_v2_qkv_shard.py +++ b/tests/models/quantization/test_mimo_v2_qkv_shard.py @@ -19,6 +19,7 @@ import torch from vllm.model_executor.models.mimo_v2 import _shard_fp8_qkv_proj from vllm.utils.math_utils import cdiv +from vllm.utils.torch_utils import set_default_torch_num_threads pytestmark = pytest.mark.cpu_test @@ -39,6 +40,13 @@ GEOMETRIES = { } +@pytest.fixture(scope="module", autouse=True) +def single_threaded_cpu_ops(): + # Do not prime the parent's thread pool before later engine forks. + with set_default_torch_num_threads(1): + yield + + def _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp): return ( num_heads // ckpt_tp * head_dim diff --git a/tools/install_snapshot_runtime.sh b/tools/install_snapshot_runtime.sh new file mode 100755 index 0000000000..40d6ad944f --- /dev/null +++ b/tools/install_snapshot_runtime.sh @@ -0,0 +1,109 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +set -euxo pipefail + +# Surface the isolated snapshot-runtime install cost in the CI build log. The +# trap captures the status first: the date substitution below would reset it. +_snapshot_install_started=$(date +%s) +code=0 +trap 'code=$?; echo "[TIMING] snapshot runtime install: $(( $(date +%s) - _snapshot_install_started )) s (exit $code)"' EXIT + +readonly CRIU_VERSION="4.2.1" +readonly CRIU_SOURCE_SHA256="feffdf4638125ebb12d2434754f80a1d7bbba85a3e6bee98c216f88fb99a5d96" +readonly DYNAMO_CHECKPOINT_HELPER_COMMIT="fdf25efad60f696a73393caf341c527a93d84190" +readonly DYNAMO_CHECKPOINT_HELPER_SHA256="7b8ea21baecb1da729202c32e0fecbbbafb789d65e5c1c0649ea253ad09b8f82" +readonly DYNAMO_LICENSE_SHA256="aeeca3d74a13b91d17ab9211c407b8d94a77bf8d9e5750c467818d3066b8adcf" + +# Keep the compiler and NL headers required by runtime JIT and RDMA. +snapshot_shared_build_deps=( + build-essential + libnl-3-dev + libnl-route-3-dev +) +snapshot_build_deps=( + pkg-config + libbsd-dev + libcap-dev + libnet1-dev + libprotobuf-dev + libprotobuf-c-dev + protobuf-c-compiler + protobuf-compiler + python3-protobuf + libgnutls28-dev + libnftables-dev + uuid-dev +) +snapshot_runtime_deps=( + libbsd0 + libcap2 + libnet1 + libnl-3-200 + libnl-route-3-200 + libprotobuf-c1 + libgnutls30t64 + libnftables1 + iproute2 + iptables + procps + uuid-runtime + util-linux +) + +apt-get update -y +apt-get install -y --no-install-recommends \ + "${snapshot_shared_build_deps[@]}" \ + "${snapshot_build_deps[@]}" \ + "${snapshot_runtime_deps[@]}" + +snapshot_build_dir="$(mktemp -d)" +curl -fsSL \ + "https://github.com/checkpoint-restore/criu/archive/refs/tags/v${CRIU_VERSION}.tar.gz" \ + -o "${snapshot_build_dir}/criu.tar.gz" +echo "${CRIU_SOURCE_SHA256} ${snapshot_build_dir}/criu.tar.gz" | sha256sum -c - +mkdir "${snapshot_build_dir}/criu" +tar -xzf "${snapshot_build_dir}/criu.tar.gz" \ + --strip-components=1 -C "${snapshot_build_dir}/criu" +make -C "${snapshot_build_dir}/criu" -j"${MAX_JOBS:-$(nproc)}" criu cuda_plugin +install -D -m 0755 "${snapshot_build_dir}/criu/criu/criu" \ + /usr/local/sbin/criu +install -D -m 0755 \ + "${snapshot_build_dir}/criu/plugins/cuda/cuda_plugin.so" \ + /usr/local/lib/criu/cuda_plugin.so + +helper_url="https://raw.githubusercontent.com/ai-dynamo/dynamo/${DYNAMO_CHECKPOINT_HELPER_COMMIT}/deploy/snapshot/cmd/cuda-checkpoint-helper/main.c" +curl -fsSL "${helper_url}" -o "${snapshot_build_dir}/cuda-checkpoint.c" +echo "${DYNAMO_CHECKPOINT_HELPER_SHA256} ${snapshot_build_dir}/cuda-checkpoint.c" \ + | sha256sum -c - +gcc -O2 -Wall -Wextra \ + -I/usr/local/cuda/include \ + -L/usr/local/cuda/lib64/stubs \ + "${snapshot_build_dir}/cuda-checkpoint.c" -lcuda \ + -o /usr/local/sbin/cuda-checkpoint + +license_url="https://raw.githubusercontent.com/ai-dynamo/dynamo/${DYNAMO_CHECKPOINT_HELPER_COMMIT}/LICENSE" +curl -fsSL "${license_url}" -o "${snapshot_build_dir}/DYNAMO-LICENSE" +echo "${DYNAMO_LICENSE_SHA256} ${snapshot_build_dir}/DYNAMO-LICENSE" \ + | sha256sum -c - + +install -D -m 0644 "${snapshot_build_dir}/criu.tar.gz" \ + /usr/share/doc/criu/source/criu.tar.gz +install -D -m 0644 "${snapshot_build_dir}/criu/COPYING" \ + /usr/share/doc/criu/COPYING +install -D -m 0644 "${snapshot_build_dir}/cuda-checkpoint.c" \ + /usr/share/doc/cuda-checkpoint-helper/main.c +install -D -m 0644 "${snapshot_build_dir}/DYNAMO-LICENSE" \ + /usr/share/doc/cuda-checkpoint-helper/LICENSE + +criu --version +ln -s /usr/local/cuda/lib64/stubs/libcuda.so \ + "${snapshot_build_dir}/libcuda.so.1" +LD_LIBRARY_PATH="${snapshot_build_dir}" \ + cuda-checkpoint --help | grep -q -- "--action" +test -x /usr/local/lib/criu/cuda_plugin.so + +apt-get purge -y --auto-remove "${snapshot_build_deps[@]}" +rm -rf "${snapshot_build_dir}" +rm -rf /var/lib/apt/lists/* diff --git a/vllm/entrypoints/cli/main.py b/vllm/entrypoints/cli/main.py index 4ac960d8cb..ab9afcb008 100644 --- a/vllm/entrypoints/cli/main.py +++ b/vllm/entrypoints/cli/main.py @@ -46,6 +46,7 @@ def main(): import vllm.entrypoints.cli.openai import vllm.entrypoints.cli.run_batch import vllm.entrypoints.cli.serve + import vllm.entrypoints.cli.snapshot from vllm.entrypoints.serve.utils.api_utils import ( VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup, @@ -59,9 +60,11 @@ def main(): vllm.entrypoints.cli.benchmark.main, vllm.entrypoints.cli.collect_env, vllm.entrypoints.cli.run_batch, + vllm.entrypoints.cli.snapshot, ] - cli_env_setup() + if sys.argv[1:2] != ["snapshot"]: + cli_env_setup() vllm.entrypoints.cli.benchmark.main.maybe_exec_rust_bench() @@ -94,7 +97,12 @@ def main(): subparsers = parser.add_subparsers(required=False, dest="subparser") cmds = {} for cmd_module in CMD_MODULES: - new_cmds = cmd_module.cmd_init() + if cmd_module is vllm.entrypoints.cli.snapshot: + new_cmds = cmd_module.cmd_init( + create_requested=sys.argv[1:3] == ["snapshot", "create"] + ) + else: + new_cmds = cmd_module.cmd_init() for cmd in new_cmds: cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd) cmds[cmd.name] = cmd diff --git a/vllm/entrypoints/cli/snapshot.py b/vllm/entrypoints/cli/snapshot.py new file mode 100644 index 0000000000..f5201675d4 --- /dev/null +++ b/vllm/entrypoints/cli/snapshot.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import argparse +import json +import os +import platform +import typing +from pathlib import Path + +from vllm.entrypoints.cli.types import CLISubcommand + +if typing.TYPE_CHECKING: + from vllm.utils.argparse_utils import FlexibleArgumentParser +else: + FlexibleArgumentParser = argparse.ArgumentParser + + +def validate_create_args(args: argparse.Namespace) -> None: + model = str(getattr(args, "model_tag", None) or args.model) + revision = str(getattr(args, "revision", None) or "") + if ( + Path(model).exists() + or len(revision) != 40 + or any(character not in "0123456789abcdefABCDEF" for character in revision) + ): + raise ValueError( + "snapshot create requires an immutable remote model and " + "40-character --revision" + ) + parallel_sizes = ( + ("tensor parallel", args.tensor_parallel_size), + ("pipeline parallel", args.pipeline_parallel_size), + ("data parallel", args.data_parallel_size), + ("prefill context parallel", args.prefill_context_parallel_size), + ) + for name, size in parallel_sizes: + if size != 1: + raise ValueError(f"snapshot create currently requires {name} size 1") + unsupported_frontend = ( + args.api_server_count not in (None, 1) + or args.grpc + or args.headless + or args.api_key + or os.environ.get("VLLM_API_KEY") + or args.uds + or args.middleware + ) + if unsupported_frontend or any( + ( + args.ssl_keyfile, + args.ssl_certfile, + args.ssl_ca_certs, + args.enable_ssl_refresh, + args.ssl_ciphers, + ) + ): + raise ValueError( + "snapshot create requires one unauthenticated plaintext HTTP server" + ) + if args.logprobs_mode in {"raw_logits", "processed_logits"}: + raise ValueError("snapshot canary requires a log-probability mode") + if any( + getattr(args, name, None) is not None + for name in ("speculative_config", "spec_method", "spec_model", "spec_tokens") + ): + raise ValueError("snapshot create does not support speculative decoding") + if platform.system() != "Linux" or platform.machine() != "x86_64": + raise ValueError("snapshot create requires Linux x86_64") + + +def run_create(args: argparse.Namespace) -> None: + validate_create_args(args) + from vllm.snapshot.controller import create_snapshot + + create_snapshot(args) + + +def run_inspect(args: argparse.Namespace) -> None: + from vllm.snapshot.manifest import inspect_snapshot + + print(json.dumps(inspect_snapshot(Path(args.snapshot_dir)), indent=2)) + + +def run_restore(args: argparse.Namespace) -> None: + from vllm.snapshot.controller import restore_snapshot + + restore_snapshot(args) + + +class SnapshotSubcommand(CLISubcommand): + """The `snapshot` subcommand for the vLLM CLI.""" + + name = "snapshot" + + def __init__(self, *, create_requested: bool = False) -> None: + self.create_requested = create_requested + + @staticmethod + def cmd(args: argparse.Namespace) -> None: + args.snapshot_dispatch(args) + + def subparser_init( + self, subparsers: argparse._SubParsersAction + ) -> FlexibleArgumentParser: + parser = subparsers.add_parser( + self.name, + help="Create, inspect, or restore an initialized vLLM snapshot.", + usage="vllm snapshot [options]", + ) + actions = parser.add_subparsers(required=True, dest="snapshot_action") + + create_parser = actions.add_parser( + "create", help="Create a snapshot from an initialized TP1 engine." + ) + if self.create_requested: + from vllm.entrypoints.launchers.cli_args import make_arg_parser + + create_parser = make_arg_parser(create_parser) + else: + create_parser.add_argument("model_tag", nargs="?") + create_parser.add_argument("--snapshot-dir", required=True) + create_parser.set_defaults(snapshot_dispatch=run_create) + + inspect_parser = actions.add_parser( + "inspect", help="Inspect a snapshot without restoring it." + ) + inspect_parser.add_argument("snapshot_dir") + inspect_parser.set_defaults(snapshot_dispatch=run_inspect) + + restore_parser = actions.add_parser( + "restore", help="Restore a same-host TP1 snapshot." + ) + restore_parser.add_argument("snapshot_dir") + restore_parser.add_argument("--host", default=None) + restore_parser.add_argument("--port", type=int, default=8000) + restore_parser.set_defaults(snapshot_dispatch=run_restore) + + return parser + + +def cmd_init(*, create_requested: bool = False) -> list[CLISubcommand]: + return [SnapshotSubcommand(create_requested=create_requested)] diff --git a/vllm/snapshot/__init__.py b/vllm/snapshot/__init__.py new file mode 100644 index 0000000000..8c4fde1282 --- /dev/null +++ b/vllm/snapshot/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Create and restore explicitly compatible vLLM engine snapshots.""" diff --git a/vllm/snapshot/controller.py b/vllm/snapshot/controller.py new file mode 100644 index 0000000000..d5564fb863 --- /dev/null +++ b/vllm/snapshot/controller.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Create and restore transaction orchestration for local snapshots.""" + +import argparse +import shutil +import sys +from pathlib import Path + +from vllm.snapshot.manifest import ( + read_manifest, + validate_artifact_root, + validate_identity, +) +from vllm.snapshot.runtime import ( + LocalSnapshotTools, + SnapshotCreateError, + SnapshotRestoreError, + _error_detail, +) +from vllm.snapshot.types import Oracle, oracles_match + + +def _child_engine_argv( + args: argparse.Namespace, engine_argv: tuple[str, ...] | None +) -> tuple[str, ...]: + """Build the child argv from the first argv source that is present. + + In order: an explicit ``engine_argv`` from the caller, the process argv + tail when invoked as ``vllm snapshot create ...``, then the bare model from + the parsed args. Whichever source wins, snapshot options are stripped and + sleep mode is appended. + """ + if not engine_argv: + process_argv = tuple(sys.argv[1:]) + if process_argv[:2] == ("snapshot", "create"): + engine_argv = process_argv[2:] + else: + model = getattr(args, "model_tag", None) or getattr(args, "model", None) + if not model: + raise SnapshotCreateError("snapshot create requires a model argument") + engine_argv = (str(model),) + child_argv: list[str] = [] + iterator = iter(engine_argv) + for item in iterator: + option, separator, _ = item.partition("=") + option = option.replace("_", "-") + if len(option) > 2 and "--snapshot-dir".startswith(option): + if not separator: + next(iterator, None) + else: + child_argv.append(item) + if "--enable-sleep-mode" not in child_argv: + child_argv.append("--enable-sleep-mode") + return tuple(child_argv) + + +def create_snapshot( + args: argparse.Namespace, + *, + engine_argv: tuple[str, ...] | None = None, + tools: LocalSnapshotTools | None = None, +) -> None: + target = Path(args.snapshot_dir).absolute() + toolset = tools or LocalSnapshotTools() + toolset.preflight("create", target) + if target.exists() or target.is_symlink(): + raise SnapshotCreateError(f"snapshot target already exists: {target}") + validate_artifact_root(target, creating=True) + target.parent.mkdir(mode=0o700, parents=False, exist_ok=True) + validate_artifact_root(target.parent, creating=False) + target.mkdir(mode=0o700) + validate_artifact_root(target, creating=False) + published = False + root_pid: int | None = None + try: + child_argv = _child_engine_argv(args, engine_argv) + root_pid = toolset.launch_child(target, child_argv) + oracle = toolset.wait_ready(target, root_pid) + inventory = toolset.inventory(root_pid, target) + toolset.dump(target, inventory) + toolset.verify_dead(inventory) + manifest = toolset.make_manifest(args, child_argv, inventory, oracle, target) + toolset.publish(target, manifest) + published = True + except BaseException: + if root_pid is not None: + toolset.abort_create(root_pid) + raise + finally: + # Runs even when abort_create raised, so no partial artifact is left. + if not published and target.exists(): + shutil.rmtree(target) + + +def restore_snapshot( + args: argparse.Namespace, *, tools: LocalSnapshotTools | None = None +) -> None: + artifact = Path(args.snapshot_dir).absolute() + toolset = tools or LocalSnapshotTools() + toolset.preflight("restore", artifact) + manifest = read_manifest(artifact) + validate_identity(manifest, toolset.current_identity(manifest.gpu_uuid)) + + root_pid: int | None = None + try: + root_pid = toolset.restore(artifact, manifest) + toolset.release(artifact, args.host, args.port) + toolset.wait_listener(root_pid, args.host, args.port) + actual = toolset.request_oracle(args.host, args.port, manifest) + expected = Oracle( + token_ids=manifest.oracle_token_ids, + text=manifest.oracle_text, + sampled_token_logprob=manifest.oracle_sampled_token_logprob, + ) + if not oracles_match(expected, actual): + raise SnapshotRestoreError( + f"snapshot oracle mismatch: expected={expected!r}, actual={actual!r}" + ) + toolset.complete_restore(root_pid) + except BaseException as error: + if root_pid is not None: + try: + toolset.cleanup(root_pid) + except BaseException as cleanup_error: + raise SnapshotRestoreError( + f"snapshot restore failed: {_error_detail(error)}; " + f"cleanup failed: {_error_detail(cleanup_error)}" + ) from error + if isinstance(error, SnapshotRestoreError): + raise + raise SnapshotRestoreError(str(error)) from error diff --git a/vllm/snapshot/manifest.py b/vllm/snapshot/manifest.py new file mode 100644 index 0000000000..a029c544c5 --- /dev/null +++ b/vllm/snapshot/manifest.py @@ -0,0 +1,277 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import math +import os +import stat +from pathlib import Path +from typing import Annotated, Any, Literal + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + StrictFloat, + StrictInt, + StrictStr, + ValidationError, + field_validator, + model_validator, +) +from typing_extensions import Self + + +class SnapshotCompatibilityError(RuntimeError): + """The snapshot identity does not match the requested runtime.""" + + +class SnapshotSecurityError(RuntimeError): + """The snapshot artifact does not meet the private-path contract.""" + + +_MAX_PID = 2**31 - 1 +_PositivePid = Annotated[int, Field(gt=0, le=_MAX_PID)] +_OracleTokenId = Annotated[int, Field(ge=0)] + + +class SnapshotRuntimeIdentity(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + vllm_version: str + python_version: str + torch_version: str + cuda_runtime: str + driver_version: str + criu_version: str + cuda_checkpoint_sha256: str + kernel_release: str + host_id: str + gpu_name: str + gpu_uuid: str + environment: tuple[tuple[str, str], ...] + + +class SnapshotManifest(SnapshotRuntimeIdentity): + schema_version: Literal[1] + boundary: Literal["post-engine-init-reloadable-state-released"] + created_at: str + artifact_bytes: Annotated[int, Field(ge=0)] + model: str + served_model_name: str + model_revision: str + tokenizer_revision: str + engine_argv: tuple[str, ...] + process_tree: Annotated[tuple[_PositivePid, ...], Field(min_length=1)] + cuda_holders: Annotated[tuple[_PositivePid, ...], Field(min_length=1)] + external_cache_files: tuple[tuple[str, str], ...] = () + oracle_token_ids: Annotated[ + tuple[_OracleTokenId, ...], Field(min_length=1, max_length=1) + ] + oracle_text: str + oracle_sampled_token_logprob: float + + @field_validator("schema_version", mode="before") + @classmethod + def _require_schema_version_one(cls, value: Any) -> Any: + if type(value) is not int or value != 1: + raise ValueError("schema version must be integer 1") + return value + + @field_validator("oracle_sampled_token_logprob", mode="before") + @classmethod + def _require_finite_float(cls, value: Any) -> Any: + if type(value) is not float or not math.isfinite(value): + raise ValueError("sampled token log probability must be a finite float") + return value + + @model_validator(mode="after") + def _validate_process_relationships(self) -> Self: + if len(set(self.process_tree)) != len(self.process_tree): + raise ValueError("process_tree contains duplicate PIDs") + if len(set(self.cuda_holders)) != len(self.cuda_holders): + raise ValueError("cuda_holders contains duplicate PIDs") + if not set(self.cuda_holders).issubset(self.process_tree): + raise ValueError("cuda_holders must be a subset of process_tree") + return self + + +class ReleaseMarker(BaseModel): + """Release marker the controller writes; unknown keys are ignored.""" + + release: Literal[True] + host: StrictStr | None = None + port: Annotated[StrictInt, Field(ge=1, le=65535)] + + @field_validator("release", mode="before") + @classmethod + def _require_json_true(cls, value: Any) -> Any: + # A literal accepts 1 as True; the barrier requires JSON true itself. + if value is not True: + raise ValueError("release must be JSON true") + return value + + +class ReadyMarker(BaseModel): + """Ready marker the child writes; unknown keys are ignored.""" + + token_ids: tuple[StrictInt, ...] + text: StrictStr + sampled_token_logprob: StrictFloat | StrictInt + + +def validate_identity( + expected: SnapshotRuntimeIdentity, actual: SnapshotRuntimeIdentity +) -> None: + """Require an exact match for every field that can affect compatibility.""" + # ``expected`` may be a manifest, so compare the identity fields only. + fields = set(SnapshotRuntimeIdentity.model_fields) + expected_values = expected.model_dump(include=fields) + actual_values = actual.model_dump(include=fields) + differing = [ + name + for name in SnapshotRuntimeIdentity.model_fields + if expected_values[name] != actual_values[name] + ] + if differing: + raise SnapshotCompatibilityError(f"snapshot mismatch: {', '.join(differing)}") + + +def _existing_path_chain(path: Path): + current = path + while True: + if current.exists() or current.is_symlink(): + yield current + if current.parent == current: + break + current = current.parent + + +def _is_trusted_sticky_directory(path_stat: os.stat_result) -> bool: + return ( + path_stat.st_uid == 0 + and stat.S_ISDIR(path_stat.st_mode) + and bool(path_stat.st_mode & stat.S_ISVTX) + ) + + +def validate_artifact_root(path: Path, *, creating: bool) -> None: + """Validate that an artifact path cannot be replaced by another user.""" + path = Path(os.path.abspath(path)) + if path.is_symlink(): + raise SnapshotSecurityError(f"snapshot path is a symlink: {path}") + if not creating and not path.exists(): + raise SnapshotSecurityError(f"snapshot path does not exist: {path}") + + effective_uid = os.geteuid() + for component in _existing_path_chain(path): + component_stat = component.lstat() + if stat.S_ISLNK(component_stat.st_mode): + raise SnapshotSecurityError( + f"snapshot path contains a symlink: {component}" + ) + if component_stat.st_uid not in (0, effective_uid): + raise SnapshotSecurityError( + f"snapshot path has an untrusted owner: {component}" + ) + if component_stat.st_mode & ( + stat.S_IWGRP | stat.S_IWOTH + ) and not _is_trusted_sticky_directory(component_stat): + raise SnapshotSecurityError( + f"snapshot path has a group- or world-writable ancestor: {component}" + ) + + if path.exists(): + path_stat = path.lstat() + if not stat.S_ISDIR(path_stat.st_mode): + raise SnapshotSecurityError(f"snapshot path is not a directory: {path}") + if path_stat.st_uid != effective_uid: + raise SnapshotSecurityError( + f"snapshot directory is not owned by the current user: {path}" + ) + if stat.S_IMODE(path_stat.st_mode) & 0o077: + raise SnapshotSecurityError( + f"snapshot directory must have mode 0700: {path}" + ) + + +def _validation_path(error: ValidationError) -> str: + detail = error.errors(include_url=False)[0] + location = detail["loc"] + if not location: + return "JSON" if detail["type"] == "json_invalid" else "root" + return ".".join(map(str, location)) + + +def _fsync_directory(path: Path) -> None: + directory_fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + + +def _write_json_atomic(path: Path, payload: object, *, overwrite: bool = False) -> None: + temporary = path.with_name(f"{path.name}.tmp") + if overwrite: + temporary.unlink(missing_ok=True) + elif path.exists() or path.is_symlink(): + raise FileExistsError(path) + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0) + descriptor = os.open(temporary, flags, 0o600) + try: + with os.fdopen(descriptor, "wb") as output: + output.write( + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + ) + output.flush() + os.fsync(output.fileno()) + os.replace(temporary, path) + _fsync_directory(path.parent) + except BaseException: + temporary.unlink(missing_ok=True) + raise + + +def write_manifest_atomic(path: Path, manifest: SnapshotManifest) -> None: + """Write a new private manifest and make its directory entry durable.""" + path = Path(path) + validate_artifact_root(path, creating=False) + destination = path / "manifest.json" + try: + _write_json_atomic(destination, manifest.model_dump(mode="json")) + except FileExistsError as error: + raise SnapshotSecurityError( + f"manifest already exists: {destination}" + ) from error + + +def read_manifest(path: Path) -> SnapshotManifest: + """Read a snapshot published by atomically installing its manifest.""" + path = Path(path) + validate_artifact_root(path, creating=False) + manifest_path = path / "manifest.json" + if manifest_path.is_symlink(): + raise SnapshotSecurityError(f"manifest is a symlink: {manifest_path}") + flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) + try: + file_descriptor = os.open(manifest_path, flags) + except FileNotFoundError as error: + raise SnapshotSecurityError( + f"manifest does not exist: {manifest_path}" + ) from error + try: + with os.fdopen(file_descriptor, encoding="utf-8") as manifest_file: + payload = manifest_file.read() + except UnicodeDecodeError as error: + raise SnapshotCompatibilityError("invalid snapshot manifest: JSON") from error + try: + return SnapshotManifest.model_validate_json(payload, strict=True) + except ValidationError as error: + raise SnapshotCompatibilityError( + f"invalid snapshot manifest: {_validation_path(error)}" + ) from error + + +def inspect_snapshot(path: Path) -> dict[str, object]: + return read_manifest(path).model_dump(mode="json") diff --git a/vllm/snapshot/runtime.py b/vllm/snapshot/runtime.py new file mode 100644 index 0000000000..582bc9a7d9 --- /dev/null +++ b/vllm/snapshot/runtime.py @@ -0,0 +1,1246 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Host runtime operations for local snapshots.""" + +import argparse +import ast +import hashlib +import importlib.metadata +import importlib.util +import json +import os +import platform +import select +import shutil +import signal +import socket +import stat +import struct +import subprocess +import sys +import tempfile +import time +import urllib.request +from collections.abc import Iterable +from contextlib import suppress +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import NoReturn + +import regex as re +from pydantic import ValidationError + +from vllm.snapshot.manifest import ( + ReadyMarker, + SnapshotManifest, + SnapshotRuntimeIdentity, + _fsync_directory, + _validation_path, + _write_json_atomic, + validate_artifact_root, + write_manifest_atomic, +) +from vllm.snapshot.types import Oracle + + +class SnapshotCreateError(RuntimeError): + """Snapshot creation did not produce a complete artifact.""" + + +class SnapshotRestoreError(RuntimeError): + """Snapshot restore failed.""" + + +_COMMON_CRIU_OPTIONS = ( + "--shell-job", + "--ext-unix-sk", + "--tcp-established", + "--link-remap", + "--file-locks", +) +_TCP_TABLES = ( + ("AF_INET", Path("/proc/net/tcp")), + ("AF_INET6", Path("/proc/net/tcp6")), +) + +_CACHE_REMEDY = "recreate the snapshot or restore with the same cache state" +# Recorded instead of a digest for files whose content is volatile by design. +_CACHE_LOCKED = "locked" +# Recorded instead of a digest for a file capture found but could not read. +_CACHE_UNREADABLE = "unreadable" + + +def _generated_cache_prefixes() -> tuple[str, ...]: + """Path prefixes of the generated caches an engine process keeps open. + + vLLM redirects Inductor and Triton under ``VLLM_CACHE_ROOT`` when it + compiles, but a process can still hold their upstream defaults open. + A cache root is resolved and separator-terminated so that it still matches + the kernel-canonical descriptor targets under a symlinked root, and so that + a sibling directory sharing its name does not. Inductor's default location + is a bare prefix instead, because its directory carries a per-user suffix. + """ + import vllm.envs as envs + + inductor_root = os.environ.get("TORCHINDUCTOR_CACHE_DIR") + roots = tuple( + f"{os.path.realpath(os.path.expanduser(root)).rstrip(os.sep)}{os.sep}" + for root in ( + envs.VLLM_CACHE_ROOT, + os.environ.get("TRITON_CACHE_DIR") or "~/.triton/cache", + "~/.cache/flashinfer", + *((inductor_root,) if inductor_root else ()), + ) + ) + if inductor_root: + return roots + return roots + ( + os.path.join(os.path.realpath(tempfile.gettempdir()), "torchinductor_"), + ) + + +def _error_detail(error: BaseException) -> str: + message = str(error) + return f"{type(error).__name__}: {message}" if message else type(error).__name__ + + +def _close_pidfds(handles: Iterable[tuple[int, int]]) -> None: + for _pid, pidfd in handles: + with suppress(OSError): + os.close(pidfd) + + +@dataclass(frozen=True) +class ProcessInventory: + root_pid: int + process_tree: tuple[int, ...] + cuda_holders: tuple[int, ...] + gpu_uuid: str + external_cache_files: tuple[tuple[str, str], ...] = () + + +@dataclass(frozen=True) +class _TcpSocketRecord: + family: str + local_raw: str + remote_raw: str + inode: int + + +def _process_comm(pid: int) -> str: + try: + return (Path("/proc") / str(pid) / "comm").read_text().strip() + except OSError: + return "unknown" + + +def _decode_endpoint(family: str, raw: str) -> str: + """Render a ``/proc/net/tcp`` ``HEXADDR:HEXPORT`` field as ``ip:port``. + + The kernel prints the address as little-endian 32-bit words: one for IPv4, + four in order for IPv6. + """ + hex_address, _, hex_port = raw.partition(":") + port = int(hex_port, 16) + if family == "AF_INET": + packed = struct.pack(" None: + owned = tuple(record for record in records if record.inode in socket_owners) + endpoints = { + (record.family, record.local_raw, record.remote_raw) for record in owned + } + offenders = tuple( + record + for record in owned + if (record.family, record.remote_raw, record.local_raw) not in endpoints + ) + if not offenders: + return + named = "; ".join( + f"pid {socket_owners[record.inode]} " + f"({_process_comm(socket_owners[record.inode])}) " + f"{_decode_endpoint(record.family, record.local_raw)} -> " + f"{_decode_endpoint(record.family, record.remote_raw)}" + for record in offenders[:5] + ) + remaining = f" (+{len(offenders) - 5} more)" if len(offenders) > 5 else "" + raise SnapshotCreateError( + f"snapshot tree has an external established TCP connection: {named}{remaining}" + ) + + +class LocalSnapshotTools: + """Host-owned CRIU and cuda-checkpoint implementation.""" + + def __init__(self) -> None: + self.criu = shutil.which("criu") or "criu" + self.cuda_checkpoint = shutil.which("cuda-checkpoint") or "cuda-checkpoint" + self.nvidia_smi = shutil.which("nvidia-smi") or "nvidia-smi" + plugin_dir = os.environ.get("CRIU_CUDA_PLUGIN_DIR", "") + self.plugin_dir = Path(plugin_dir) if plugin_dir else None + self.shm_dir = Path("/dev/shm") + self.timeout_s = float(os.environ.get("VLLM_SNAPSHOT_TIMEOUT_S", "900")) + if not 0 < self.timeout_s < float("inf"): + raise ValueError("VLLM_SNAPSHOT_TIMEOUT_S must be a positive finite number") + self._children: dict[int, subprocess.Popen[bytes]] = {} + self._restored_processes: dict[int, tuple[tuple[int, int], ...]] = {} + + def _privileged(self) -> list[str]: + return [] if os.geteuid() == 0 else ["sudo", "-n"] + + def _run( + self, + command: list[str], + *, + timeout: float | None = None, + ) -> subprocess.CompletedProcess[str]: + return subprocess.run( + command, + check=True, + capture_output=True, + text=True, + timeout=timeout or self.timeout_s, + ) + + def preflight(self, action: str, artifact: Path) -> None: + if platform.system() != "Linux" or platform.machine() != "x86_64": + raise RuntimeError("snapshot requires Linux x86_64") + if action == "restore" and ( + not callable(getattr(os, "pidfd_open", None)) + or not callable(getattr(signal, "pidfd_send_signal", None)) + ): + raise RuntimeError("snapshot requires Linux pidfd support") + for command in (self.criu, self.cuda_checkpoint, self.nvidia_smi): + if shutil.which(command) is None: + raise RuntimeError(f"snapshot dependency not found: {command}") + if self.plugin_dir is None: + raise RuntimeError("CRIU_CUDA_PLUGIN_DIR is required") + if not (self.plugin_dir / "cuda_plugin.so").is_file(): + raise RuntimeError(f"CRIU CUDA plugin not found in {self.plugin_dir}") + if os.geteuid() != 0: + self._run(["sudo", "-n", "true"], timeout=5) + if action == "restore": + validate_artifact_root(artifact, creating=False) + else: + validate_artifact_root(artifact, creating=True) + + def launch_child( + self, + workdir: Path, + engine_argv: tuple[str, ...], + ) -> int: + command = [ + sys.executable, + "-m", + "vllm.snapshot.server", + "--ready-file", + str(workdir / "ready.json"), + "--release-file", + str(workdir / "release.json"), + "--release-timeout-s", + str(self.timeout_s), + "--", + *engine_argv, + ] + child_environment = os.environ.copy() + child_environment.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + with (workdir / "child.log").open("wb") as log_file: + process = subprocess.Popen( + command, + env=child_environment, + stdin=subprocess.DEVNULL, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + self._children[process.pid] = process + return process.pid + + def wait_ready(self, workdir: Path, root_pid: int) -> Oracle: + ready_file = workdir / "ready.json" + deadline = time.monotonic() + self.timeout_s + waitid = os.waitid # type: ignore[attr-defined] + while time.monotonic() < deadline: + if ready_file.is_file() and ready_file.stat().st_size: + try: + marker = ReadyMarker.model_validate_json(ready_file.read_text()) + except ValidationError as error: + raise SnapshotCreateError( + f"snapshot ready marker invalid: {_validation_path(error)}" + ) from error + return Oracle( + token_ids=marker.token_ids, + text=marker.text, + sampled_token_logprob=marker.sampled_token_logprob, + ) + process = self._children.get(root_pid) + if process is not None and waitid( + os.P_PID, + root_pid, + os.WEXITED | os.WNOHANG | os.WNOWAIT, + ): + log = (workdir / "child.log").read_text(errors="replace") + raise SnapshotCreateError(f"snapshot child exited before ready:\n{log}") + try: + os.kill(root_pid, 0) + except ProcessLookupError as error: + log = (workdir / "child.log").read_text(errors="replace") + raise SnapshotCreateError( + f"snapshot child exited before ready:\n{log}" + ) from error + time.sleep(0.05) + raise SnapshotCreateError("snapshot child did not become ready before timeout") + + def _tree_pids(self, root_pid: int) -> tuple[int, ...]: + rows = self._run(["ps", "-eo", "pid=,ppid="]).stdout.splitlines() + children: dict[int, list[int]] = {} + for row in rows: + pid, parent_pid = map(int, row.split()) + children.setdefault(parent_pid, []).append(pid) + stack = [root_pid] + found: list[int] = [] + while stack: + pid = stack.pop() + if pid in found: + continue + found.append(pid) + stack.extend(children.get(pid, [])) + return tuple(found) + + def _cuda_process_rows(self) -> tuple[str, ...]: + output = self._run( + [ + self.nvidia_smi, + "--query-compute-apps=pid,gpu_uuid", + "--format=csv,noheader,nounits", + ] + ).stdout + return tuple(line for line in output.splitlines() if line.strip()) + + def _cuda_pids(self, rows: tuple[str, ...]) -> tuple[int, ...]: + return tuple(sorted({int(line.split(",", 1)[0].strip()) for line in rows})) + + def _descriptor_targets(self, pid: int) -> tuple[str, ...]: + descriptor_dir = Path("/proc") / str(pid) / "fd" + try: + descriptors = tuple(descriptor_dir.iterdir()) + except FileNotFoundError as error: + raise SnapshotCreateError( + f"snapshot process exited during descriptor inventory: {pid}" + ) from error + targets: list[str] = [] + for descriptor in descriptors: + try: + targets.append(os.readlink(descriptor)) + except FileNotFoundError: + continue + return tuple(targets) + + def _descriptor_inventory( + self, process_tree: tuple[int, ...] + ) -> tuple[tuple[int, ...], dict[int, int]]: + """Collect io_uring holders and each socket inode's owning pid.""" + io_uring_pids: list[int] = [] + socket_owners: dict[int, int] = {} + for pid in process_tree: + targets = self._descriptor_targets(pid) + if "anon_inode:[io_uring]" in targets: + io_uring_pids.append(pid) + for target in targets: + if target.startswith("socket:[") and target.endswith("]"): + # Shared descriptors name one holder; any of them locates + # the connection for the operator. + socket_owners.setdefault(int(target[len("socket:[") : -1]), pid) + return tuple(io_uring_pids), socket_owners + + def _tcp_records(self) -> tuple[_TcpSocketRecord, ...]: + records: list[_TcpSocketRecord] = [] + for family, table in _TCP_TABLES: + if not table.exists(): + continue + for row in table.read_text().splitlines()[1:]: + fields = row.split() + if fields[3] != "01": + continue + records.append( + _TcpSocketRecord( + family=family, + local_raw=fields[1], + remote_raw=fields[2], + inode=int(fields[9]), + ) + ) + return tuple(records) + + def _cache_fingerprint(self, path: Path) -> str | None: + """Digest a cache file, or mark the ones whose content is volatile.""" + try: + if not stat.S_ISREG(path.stat().st_mode): + return None + except OSError: + return None + if path.name.endswith(".lock"): + return _CACHE_LOCKED + try: + return self._sha256(path) + except OSError: + # Restore can only check that this one still exists, so record the + # downgrade rather than passing it off as a volatile file. + return _CACHE_UNREADABLE + + def _external_cache_files( + self, process_tree: tuple[int, ...], artifact: Path + ) -> tuple[tuple[str, str], ...]: + """Fingerprint the generated-cache files the captured tree holds open. + + CRIU reopens these by path on restore, so a JIT cache that rotates + between capture and restore fails the restore permanently. Recording + them lets restore fail early and name the file that changed. Files + inside the artifact are excluded: they belong to the snapshot itself. + Deleted targets are excluded too, because CRIU remaps those links. + """ + prefixes = _generated_cache_prefixes() + artifact_prefix = f"{os.path.abspath(artifact)}{os.sep}" + candidates: set[str] = set() + for pid in process_tree: + for target in self._descriptor_targets(pid): + if target.endswith(" (deleted)") or target.startswith(artifact_prefix): + continue + if target.endswith(".lock") or target.startswith(prefixes): + candidates.add(target) + recorded = ( + (path, self._cache_fingerprint(Path(path))) for path in sorted(candidates) + ) + return tuple( + (path, fingerprint) + for path, fingerprint in recorded + if fingerprint is not None + ) + + def inventory(self, root_pid: int, artifact: Path) -> ProcessInventory: + process_tree = self._tree_pids(root_pid) + cuda_rows = self._cuda_process_rows() + cuda_holders = tuple( + pid for pid in self._cuda_pids(cuda_rows) if pid in process_tree + ) + if not cuda_holders: + raise SnapshotCreateError("snapshot tree has no CUDA-holding process") + io_uring_pids, socket_owners = self._descriptor_inventory(process_tree) + if io_uring_pids: + raise SnapshotCreateError( + "snapshot tree owns io_uring state that CRIU cannot dump; set " + "kernel.io_uring_disabled=1 before starting an unprivileged " + "vLLM process, or use =2 to disable it host-wide " + f"(pids: {', '.join(map(str, io_uring_pids))})" + ) + _validate_tcp_connections(self._tcp_records(), socket_owners) + gpu_uuid = self._gpu_uuid_for_pids(cuda_holders, cuda_rows) + return ProcessInventory( + root_pid=root_pid, + process_tree=process_tree, + cuda_holders=cuda_holders, + gpu_uuid=gpu_uuid, + external_cache_files=self._external_cache_files(process_tree, artifact), + ) + + def _gpu_uuid_for_pids(self, pids: tuple[int, ...], rows: tuple[str, ...]) -> str: + wanted = set(pids) + matches: set[str] = set() + try: + for line in rows: + pid, gpu_uuid = (item.strip() for item in line.split(",", 1)) + if int(pid) in wanted: + matches.add(gpu_uuid) + except (TypeError, ValueError) as error: + raise SnapshotCreateError( + "could not identify the GPU used by the snapshot process tree" + ) from error + if len(matches) != 1: + raise SnapshotCreateError( + "snapshot process tree must use exactly one identifiable GPU" + ) + return matches.pop() + + def _criu( + self, + action: str, + artifact: Path, + arguments: list[str], + ) -> None: + if self.plugin_dir is None: + raise RuntimeError("CRIU_CUDA_PLUGIN_DIR is required") + env = os.environ.copy() + env["PATH"] = f"{Path(self.cuda_checkpoint).parent}:{env['PATH']}" + env["VLLM_SNAPSHOT_IGNORED_FD_MAP"] = str(artifact / "ignored-nvidia-fds.tsv") + env["RIG_CRIU_IGNORED_FD_MAP"] = env["VLLM_SNAPSHOT_IGNORED_FD_MAP"] + command = [ + *self._privileged(), + "env", + f"PATH={env['PATH']}", + f"VLLM_SNAPSHOT_IGNORED_FD_MAP={env['VLLM_SNAPSHOT_IGNORED_FD_MAP']}", + f"RIG_CRIU_IGNORED_FD_MAP={env['RIG_CRIU_IGNORED_FD_MAP']}", + self.criu, + action, + "-L", + str(self.plugin_dir), + *arguments, + ] + self._run(command) + + def _record_child_log_size(self, artifact: Path) -> None: + child_log = artifact / "child.log" + size_path = artifact / "child.log.snapshot-size" + payload = f"{child_log.stat().st_size}\n".encode() + descriptor = os.open( + size_path, + os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), + 0o600, + ) + with os.fdopen(descriptor, "wb") as output: + output.write(payload) + output.flush() + os.fsync(output.fileno()) + + def _link_remap_names(self) -> set[str]: + return { + path.name + for path in self.shm_dir.glob("link_remap.*") + if path.name.removeprefix("link_remap.").isdigit() + } + + def _capture_link_remaps(self, artifact: Path, names: set[str]) -> None: + if not names: + return + destination_dir = artifact / "link-remaps" + destination_dir.mkdir(mode=0o700) + for name in sorted(names): + source = self.shm_dir / name + source_fd = os.open( + source, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + ) + try: + metadata = os.fstat(source_fd) + if not stat.S_ISREG(metadata.st_mode) or metadata.st_uid != os.getuid(): + raise SnapshotCreateError(f"invalid CRIU link remap: {source}") + destination_fd = os.open( + destination_dir / name, + os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), + 0o600, + ) + with os.fdopen(destination_fd, "wb") as destination: + with os.fdopen(os.dup(source_fd), "rb") as contents: + shutil.copyfileobj(contents, destination) + destination.flush() + os.fsync(destination.fileno()) + finally: + os.close(source_fd) + source.unlink() + + def _rollback_link_remaps(self, created: dict[str, tuple[int, int]]) -> None: + for name, identity in created.items(): + target = self.shm_dir / name + try: + metadata = target.lstat() + except FileNotFoundError: + continue + if (metadata.st_dev, metadata.st_ino) == identity: + target.unlink() + + def _restore_link_remaps(self, artifact: Path) -> dict[str, tuple[int, int]]: + source_dir = artifact / "link-remaps" + if not source_dir.exists(): + return {} + created: dict[str, tuple[int, int]] = {} + try: + for source in sorted(source_dir.iterdir()): + if ( + not source.name.removeprefix("link_remap.").isdigit() + or source.is_symlink() + or not source.is_file() + ): + raise SnapshotRestoreError(f"invalid saved link remap: {source}") + payload = source.read_bytes() + target = self.shm_dir / source.name + try: + descriptor = os.open( + target, + os.O_WRONLY + | os.O_CREAT + | os.O_EXCL + | getattr(os, "O_NOFOLLOW", 0), + 0o600, + ) + except FileExistsError as error: + try: + descriptor = os.open( + target, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + ) + except OSError: + raise SnapshotRestoreError( + f"conflicting CRIU link remap exists: {target}" + ) from error + with os.fdopen(descriptor, "rb") as existing: + metadata = os.fstat(existing.fileno()) + if ( + not stat.S_ISREG(metadata.st_mode) + or metadata.st_uid != os.getuid() + or existing.read() != payload + ): + raise SnapshotRestoreError( + f"conflicting CRIU link remap exists: {target}" + ) from error + continue + metadata = os.fstat(descriptor) + created[source.name] = (metadata.st_dev, metadata.st_ino) + with os.fdopen(descriptor, "wb") as output: + output.write(payload) + output.flush() + os.fsync(output.fileno()) + except BaseException: + self._rollback_link_remaps(created) + raise + return created + + def _reset_child_log(self, artifact: Path) -> None: + size_path = artifact / "child.log.snapshot-size" + descriptor = os.open( + size_path, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + ) + with os.fdopen(descriptor, "r") as source: + size = int(source.read().strip()) + if size < 0: + raise SnapshotRestoreError("captured child log size is invalid") + + child_log = artifact / "child.log" + descriptor = os.open( + child_log, + os.O_WRONLY | getattr(os, "O_NOFOLLOW", 0), + ) + try: + metadata = os.fstat(descriptor) + if not stat.S_ISREG(metadata.st_mode) or metadata.st_size < size: + raise SnapshotRestoreError("captured child log is invalid") + os.ftruncate(descriptor, size) + os.fsync(descriptor) + finally: + os.close(descriptor) + + def dump(self, workdir: Path, inventory: ProcessInventory) -> None: + images = workdir / "images" + images.mkdir(mode=0o700) + fd_map = workdir / "ignored-nvidia-fds.tsv" + fd_map.touch(mode=0o600) + remaps_before = self._link_remap_names() + try: + self._criu( + "dump", + workdir, + [ + "--tree", + str(inventory.root_pid), + "--images-dir", + str(images), + "--log-file", + "dump.log", + *_COMMON_CRIU_OPTIONS, + ], + ) + self._capture_link_remaps( + workdir, + self._link_remap_names() - remaps_before, + ) + except BaseException: + for name in self._link_remap_names() - remaps_before: + (self.shm_dir / name).unlink(missing_ok=True) + raise + self._record_child_log_size(workdir) + + def verify_dead(self, inventory: ProcessInventory) -> None: + process = self._children.get(inventory.root_pid) + if process is not None: + process.wait(timeout=10) + for pid in inventory.process_tree: + try: + os.kill(pid, 0) + except ProcessLookupError: + continue + raise SnapshotCreateError(f"process survived CRIU dump: {pid}") + self._children.pop(inventory.root_pid, None) + + def _version(self, command: list[str]) -> str: + try: + result = self._run(command, timeout=10) + except subprocess.CalledProcessError as error: + return (error.stdout or error.stderr).splitlines()[0] + return result.stdout.splitlines()[0] if result.stdout else "unknown" + + def _sha256(self, path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + def _plugin_identity(self) -> tuple[tuple[str, str], ...]: + if self.plugin_dir is None: + return () + return tuple( + (f"snapshot_plugin_sha256:{path.name}", self._sha256(path)) + for path in sorted(self.plugin_dir.glob("*.so")) + ) + + def _torch_identity(self) -> tuple[str, str]: + torch_version = importlib.metadata.version("torch") + torch_spec = importlib.util.find_spec("torch") + locations = torch_spec and torch_spec.submodule_search_locations + if not locations: + raise RuntimeError("installed torch package could not be located") + version_path = Path(next(iter(locations))) / "version.py" + # Parse the version file rather than importing torch: snapshot create's + # CLI path stays import-light (vllm/entrypoints/cli/snapshot.py). + syntax = ast.parse(version_path.read_text(), filename=str(version_path)) + cuda_runtime = "" + for statement in syntax.body: + target = None + value = None + if isinstance(statement, ast.Assign) and len(statement.targets) == 1: + target = statement.targets[0] + value = statement.value + elif isinstance(statement, ast.AnnAssign): + target = statement.target + value = statement.value + if isinstance(target, ast.Name) and target.id == "cuda" and value: + parsed = ast.literal_eval(value) + cuda_runtime = str(parsed or "") + break + return torch_version, cuda_runtime + + def _gpu_identity(self, gpu_uuid: str) -> tuple[str, str, str]: + output = self._run( + [ + self.nvidia_smi, + "--query-gpu=name,uuid,driver_version", + "--format=csv,noheader,nounits", + f"--id={gpu_uuid}", + ] + ).stdout.strip() + name, uuid, driver = (item.strip() for item in output.split(",", 2)) + return name, uuid, driver + + def _artifact_bytes(self, workdir: Path) -> int: + return sum(path.stat().st_size for path in workdir.rglob("*") if path.is_file()) + + def _environment_identity(self) -> tuple[tuple[str, str], ...]: + prefixes = ("VLLM_", "CUDA_", "NCCL_", "TORCH_", "TRITON_") + environment = os.environ.copy() + environment.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + selected = tuple( + sorted( + (key, hashlib.sha256(key.encode() + b"\0" + value.encode()).hexdigest()) + for key, value in environment.items() + if key.startswith(prefixes) + and key not in {"VLLM_API_KEY", "VLLM_SNAPSHOT_TIMEOUT_S"} + ) + ) + return selected + self._plugin_identity() + + def make_manifest( + self, + args: argparse.Namespace, + engine_argv: tuple[str, ...], + inventory: ProcessInventory, + oracle: Oracle, + workdir: Path, + ) -> SnapshotManifest: + identity = self.current_identity(inventory.gpu_uuid) + redacted_argv = list(engine_argv) + redact_api_keys = False + for index, argument in enumerate(redacted_argv): + option, separator, _ = argument.partition("=") + normalized_option = option.replace("_", "-") + # Parsed argv can contain unambiguous long-option abbreviations. + if len(normalized_option) > 2 and "--api-key".startswith(normalized_option): + redact_api_keys = not separator + if separator: + redacted_argv[index] = f"{option}=***" + elif len(normalized_option) > 2 and "--hf-token".startswith( + normalized_option + ): + if separator: + redacted_argv[index] = f"{option}=***" + elif index + 1 < len(redacted_argv): + redacted_argv[index + 1] = "***" + redact_api_keys = False + elif argument.startswith("-"): + redact_api_keys = False + elif redact_api_keys: + redacted_argv[index] = "***" + engine_argv = tuple(redacted_argv) + revision = str(getattr(args, "revision", None) or "") + tokenizer_revision = str(getattr(args, "tokenizer_revision", None) or revision) + source_model = str(getattr(args, "model_tag", None) or args.model) + from vllm.config.model import get_served_model_name + + served_model_name = get_served_model_name( + source_model, getattr(args, "served_model_name", None) + ) + return SnapshotManifest( + schema_version=1, + boundary="post-engine-init-reloadable-state-released", + created_at=datetime.now(timezone.utc).isoformat(), + artifact_bytes=self._artifact_bytes(workdir), + **identity.model_dump(), + model=source_model, + served_model_name=str(served_model_name), + model_revision=revision, + tokenizer_revision=tokenizer_revision, + engine_argv=engine_argv, + process_tree=inventory.process_tree, + cuda_holders=inventory.cuda_holders, + external_cache_files=inventory.external_cache_files, + oracle_token_ids=oracle.token_ids, + oracle_text=oracle.text, + oracle_sampled_token_logprob=oracle.sampled_token_logprob, + ) + + def publish(self, workdir: Path, manifest: SnapshotManifest) -> None: + write_manifest_atomic(workdir, manifest) + _fsync_directory(workdir.parent) + + def current_identity(self, gpu_uuid: str) -> SnapshotRuntimeIdentity: + gpu_name, gpu_uuid, driver_version = self._gpu_identity(gpu_uuid) + torch_version, cuda_runtime = self._torch_identity() + host_id = Path("/etc/machine-id").read_text().strip() + return SnapshotRuntimeIdentity( + vllm_version=importlib.metadata.version("vllm"), + python_version=platform.python_version(), + torch_version=torch_version, + cuda_runtime=cuda_runtime, + driver_version=driver_version, + criu_version=self._version([self.criu, "--version"]), + cuda_checkpoint_sha256=self._sha256(Path(self.cuda_checkpoint)), + kernel_release=platform.release(), + host_id=host_id, + gpu_name=gpu_name, + gpu_uuid=gpu_uuid, + environment=self._environment_identity(), + ) + + def _read_restored_pid(self, pidfile: Path) -> int: + try: + payload = self._run([*self._privileged(), "cat", str(pidfile)]).stdout + except Exception as error: + raise SnapshotRestoreError("restored PID file is missing") from error + match = re.fullmatch(r"([1-9][0-9]{0,9})\n?", payload) + if match is None: + raise SnapshotRestoreError("restored PID file is invalid") + restored_pid = int(match.group(1)) + if restored_pid > 2**31 - 1: + raise SnapshotRestoreError("restored PID file is invalid") + return restored_pid + + def _process_state(self, pid: int) -> tuple[int, int, int, int]: + payload = (Path("/proc") / str(pid) / "stat").read_text() + command_end = payload.rfind(")") + fields = payload[command_end + 2 :].split() + if command_end < 0 or len(fields) < 20: + raise ValueError(f"invalid process state for PID {pid}") + return (int(fields[1]), int(fields[2]), int(fields[3]), int(fields[19])) + + def _process_states(self) -> dict[int, tuple[int, int, int, int]]: + states: dict[int, tuple[int, int, int, int]] = {} + for entry in Path("/proc").iterdir(): + if not entry.name.isdigit(): + continue + pid = int(entry.name) + try: + states[pid] = self._process_state(pid) + except (OSError, ValueError): + continue + return states + + def _process_command(self, pid: int) -> tuple[str, ...]: + payload = (Path("/proc") / str(pid) / "cmdline").read_bytes() + return tuple( + os.fsdecode(argument) for argument in payload.split(b"\0") if argument + ) + + def _pidfd_open(self, pid: int) -> int: + pidfd_open = getattr(os, "pidfd_open", None) + if not callable(pidfd_open): + raise SnapshotRestoreError("snapshot requires Linux pidfd support") + return pidfd_open(pid) + + def _pidfd_exited(self, pidfd: int) -> bool: + poller = select.poll() + poller.register(pidfd, select.POLLIN | select.POLLHUP | select.POLLERR) + return any(self._pidfd_event_is_exit(event) for _fd, event in poller.poll(0)) + + def _pidfd_event_is_exit(self, event: int) -> bool: + if event & (select.POLLERR | select.POLLNVAL): + raise SnapshotRestoreError("restored process pidfd poll failed") + return bool(event & (select.POLLIN | select.POLLHUP)) + + def _pin_restored_tree( + self, + artifact: Path, + process_tree: tuple[int, ...], + cuda_holders: tuple[int, ...], + ) -> None: + root_pid = process_tree[0] + expected_pids = set(process_tree) + expected_cuda_holders = set(cuda_holders) + if not expected_cuda_holders or not expected_cuda_holders.issubset( + expected_pids + ): + raise SnapshotRestoreError( + "captured CUDA holders are outside the captured process tree" + ) + if root_pid in self._restored_processes: + raise SnapshotRestoreError("restored PID already has a pinned transaction") + handles: list[tuple[int, int]] = [] + try: + for pid in process_tree: + handles.append((pid, self._pidfd_open(pid))) + handles_by_pid = dict(handles) + if any(self._pidfd_exited(pidfd) for _pid, pidfd in handles): + raise SnapshotRestoreError( + "restored process exited during verification" + ) + + root_state = self._process_state(root_pid) + if root_state[1:3] != (root_pid, root_pid): + raise SnapshotRestoreError( + "restored root is not the session and process-group leader" + ) + command = self._process_command(root_pid) + release_file = str(artifact / "release.json") + if not any( + command[index : index + 2] == ("-m", "vllm.snapshot.server") + for index in range(len(command) - 1) + ) or not any( + command[index : index + 2] == ("--release-file", release_file) + for index in range(len(command) - 1) + ): + raise SnapshotRestoreError( + "restored root command does not match the snapshot barrier" + ) + if self._process_state(root_pid) != root_state: + raise SnapshotRestoreError("restored PID changed during verification") + if self._pidfd_exited(handles[0][1]): + raise SnapshotRestoreError("restored root exited during verification") + + first_states = self._process_states() + live_session = { + pid for pid, state in first_states.items() if state[2] == root_pid + } + if not expected_cuda_holders.issubset(live_session): + raise SnapshotRestoreError( + "restored session is missing captured CUDA holders" + ) + if live_session != expected_pids: + raise SnapshotRestoreError( + "restored session does not match the captured process tree" + ) + if first_states[root_pid] != root_state: + raise SnapshotRestoreError("restored PID changed during verification") + for pid in process_tree[1:]: + state = first_states[pid] + pidfd = handles_by_pid[pid] + if self._process_state(pid) != state or self._pidfd_exited(pidfd): + raise SnapshotRestoreError( + "restored process tree changed during verification" + ) + + second_states = self._process_states() + second_session = { + pid for pid, state in second_states.items() if state[2] == root_pid + } + if second_session != expected_pids or any( + second_states[pid] != first_states[pid] for pid in process_tree + ): + raise SnapshotRestoreError( + "restored process tree changed during verification" + ) + if any(self._pidfd_exited(pidfd) for _pid, pidfd in handles): + raise SnapshotRestoreError( + "restored process exited during verification" + ) + except BaseException as error: + _close_pidfds(handles) + if isinstance(error, SnapshotRestoreError): + raise + if not isinstance(error, Exception): + raise + raise SnapshotRestoreError( + "restored process tree could not be pinned" + ) from error + self._restored_processes[root_pid] = tuple(handles) + + def _abort_failed_restore( + self, + artifact: Path, + root_pid: int, + staged_remaps: dict[str, tuple[int, int]], + primary: BaseException, + ) -> NoReturn: + secondary: list[str] = [] + try: + _write_json_atomic( + artifact / "release.json", {"release": False}, overwrite=True + ) + except BaseException as error: + secondary.append(f"abort marker failed: {_error_detail(error)}") + # The marker above is only a request. Terminate and wait for the tree we + # own, before releasing the shared state it may still be using. A tree + # that could not be pinned is not ours to signal. + if root_pid in self._restored_processes: + try: + self.cleanup(root_pid) + except BaseException as error: + secondary.append(f"restore cleanup failed: {_error_detail(error)}") + try: + self._rollback_link_remaps(staged_remaps) + except BaseException as error: + secondary.append(f"remap rollback failed: {_error_detail(error)}") + if secondary: + raise SnapshotRestoreError( + f"snapshot restore failed: {_error_detail(primary)}; " + + "; ".join(secondary) + ) from primary + raise primary + + def _verify_external_cache_files(self, manifest: SnapshotManifest) -> None: + """Reject a rotated generated cache before CRIU reopens it by path. + + Like the PID probe above this reports a stale host early rather than + guarding it: CRIU stays authoritative for a cache replaced afterwards. + """ + for path, fingerprint in manifest.external_cache_files: + target = Path(path) + if not target.is_file(): + raise SnapshotRestoreError( + f"external cache file vanished since capture: {path}; " + f"{_CACHE_REMEDY}" + ) + if fingerprint in (_CACHE_LOCKED, _CACHE_UNREADABLE): + continue + try: + current = self._sha256(target) + except OSError as error: + raise SnapshotRestoreError( + f"external cache file is unreadable since capture: {path}; " + f"{_CACHE_REMEDY}" + ) from error + if current != fingerprint: + raise SnapshotRestoreError( + f"external cache file changed since capture: {path}; " + f"{_CACHE_REMEDY}" + ) + + def restore(self, artifact: Path, manifest: SnapshotManifest) -> int: + # This is necessarily a best-effort TOCTOU check. CRIU restore and the + # rollback below remain authoritative if a PID is claimed after it. + for pid in manifest.process_tree: + try: + os.kill(pid, 0) + except ProcessLookupError: + continue + except PermissionError: + pass + except Exception as error: + raise SnapshotRestoreError( + f"captured PID availability probe failed: {pid}" + ) from error + raise SnapshotRestoreError(f"captured PID is already occupied: {pid}") + self._verify_external_cache_files(manifest) + expected_root_pid = manifest.process_tree[0] + (artifact / "release.json").unlink(missing_ok=True) + pidfile = artifact / "restored.pid" + pidfile.unlink(missing_ok=True) + self._reset_child_log(artifact) + staged_remaps = self._restore_link_remaps(artifact) + try: + self._criu( + "restore", + artifact, + [ + "--images-dir", + str(artifact / "images"), + "--log-file", + "restore.log", + *_COMMON_CRIU_OPTIONS, + "--restore-detached", + "--pidfile", + str(pidfile), + ], + ) + # Pin before reading the PID file so that a failure between restore + # and the read cannot orphan the tree. + self._pin_restored_tree( + artifact, manifest.process_tree, manifest.cuda_holders + ) + restored_pid = self._read_restored_pid(pidfile) + if restored_pid != expected_root_pid: + raise SnapshotRestoreError( + "restored PID does not match the captured process tree" + ) + except BaseException as error: + self._abort_failed_restore( + artifact, expected_root_pid, staged_remaps, error + ) + return restored_pid + + def release(self, artifact: Path, host: str | None, port: int) -> None: + if not 1 <= port <= 65535: + raise SnapshotRestoreError(f"snapshot restore port is invalid: {port}") + family = socket.AF_INET6 if host and ":" in host else socket.AF_INET + with socket.socket(family=family, type=socket.SOCK_STREAM) as listener_probe: + listener_probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + listener_probe.bind((host or "", port)) + except OSError as error: + address = f"{host or '0.0.0.0'}:{port}" + raise SnapshotRestoreError( + f"snapshot restore address is already in use: {address}" + ) from error + + _write_json_atomic( + artifact / "release.json", + {"release": True, "host": host, "port": port}, + overwrite=True, + ) + + def _connect_host(self, host: str | None) -> str: + if not host or host == "0.0.0.0": + return "127.0.0.1" + if host == "::": + return "::1" + return host + + def wait_listener(self, root_pid: int, host: str | None, port: int) -> None: + deadline = time.monotonic() + self.timeout_s + connect_host = self._connect_host(host) + while time.monotonic() < deadline: + try: + with socket.create_connection((connect_host, port), timeout=1): + return + except OSError as error: + handles = self._restored_processes.get(root_pid) + if handles is None or self._pidfd_exited(handles[0][1]): + raise SnapshotRestoreError( + "restored process exited before binding HTTP" + ) from error + time.sleep(0.1) + raise SnapshotRestoreError("restored HTTP listener did not become ready") + + def request_oracle( + self, host: str | None, port: int, manifest: SnapshotManifest + ) -> Oracle: + connect_host = self._connect_host(host) + url_host = f"[{connect_host}]" if ":" in connect_host else connect_host + payload = json.dumps( + { + "model": manifest.served_model_name, + "prompt": "The capital of France is", + "min_tokens": 1, + "max_tokens": 1, + "temperature": 0, + "seed": 0, + "logprobs": 0, + "return_token_ids": True, + } + ).encode() + request = urllib.request.Request( + f"http://{url_host}:{port}/v1/completions", + data=payload, + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=120) as response: + body = json.load(response) + choice = body["choices"][0] + try: + return Oracle( + token_ids=tuple(choice["token_ids"]), + text=choice["text"], + sampled_token_logprob=choice["logprobs"]["token_logprobs"][0], + ) + except ValueError as error: + raise SnapshotRestoreError( + "snapshot HTTP canary must return exactly one finite token logprob" + ) from error + except (IndexError, KeyError, TypeError) as error: + raise SnapshotRestoreError( + "snapshot HTTP canary response is missing sampled token logprob" + ) from error + + def cleanup(self, root_pid: int) -> None: + handles = self._restored_processes.get(root_pid) + if handles is None: + raise SnapshotRestoreError("no pinned restore transaction for cleanup") + pidfd_send_signal = getattr(signal, "pidfd_send_signal", None) + if not callable(pidfd_send_signal): + raise SnapshotRestoreError("restored process cleanup is incomplete") + try: + for _pid, pidfd in handles: + with suppress(ProcessLookupError): + pidfd_send_signal(pidfd, signal.SIGTERM, None, 0) + time.sleep(min(0.5, self.timeout_s)) + for _pid, pidfd in handles: + with suppress(ProcessLookupError): + pidfd_send_signal(pidfd, signal.SIGKILL, None, 0) + except BaseException as error: + raise SnapshotRestoreError( + f"restored process cleanup is incomplete: {_error_detail(error)}" + ) from error + + poller = select.poll() + for _pid, pidfd in handles: + poller.register(pidfd, select.POLLIN | select.POLLHUP | select.POLLERR) + pending = {pidfd for _pid, pidfd in handles} + deadline = time.monotonic() + self.timeout_s + while pending: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise SnapshotRestoreError("restored process cleanup is incomplete") + for pidfd, event in poller.poll(max(1, int(remaining * 1000))): + if not self._pidfd_event_is_exit(event): + continue + pending.discard(pidfd) + with suppress(KeyError, OSError): + poller.unregister(pidfd) + self._restored_processes.pop(root_pid) + _close_pidfds(handles) + + def complete_restore(self, root_pid: int) -> None: + handles = self._restored_processes.pop(root_pid, None) + if handles is None: + raise SnapshotRestoreError("no pinned restore transaction to complete") + _close_pidfds(handles) + + def abort_create(self, root_pid: int) -> None: + process = self._children.get(root_pid) + if process is None: + return + with suppress(ProcessLookupError): + os.killpg(root_pid, signal.SIGKILL) + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired as error: + raise SnapshotCreateError( + f"snapshot child process group {root_pid} survived SIGKILL" + ) from error + self._children.pop(root_pid, None) diff --git a/vllm/snapshot/server.py b/vllm/snapshot/server.py new file mode 100644 index 0000000000..2dc0f3384a --- /dev/null +++ b/vllm/snapshot/server.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import argparse +import asyncio +import json +import os +import stat +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from pydantic import ValidationError + +from vllm.snapshot.manifest import ReleaseMarker, _validation_path, _write_json_atomic +from vllm.snapshot.types import Oracle, oracles_match + +_CANARY_PROMPT = "The capital of France is" + + +class SnapshotCanaryError(RuntimeError): + """The initialized engine did not produce a valid snapshot oracle.""" + + +class SnapshotBarrierError(RuntimeError): + """The controller supplied an invalid snapshot release marker.""" + + +@dataclass(frozen=True) +class ListenerConfig: + host: str | None + port: int + + +@dataclass(frozen=True) +class ControlArgs: + ready_file: Path + release_file: Path + release_timeout_s: float + + +def write_ready_atomic(path: Path, oracle: Oracle) -> None: + _write_json_atomic( + path, + { + "sampled_token_logprob": oracle.sampled_token_logprob, + "token_ids": oracle.token_ids, + "text": oracle.text, + }, + ) + + +def read_release_marker(path: Path) -> ListenerConfig: + path = Path(path) + path_stat = path.lstat() + if stat.S_ISLNK(path_stat.st_mode): + raise SnapshotBarrierError("release marker must not be a symlink") + try: + payload = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError) as error: + raise SnapshotBarrierError("release marker is not valid JSON") from error + try: + marker = ReleaseMarker.model_validate(payload) + except ValidationError as error: + raise SnapshotBarrierError( + f"release marker is invalid: {_validation_path(error)}" + ) from error + return ListenerConfig(host=marker.host, port=marker.port) + + +def parse_control_args(argv: list[str]) -> tuple[ControlArgs, list[str]]: + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--ready-file", type=Path, required=True) + parser.add_argument("--release-file", type=Path, required=True) + parser.add_argument("--release-timeout-s", type=float, default=900.0) + control_args, remaining = parser.parse_known_args(argv) + if remaining and remaining[0] == "--": + remaining = remaining[1:] + if control_args.release_timeout_s <= 0: + raise ValueError("release timeout must be positive") + return ( + ControlArgs( + ready_file=control_args.ready_file, + release_file=control_args.release_file, + release_timeout_s=control_args.release_timeout_s, + ), + remaining, + ) + + +async def _release_reloadable_state(engine: Any) -> None: + """Discard model and KV state before the process image is captured.""" + await engine.sleep(level=2) + + +async def _restore_reloadable_state(engine: Any) -> None: + """Rebuild state discarded by ``_release_reloadable_state``.""" + await engine.wake_up(tags=["weights"]) + await engine.collective_rpc("reload_weights") + await engine.wake_up(tags=["kv_cache"]) + + +def oracle_from_request_output(request_output: Any) -> Oracle: + try: + candidate = request_output.outputs[0] + token_ids = tuple(candidate.token_ids) + (sampled_token_id,) = token_ids + return Oracle( + token_ids=token_ids, + text=candidate.text, + sampled_token_logprob=candidate.logprobs[0][sampled_token_id].logprob, + ) + except ValueError as error: + raise SnapshotCanaryError( + "snapshot canary must produce exactly one finite token logprob" + ) from error + except (AttributeError, IndexError, KeyError, TypeError) as error: + raise SnapshotCanaryError( + "snapshot canary did not return sampled token logprob" + ) from error + + +async def run_engine_canary(engine: Any) -> Oracle: + from vllm import SamplingParams + + final_output = None + sampling_params = SamplingParams( + temperature=0, + min_tokens=1, + max_tokens=1, + seed=0, + logprobs=0, + ) + async for output in engine.generate( + _CANARY_PROMPT, + sampling_params, + request_id="vllm-snapshot-canary", + ): + final_output = output + if final_output is None: + raise SnapshotCanaryError("canary generation returned no output") + return oracle_from_request_output(final_output) + + +def detach_snapshot_streams() -> None: + """Remove launch-log descriptors from the reusable process image.""" + sys.stdout.flush() + sys.stderr.flush() + sink = os.open(os.devnull, os.O_WRONLY) + try: + os.dup2(sink, sys.stdout.fileno()) + os.dup2(sink, sys.stderr.fileno()) + finally: + os.close(sink) + + +async def wait_for_release_marker( + path: Path, + *, + timeout_s: float, + poll_interval_s: float = 0.05, +) -> ListenerConfig: + if timeout_s <= 0: + raise ValueError("release timeout must be positive") + if poll_interval_s <= 0: + raise ValueError("release poll interval must be positive") + remaining_s = timeout_s + while not path.exists(): + if remaining_s <= 0: + raise TimeoutError(f"release marker not found before timeout: {path}") + delay_s = min(poll_interval_s, remaining_s) + await asyncio.sleep(delay_s) + remaining_s -= delay_s + return read_release_marker(path) + + +def parse_vllm_args(argv: list[str]) -> Any: + from vllm.entrypoints.launchers.cli_args import ( + make_arg_parser, + validate_parsed_serve_args, + ) + from vllm.utils.argparse_utils import FlexibleArgumentParser + + parser = FlexibleArgumentParser(prog="python -m vllm.snapshot.server") + args = make_arg_parser(parser).parse_args(argv) + if args.model_tag is not None: + args.model = args.model_tag + if args.grpc or args.headless: + raise ValueError("snapshot server requires the HTTP frontend") + if args.api_server_count not in (None, 1): + raise ValueError("snapshot server supports one API server") + args.api_server_count = None + validate_parsed_serve_args(args) + return args + + +async def run_vllm_snapshot_child(control: ControlArgs, args: Any) -> None: + from vllm.entrypoints.launchers.api_server.entry import ( + build_and_serve, + build_async_engine_client, + ) + from vllm.entrypoints.launchers.launcher import setup_server + + args.enable_sleep_mode = True + + async with build_async_engine_client(args) as engine: + oracle = await run_engine_canary(engine) + try: + await _release_reloadable_state(engine) + await _restore_reloadable_state(engine) + rehearsal_oracle = await run_engine_canary(engine) + if not oracles_match(oracle, rehearsal_oracle): + raise SnapshotCanaryError("snapshot rehearsal changed canary output") + await _release_reloadable_state(engine) + except SnapshotCanaryError: + raise + except Exception as error: + raise SnapshotCanaryError("snapshot rehearsal failed") from error + + detach_snapshot_streams() + write_ready_atomic(control.ready_file, oracle) + listener = await wait_for_release_marker( + control.release_file, + timeout_s=control.release_timeout_s, + ) + args.host = listener.host + args.port = listener.port + await _restore_reloadable_state(engine) + listen_address, sock = setup_server(args, reuse_port=False) + try: + shutdown_task = await build_and_serve( + engine, + listen_address, + sock, + args, + ) + await shutdown_task + finally: + sock.close() + + +def main(argv: list[str] | None = None) -> None: + control, remaining = parse_control_args(sys.argv[1:] if argv is None else argv) + args = parse_vllm_args(remaining) + import uvloop + + uvloop.run(run_vllm_snapshot_child(control, args)) + + +if __name__ == "__main__": + main() diff --git a/vllm/snapshot/types.py b/vllm/snapshot/types.py new file mode 100644 index 0000000000..236ba12436 --- /dev/null +++ b/vllm/snapshot/types.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +from dataclasses import dataclass + + +@dataclass(frozen=True) +class Oracle: + token_ids: tuple[int, ...] + text: str + sampled_token_logprob: float + + def __post_init__(self) -> None: + if len(self.token_ids) != 1: + raise ValueError("snapshot oracle must contain exactly one token") + value = self.sampled_token_logprob + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + ): + raise ValueError("sampled token logprob must be finite") + object.__setattr__(self, "sampled_token_logprob", float(value)) + + +def oracles_match(expected: Oracle, actual: Oracle) -> bool: + return ( + expected.token_ids == actual.token_ids + and expected.text == actual.text + and math.isclose( + expected.sampled_token_logprob, + actual.sampled_token_logprob, + rel_tol=0.0, + abs_tol=1e-3, + ) + ) diff --git a/vllm/utils/argparse_utils.py b/vllm/utils/argparse_utils.py index d214da3402..b91b1f26fd 100644 --- a/vllm/utils/argparse_utils.py +++ b/vllm/utils/argparse_utils.py @@ -519,7 +519,7 @@ class FlexibleArgumentParser(ArgumentParser): config_args = self.load_config_file(file_path) # 0th index might be the sub command {serve,chat,complete,...} - # optionally followed by model_tag (only for serve) + # optionally followed by model_tag (serve or snapshot create) # followed by config args # followed by rest of cli args. # maintaining this order will enforce the precedence @@ -527,8 +527,11 @@ class FlexibleArgumentParser(ArgumentParser): if args[0].startswith("-"): # No sub command (e.g., api_server entry point) args = config_args + args[0:index] + args[index + 2 :] - elif args[0] == "serve": - model_in_cli = len(args) > 1 and not args[1].startswith("-") + elif args[0] == "serve" or args[:2] == ["snapshot", "create"]: + model_index = 1 if args[0] == "serve" else 2 + model_in_cli = len(args) > model_index and not args[model_index].startswith( + "-" + ) model_in_config = any(arg == "--model" for arg in config_args) if not model_in_cli and not model_in_config: @@ -540,15 +543,19 @@ class FlexibleArgumentParser(ArgumentParser): if model_in_cli: # Model specified as positional arg, keep CLI version args = ( - [args[0]] - + [args[1]] + args[: model_index + 1] + config_args - + args[2:index] + + args[model_index + 1 : index] + args[index + 2 :] ) else: # No model in CLI, use config if available - args = [args[0]] + config_args + args[1:index] + args[index + 2 :] + args = ( + args[:model_index] + + config_args + + args[model_index:index] + + args[index + 2 :] + ) else: args = [args[0]] + config_args + args[1:index] + args[index + 2 :]