@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 13058faa49

[Frontend] Add reusable TP1 initialized-engine snapshots (#51360)

Signed-off-by: Nils Matteson <nilsmatteson@icloud.com> Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: elehayym <52448798+Yuzu23@users.noreply.github.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Nils Matteson committed Sep 21, 2026 at 16:18 UTC 13058faa49f5e4a4d50e1ca562ba0c2fdd73c7bc
16 files changed +3773 -10
.buildkite/scripts/initialized-snapshot-e2e.sh new
+311
@@ -0,0 +1,311 @@
1 +#!/usr/bin/env bash
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +set -Eeuo pipefail
5 +readonly MODEL="Qwen/Qwen3-0.6B"
6 +readonly MODEL_REVISION="c1899de289a04d12100db370d81485cdf75e47ca"
7 +readonly EXPECTED_TOKEN_ID="12095"
8 +readonly EXPECTED_TEXT=" Paris"
9 +readonly CLEANUP_RESERVE_S=300
10 +RUN_LIMIT_S="${SNAPSHOT_E2E_RUN_LIMIT_S:-3300}"
11 +[[ "$RUN_LIMIT_S" =~ ^[1-9][0-9]{2,3}$ ]] \
12 + || { echo "SNAPSHOT_E2E_RUN_LIMIT_S must be a bounded decimal integer" >&2; exit 1; }
13 +(( RUN_LIMIT_S >= CLEANUP_RESERVE_S + 60 && RUN_LIMIT_S <= 3300 )) \
14 + || { echo "SNAPSHOT_E2E_RUN_LIMIT_S must be between 360 and 3300" >&2; exit 1; }
15 +readonly RUN_LIMIT_S
16 +readonly HOST_HF_CACHE="${SNAPSHOT_E2E_HOST_HF_CACHE:-}"
17 +HF_CACHE_MOUNT=()
18 +if [[ -n "$HOST_HF_CACHE" ]]; then
19 + [[ "$HOST_HF_CACHE" == /* && -d "$HOST_HF_CACHE" ]] \
20 + || { echo "SNAPSHOT_E2E_HOST_HF_CACHE must be an existing absolute directory" >&2; exit 1; }
21 + HF_CACHE_MOUNT=(--mount "type=bind,source=$HOST_HF_CACHE,target=/e2e/hf,readonly")
22 + PREPARE_ARTIFACT_ROOT='install -d -m 0700 /e2e'
23 +else
24 + PREPARE_ARTIFACT_ROOT='mkdir -m 0700 /e2e && mkdir -m 0700 /e2e/hf'
25 +fi
26 +readonly PREPARE_ARTIFACT_ROOT
27 +if (( $# != 2 )); then
28 + echo "usage: $0 IMAGE_TAG BUILDKITE_COMMIT" >&2
29 + exit 2
30 +fi
31 +readonly IMAGE_TAG="$1"
32 +readonly EXPECTED_COMMIT="$2"
33 +readonly DEADLINE_EPOCH="$(( $(date +%s) + RUN_LIMIT_S ))"
34 +RUN_ID="" CONTAINER_NAME="" CONTAINER_ID=""
35 +GPU_UUID="" PORT_ONE="" PORT_TWO=""
36 +LINK_REMAP_BASELINE=""
37 +SNAPSHOT_PIDS=()
38 +die() {
39 + echo "initialized snapshot E2E: $*" >&2
40 + exit 1
41 +}
42 +
43 +run() {
44 + local label="$1" maximum_s="$2" available limit
45 + shift 2
46 + available="$((DEADLINE_EPOCH - $(date +%s) - CLEANUP_RESERVE_S))"
47 + (( available > 0 )) || die "no time remains before the cleanup reserve"
48 + (( available < maximum_s )) && limit="$available" || limit="$maximum_s"
49 + echo "--- ${label} (timeout ${limit}s)" >&2
50 + timeout --signal=TERM --kill-after=30s "${limit}s" "$@"
51 +}
52 +
53 +link_remaps() {
54 + find /dev/shm -maxdepth 1 -name 'link_remap.*' \
55 + -printf '%p|%D|%i|%m|%s\n' | sort
56 +}
57 +
58 +gpu_pids() {
59 + nvidia-smi --id="$GPU_UUID" --query-compute-apps=pid \
60 + --format=csv,noheader,nounits | sed '/^[[:space:]]*$/d'
61 +}
62 +
63 +port_listeners() {
64 + ss -H -ltn "sport = :$1"
65 +}
66 +
67 +state_clean() {
68 + local pid port output
69 + for pid in "${SNAPSHOT_PIDS[@]}"; do
70 + [[ ! -e "/proc/$pid" ]] || return 1
71 + done
72 + output="$(gpu_pids)" || return 1
73 + [[ -z "$output" ]] || return 1
74 + output="$(link_remaps)" || return 1
75 + [[ "$output" == "$LINK_REMAP_BASELINE" ]] || return 1
76 + for port in "$PORT_ONE" "$PORT_TWO"; do
77 + [[ -n "$port" ]] || continue
78 + output="$(port_listeners "$port")" || return 1
79 + [[ -z "$output" ]] || return 1
80 + done
81 +}
82 +
83 +wait_clean() {
84 + local _attempt
85 + for ((_attempt = 0; _attempt < 60; _attempt++)); do
86 + state_clean && return 0
87 + sleep 1
88 + done
89 + return 1
90 +}
91 +
92 +container_names() {
93 + timeout 15s docker container ls --all --format '{{.Names}}'
94 +}
95 +
96 +cleanup() {
97 + local failed=0 names="" actual="" expected=""
98 + if [[ -n "$CONTAINER_NAME" ]]; then
99 + names="$(container_names 2>/dev/null)" || return 1
100 + if grep -Fxq "$CONTAINER_NAME" <<< "$names"; then
101 + actual="$(docker container inspect \
102 + --format '{{.Id}}|{{index .Config.Labels "ai.vllm.snapshot.e2e.run"}}' \
103 + "$CONTAINER_NAME" 2>/dev/null)" || return 1
104 + expected="${CONTAINER_ID:-${actual%%|*}}|$RUN_ID"
105 + if [[ "$actual" != "$expected" ]]; then
106 + echo "refusing to remove a container with mismatched identity" >&2
107 + failed=1
108 + else
109 + timeout 60s docker rm --force "$CONTAINER_NAME" >/dev/null \
110 + || failed=1
111 + fi
112 + fi
113 + names="$(container_names 2>/dev/null)" || return 1
114 + grep -Fxq "$CONTAINER_NAME" <<< "$names" && failed=1
115 + fi
116 + if [[ -n "$PORT_TWO" ]] && ! wait_clean; then
117 + echo "process, GPU, port, or link-remap residue remained" >&2
118 + failed=1
119 + fi
120 + return "$failed"
121 +}
122 +
123 +on_exit() {
124 + local status=$?
125 + trap - EXIT INT TERM
126 + set +e
127 + cleanup || status=1
128 + exit "$status"
129 +}
130 +
131 +trap on_exit EXIT
132 +trap 'exit 130' INT
133 +trap 'exit 143' TERM
134 +
135 +parse_inspect() {
136 + timeout 30s python3 -c '
137 +import json, math, sys
138 +data = json.load(sys.stdin)
139 +model, revision, gpu_uuid, token_id, text = sys.argv[1:]
140 +expected = {
141 + "boundary": "post-engine-init-reloadable-state-released",
142 + "model": model, "model_revision": revision,
143 + "gpu_uuid": gpu_uuid, "oracle_token_ids": [int(token_id)],
144 + "oracle_text": text,
145 +}
146 +bad = next((field for field, value in expected.items() if data.get(field) != value), None)
147 +if bad:
148 + raise SystemExit(f"inspect field {bad} mismatch: {data.get(bad)!r}")
149 +logprob = data.get("oracle_sampled_token_logprob")
150 +if type(logprob) is not float or not math.isfinite(logprob):
151 + raise SystemExit("inspect logprob is not finite")
152 +pids = data.get("process_tree")
153 +if not pids:
154 + raise SystemExit("inspect process_tree is empty")
155 +for pid in pids:
156 + print(f"pid={pid}")
157 +print(f"oracle_logprob={logprob}")
158 +' "$MODEL" "$MODEL_REVISION" "$GPU_UUID" "$EXPECTED_TOKEN_ID" \
159 + "$EXPECTED_TEXT"
160 +}
161 +
162 +assert_oracle() {
163 + local port="$1" ordinal="$2" response
164 + response="$(run "restore ${ordinal} public completion" 120 curl \
165 + --fail --silent --show-error --max-time 120 \
166 + --header 'Content-Type: application/json' \
167 + --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}" \
168 + "http://127.0.0.1:${port}/v1/completions")"
169 + printf '%s' "$response" | timeout 30s python3 -c '
170 +import json, math, sys
171 +choice = json.load(sys.stdin)["choices"][0]
172 +token_id, text, expected, ordinal = sys.argv[1:]
173 +values = choice["logprobs"]["token_logprobs"]
174 +actual = float(values[0]) if len(values) == 1 else math.nan
175 +if (choice["token_ids"] != [int(token_id)] or choice["text"] != text
176 + or not math.isfinite(actual) or abs(actual - float(expected)) > 1e-3):
177 + raise SystemExit(f"public oracle mismatch: {choice!r}")
178 +print(f"restore={ordinal} token_id={token_id} text={text!r} logprob={actual}")
179 +' "$EXPECTED_TOKEN_ID" "$EXPECTED_TEXT" "$INSPECT_LOGPROB" "$ordinal"
180 +}
181 +
182 +[[ "$(uname -s)" == "Linux" && "$(uname -m)" == "x86_64" ]] \
183 + || die "host must be Linux x86_64"
184 +[[ "$EXPECTED_COMMIT" =~ ^[0-9a-f]{40}$ ]] \
185 + || die "BUILDKITE_COMMIT must be a full lowercase commit hash"
186 +[[ "$(< /proc/sys/kernel/io_uring_disabled)" == "2" ]] \
187 + || die "kernel.io_uring_disabled must already equal 2"
188 +run "Docker preflight" 60 docker info >/dev/null
189 +
190 +VISIBLE_GPU="${NVIDIA_VISIBLE_DEVICES:-}"
191 +if [[ -z "$VISIBLE_GPU" || "$VISIBLE_GPU" == "all" || "$VISIBLE_GPU" == "void" ]]; then
192 + VISIBLE_GPU="${CUDA_VISIBLE_DEVICES:-}"
193 +fi
194 +if [[ -z "$VISIBLE_GPU" || "$VISIBLE_GPU" == "all" || "$VISIBLE_GPU" == "void" ]]; then
195 + CDI_GPU="${BUILDKITE_PLUGIN_DOCKER_DEVICES_0:-}"
196 + if [[ "$CDI_GPU" == nvidia.com/gpu=* ]]; then
197 + VISIBLE_GPU="${CDI_GPU#nvidia.com/gpu=}"
198 + else
199 + VISIBLE_GPU=""
200 + fi
201 +fi
202 +[[ -n "$VISIBLE_GPU" && "$VISIBLE_GPU" != all && "$VISIBLE_GPU" != void \
203 + && "$VISIBLE_GPU" != *","* ]] \
204 + || die "the runner must expose exactly one selected GPU"
205 +GPU_ROW="$(nvidia-smi --id="$VISIBLE_GPU" \
206 + --query-gpu=name,uuid,mig.mode.current,memory.total \
207 + --format=csv,noheader,nounits)"
208 +IFS=',' read -r GPU_NAME GPU_UUID MIG_MODE GPU_MEMORY_MIB <<< "$GPU_ROW"
209 +for variable in GPU_NAME GPU_UUID MIG_MODE GPU_MEMORY_MIB; do
210 + printf -v "$variable" '%s' "${!variable#"${!variable%%[![:space:]]*}"}"
211 +done
212 +[[ "$GPU_NAME" == *H200* && "$GPU_UUID" == GPU-* && "$MIG_MODE" == Disabled ]] \
213 + || die "selected device is not a full non-MIG H200: $GPU_ROW"
214 +(( GPU_MEMORY_MIB >= 130000 )) || die "selected H200 does not expose full memory"
215 +GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU"
216 +[[ -z "$GPU_PIDS" ]] || die "selected H200 is not idle"
217 +
218 +RUN_ID="snapshot-e2e-$(date +%s)-$$-$RANDOM"
219 +CONTAINER_NAME="vllm-$RUN_ID"
220 +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))')
221 +LINK_REMAP_BASELINE="$(link_remaps)"
222 +
223 +run "pull exact candidate image" 900 docker pull --platform linux/amd64 \
224 + "$IMAGE_TAG" >/dev/null
225 +IFS='|' read -r IMAGE_ID IMAGE_PLATFORM OCI_COMMIT VLLM_COMMIT < <(
226 + docker image inspect \
227 + --format '{{.Id}}|{{.Os}}/{{.Architecture}}|{{index .Config.Labels "org.opencontainers.image.revision"}}|{{index .Config.Labels "ai.vllm.build.commit"}}' \
228 + "$IMAGE_TAG"
229 +)
230 +[[ "$IMAGE_ID" =~ ^sha256:[0-9a-f]{64}$ && "$IMAGE_PLATFORM" == linux/amd64 ]] \
231 + || die "candidate image identity is invalid"
232 +[[ "$OCI_COMMIT" == "$EXPECTED_COMMIT" && "$VLLM_COMMIT" == "$EXPECTED_COMMIT" ]] \
233 + || die "candidate image commit labels do not match"
234 +CONTAINER_ID="$(run "start exact candidate container" 90 docker run --detach \
235 + --name "$CONTAINER_NAME" --label "ai.vllm.snapshot.e2e.run=$RUN_ID" \
236 + --gpus "device=$GPU_UUID" --user 0 --privileged --pid=host --ipc=host \
237 + --network=host --env "CUDA_VISIBLE_DEVICES=$GPU_UUID" --env HF_HUB_DISABLE_TELEMETRY=1 \
238 + --env VLLM_NO_USAGE_STATS=1 --env VLLM_USE_V2_MODEL_RUNNER=1 \
239 + ${HF_CACHE_MOUNT[@]+"${HF_CACHE_MOUNT[@]}"} \
240 + --entrypoint sleep "$IMAGE_ID" infinity)"
241 +[[ "$CONTAINER_ID" =~ ^[0-9a-f]{64}$ ]] || die "invalid container ID"
242 +run "prepare private artifact root" 60 docker exec "$CONTAINER_NAME" sh -c \
243 + "$PREPARE_ARTIFACT_ROOT && df -Pk / | awk 'NR == 2 && \$4 >= 10485760 {ok=1} END {exit !ok}'"
244 +run "verify snapshot runtime" 60 docker exec "$CONTAINER_NAME" sh -c \
245 + 'criu --version; test -x /usr/local/sbin/cuda-checkpoint; test -f /usr/local/lib/criu/cuda_plugin.so'
246 +run "verify selected CUDA device" 60 docker exec "$CONTAINER_NAME" python3 -c '
247 +import sys, torch
248 +if torch.cuda.device_count() != 1:
249 + raise SystemExit("expected exactly one visible CUDA device")
250 +actual = str(torch.cuda.get_device_properties(0).uuid).removeprefix("GPU-")
251 +if actual != sys.argv[1].removeprefix("GPU-"):
252 + raise SystemExit(f"CUDA device {actual} does not match {sys.argv[1]}")
253 +print(f"selected_cuda_uuid=GPU-{actual}")
254 +' "$GPU_UUID"
255 +OFFLINE_EXEC=(docker exec --env HF_HOME=/e2e/hf --env HF_HUB_OFFLINE=1 \
256 + --env TRANSFORMERS_OFFLINE=1 --env VLLM_SNAPSHOT_TIMEOUT_S=900 \
257 + "$CONTAINER_NAME")
258 +if [[ -n "$HOST_HF_CACHE" ]]; then
259 + run "validate pinned model in read-only host cache" 120 \
260 + "${OFFLINE_EXEC[@]}" python3 -c \
261 + '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]' \
262 + "$MODEL" "$MODEL_REVISION"
263 +else
264 + run "prefetch pinned public model" 900 docker exec --env HF_HOME=/e2e/hf \
265 + "$CONTAINER_NAME" python3 -c \
266 + '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]' \
267 + "$MODEL" "$MODEL_REVISION"
268 +fi
269 +
270 +run "create compact initialized snapshot" 1200 "${OFFLINE_EXEC[@]}" \
271 + vllm snapshot create "$MODEL" --revision "$MODEL_REVISION" \
272 + --tokenizer-revision "$MODEL_REVISION" --snapshot-dir /e2e/artifact \
273 + --dtype float16 --max-model-len 512 --gpu-memory-utilization 0.50
274 +wait_clean || die "create left process, GPU, port, or link-remap residue"
275 +
276 +INSPECT_JSON="$(run "inspect compact initialized snapshot" 120 \
277 + "${OFFLINE_EXEC[@]}" vllm snapshot inspect /e2e/artifact)"
278 +INSPECT_RECEIPT="$(printf '%s' "$INSPECT_JSON" | parse_inspect)"
279 +mapfile -t SNAPSHOT_PIDS < <(sed -n 's/^pid=//p' <<< "$INSPECT_RECEIPT")
280 +INSPECT_LOGPROB="$(sed -n 's/^oracle_logprob=//p' <<< "$INSPECT_RECEIPT")"
281 +(( ${#SNAPSHOT_PIDS[@]} > 0 )) || die "inspect returned no PIDs"
282 +state_clean || die "create left recorded state alive"
283 +echo "$INSPECT_RECEIPT"
284 +
285 +run "restore compact snapshot 1" 900 "${OFFLINE_EXEC[@]}" \
286 + vllm snapshot restore /e2e/artifact --host 127.0.0.1 --port "$PORT_ONE"
287 +PORT_LISTENERS="$(port_listeners "$PORT_ONE")" || die "failed to inspect restore 1 port"
288 +[[ -n "$PORT_LISTENERS" ]] || die "restore 1 did not bind its port"
289 +GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU"
290 +[[ -n "$GPU_PIDS" ]] || die "restore 1 did not engage the selected GPU"
291 +assert_oracle "$PORT_ONE" 1
292 +run "stop container after restore 1" 60 docker stop --time 20 "$CONTAINER_NAME" >/dev/null
293 +[[ "$(docker container inspect --format '{{.Id}}' "$CONTAINER_NAME")" == "$CONTAINER_ID" ]] \
294 + || die "docker stop replaced the container"
295 +wait_clean || die "restore 1 teardown left residue"
296 +
297 +run "start the same container" 60 docker start "$CONTAINER_NAME" >/dev/null
298 +[[ "$(docker container inspect --format '{{.Id}}' "$CONTAINER_NAME")" == "$CONTAINER_ID" ]] \
299 + || die "docker start replaced the container"
300 +run "restore compact snapshot 2" 900 "${OFFLINE_EXEC[@]}" \
301 + vllm snapshot restore /e2e/artifact --host 127.0.0.1 --port "$PORT_TWO"
302 +PORT_LISTENERS="$(port_listeners "$PORT_TWO")" || die "failed to inspect restore 2 port"
303 +[[ -n "$PORT_LISTENERS" ]] || die "restore 2 did not bind its port"
304 +GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU"
305 +[[ -n "$GPU_PIDS" ]] || die "restore 2 did not engage the selected GPU"
306 +assert_oracle "$PORT_TWO" 2
307 +run "stop container after restore 2" 60 docker stop --time 20 "$CONTAINER_NAME" >/dev/null
308 +wait_clean || die "restore 2 teardown left residue"
309 +cleanup || die "exact cleanup failed"
310 +echo "container_id=$CONTAINER_ID restore_1_port=$PORT_ONE restore_2_port=$PORT_TWO"
311 +echo "result=PASS"
.buildkite/test_areas/entrypoints.yaml
+20
@@ -2,6 +2,26 @@ group: Entrypoints
2 depends_on:
3 - image-build
4 steps:
5 +- label: ":nvidia: (H200) Initialized Snapshot E2E"
6 + key: initialized-snapshot-e2e
7 + device: h200
8 + num_devices: 1
9 + no_plugin: true
10 + optional: true
11 + timeout_in_minutes: 60
12 + working_dir: "."
13 + env:
14 + SNAPSHOT_E2E_RUN_LIMIT_S: "1200"
15 + source_file_dependencies:
16 + - .buildkite/scripts/initialized-snapshot-e2e.sh
17 + - .buildkite/test_areas/entrypoints.yaml
18 + - docker/Dockerfile
19 + - tools/install_snapshot_runtime.sh
20 + - vllm/entrypoints/cli/
21 + - vllm/snapshot/
22 + commands:
23 + - bash .buildkite/scripts/initialized-snapshot-e2e.sh "$IMAGE_TAG" "$BUILDKITE_COMMIT"
24 +
25 - label: ":nvidia: (H200 MIG 35GB) Entrypoints Unit"
26 device: h200_35gb
27 key: entrypoints-unit-tests
docker/Dockerfile
+16
@@ -735,6 +735,8 @@ RUN --mount=type=cache,target=/opt/uv/cache \
735 FROM ${FINAL_BASE_IMAGE} AS vllm-runtime-base
736
737 ARG INSTALL_RUBIN_PRERELEASE
738 +ARG TARGETPLATFORM
739 +ARG max_jobs=2
740 ARG CUDA_VERSION
741 ARG PYTHON_VERSION
742 ARG NCCL_VERSION
@@ -854,6 +856,20 @@ RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
856 && ldconfig; \
857 fi
858
859 +# Install the initialized-snapshot runtime on the architecture supported by
860 +# `vllm snapshot`. CRIU and its CUDA plugin are built from the same pinned
861 +# source. The `cuda-checkpoint`-compatible helper is built from NVIDIA Dynamo's
862 +# Apache-2.0 source against the public CUDA Driver checkpoint APIs instead of
863 +# redistributing NVIDIA's proprietary cuda-checkpoint binary.
864 +COPY tools/install_snapshot_runtime.sh /usr/local/bin/install_snapshot_runtime.sh
865 +RUN if [ "${TARGETPLATFORM}" = "linux/amd64" ]; then \
866 + MAX_JOBS="${max_jobs}" /usr/local/bin/install_snapshot_runtime.sh; \
867 + else \
868 + echo "Initialized snapshots are not installed for ${TARGETPLATFORM}"; \
869 + fi
870 +
871 +ENV CRIU_CUDA_PLUGIN_DIR=/usr/local/lib/criu
872 +
873 # Install uv for faster pip installs
874 RUN python3 -m pip install uv
875
docs/features/initialized_snapshots.md new
+152
@@ -0,0 +1,152 @@
1 +# Initialized engine snapshots
2 +
3 +Initialized engine snapshots are an experimental way to trade local disk space
4 +and host privileges for a faster vLLM activation. Snapshot creation initializes
5 +the engine and records deterministic generation output before CRIU captures the
6 +process tree and CUDA state. Restore validates the saved environment, restores
7 +the engine, binds the HTTP server, and reproduces the recorded token and sampled
8 +token log probability before returning.
9 +
10 +This path is intended for repeatedly activating the same model and engine
11 +configuration on the same machine. It is not a portable model artifact.
12 +
13 +## Requirements
14 +
15 +Snapshots currently require:
16 +
17 +- Linux on x86-64 with one NVIDIA GPU.
18 +- TP1 with one unauthenticated plaintext HTTP server. Other parallel sizes,
19 + TLS, middleware, Unix sockets, and speculative decoding are unsupported.
20 +- Because snapshot mode requires an unauthenticated plaintext HTTP server, use
21 + a trusted host or network boundary, or external controls. See the vLLM
22 + [Security guide](../usage/security.md).
23 +- [CRIU](https://github.com/checkpoint-restore/criu), its CUDA plugin, a
24 + `cuda-checkpoint`-compatible helper, and `nvidia-smi` on `PATH`.
25 +- Root or passwordless `sudo` for CRIU.
26 +- `io_uring` disabled before launch because CRIU cannot dump it. Use
27 + `kernel.io_uring_disabled=1` for an unprivileged process. A process running
28 + as root, including the `docker exec` flow below, bypasses `=1`, so use `=2`
29 + host-wide.
30 +- No established TCP connection to a peer outside the captured process tree.
31 + Current Hugging Face hub clients hold their connections for the process
32 + lifetime, so download the model in a separate step and run create with
33 + `HF_HUB_OFFLINE=1`, as the quickstart below does.
34 +- A remote model ID and an immutable 40-character `--revision`. Local model
35 + directories and mutable revisions are not supported.
36 +- Enough disk for the artifact, with the same installed vLLM package, model
37 + files, container filesystem, and generated-cache paths available at restore.
38 + Generated-cache files are not copied into the artifact. The manifest
39 + fingerprints the ones the captured tree holds open, so restore fails early
40 + and names the file when one is removed or replaced. Only open descriptors
41 + are recorded. A library the tree mapped and then closed is not fingerprinted,
42 + and CRIU still reopens it by path, so replacing that file permanently
43 + invalidates the artifact without an early error.
44 +
45 +The official `vllm/vllm-openai` Linux x86-64 image includes the snapshot
46 +runtime. It still requires a compatible host driver, kernel, and privileges.
47 +Arm64 images omit it. Source installs must set `CRIU_CUDA_PLUGIN_DIR` to the
48 +directory containing `cuda_plugin.so`.
49 +
50 +Run snapshot commands with `docker exec` inside a long-lived container. Restore
51 +hands the API server off as a detached process, so a one-shot container would
52 +stop that server when its PID 1 exits. Snapshot preflight requires every
53 +component of the artifact path to be owned by the invoking user or root, with
54 +the directory itself at mode 0700, and the commands below run as root inside
55 +the official image, so the bind-mounted host directory is created with `sudo`
56 +and root ownership. The model downloads in its own step so the captured tree
57 +holds no hub connection, and create runs offline. This example also keeps the
58 +container filesystem and `/dev/shm` namespace stable for the lifetime of the
59 +artifact:
60 +
61 +```bash
62 +sudo sysctl kernel.io_uring_disabled=2
63 +
64 +snapshot_root="$(pwd)/vllm-snapshots"
65 +sudo install -d -m 0700 -o root -g root "${snapshot_root}"
66 +
67 +docker run --detach --name vllm-snapshot \
68 + --gpus all \
69 + --privileged \
70 + --pid=host \
71 + --ipc=host \
72 + --network=host \
73 + --mount "type=bind,source=${snapshot_root},target=/snapshots" \
74 + --entrypoint sleep \
75 + vllm/vllm-openai:latest infinity
76 +
77 +docker exec vllm-snapshot hf download Qwen/Qwen3-0.6B \
78 + --revision c1899de289a04d12100db370d81485cdf75e47ca
79 +
80 +docker exec -e HF_HUB_OFFLINE=1 vllm-snapshot vllm snapshot create Qwen/Qwen3-0.6B \
81 + --snapshot-dir /snapshots/qwen3-0.6b \
82 + --revision c1899de289a04d12100db370d81485cdf75e47ca \
83 + --dtype float16 \
84 + --max-model-len 512
85 +
86 +docker exec vllm-snapshot vllm snapshot inspect /snapshots/qwen3-0.6b
87 +
88 +docker exec vllm-snapshot vllm snapshot restore \
89 + /snapshots/qwen3-0.6b --host 0.0.0.0 --port 8000
90 +```
91 +
92 +Keep that container and its mounts available while the snapshot is in use.
93 +Stop and remove it only after the restored API server is no longer needed.
94 +
95 +Create initializes the engine, records a one-token canary, releases and reloads
96 +weights and KV cache to rehearse restore, then releases them again for capture.
97 +The manifest is published only after CRIU completes and the source tree stops.
98 +Creation is offline preparation and is not part of restore latency.
99 +
100 +The private `0700` artifact contains process memory, CUDA state, engine
101 +arguments, compatibility identity, and canary output; its manifest is `0600`.
102 +Literal API keys and Hugging Face tokens are redacted from the manifest's engine
103 +arguments. The manifest records selected environment names plus deterministic
104 +name-bound fingerprints, not their values. The protected CRIU artifact can
105 +still contain process secrets, so treat it as sensitive data. Restore reloads
106 +model files and KV cache before binding HTTP, then reproduces the canary or
107 +tears down the restored tree. The inspect command prints that identity and
108 +canary without executing the saved process.
109 +
110 +## Restore behavior
111 +
112 +Restore fails before CRIU runs if the saved identity does not match the current
113 +host. It does not silently fall back to ordinary startup. After CRIU restores
114 +the process tree, vLLM releases the saved engine to bind the requested HTTP
115 +address and checks the first generated token and sampled-token log probability
116 +against the snapshot canary. The command returns only after that check passes.
117 +The restored API server continues to run as a detached process.
118 +
119 +The pre-release port probe is best-effort only. It neither reserves the port
120 +nor authenticates the listener that appears afterward.
121 +
122 +Rollback terminates and waits for the restored tree after process identity
123 +verification succeeds. If verification fails, vLLM writes an abort marker for
124 +the snapshot server instead of signaling unverified PIDs. If that server has
125 +already exited, surviving engine processes can require operator cleanup.
126 +Verify that the previous tree has stopped before retrying; do not kill a process
127 +solely because its PID appears in the manifest, since PIDs can be reused.
128 +
129 +The artifact is reusable after its previous restored tree stops. Only one
130 +snapshot or external CRIU operation may use a shared `/dev/shm` mount at a time.
131 +
132 +## Tradeoffs and limitations
133 +
134 +- Creation has its own latency and briefly requires the full engine. Artifact
135 + size can approach the captured process and GPU memory.
136 +- Artifacts are not guaranteed to survive a power loss. The manifest is
137 + fsynced, but publication does not explicitly flush all captured payload files
138 + and directories. Recreate the artifact after an unclean host shutdown.
139 +- Restore currently requires the same host, GPU, driver, kernel, Python,
140 + PyTorch, installed vLLM version, model revision, engine arguments,
141 + selected environment variables, and CRIU plugin binaries.
142 +- Only dense float16 TP1 has been validated. Other model formats depend on their
143 + existing sleep level 2 reload support; NCCL state is not restored.
144 +- CRIU support varies by kernel and driver. Preserve package, library, model,
145 + and generated-cache paths for the artifact lifetime.
146 +- A snapshot can include application secrets or request state present in the
147 + process. Build it before traffic and protect it like process memory.
148 +- Any feature that opens an external connection must close it before capture.
149 +
150 +For a lower-complexity option that retains a live process, see
151 +[Sleep mode](sleep_mode.md). Sleep mode and initialized snapshots retain
152 +different amounts of state and have different idle resource costs.
tests/entrypoints/launchers/test_launch_cli.py
+1042 -1
@@ -3,15 +3,48 @@
3 """Unit tests for the `vllm launch` CLI subcommand."""
4
5 import argparse
6 -from unittest.mock import patch
6 +import hashlib
7 +import json
8 +import os
9 +import signal
10 +import subprocess
11 +import sys
12 +from contextlib import asynccontextmanager
13 +from pathlib import Path
14 +from typing import Any
15 +from unittest.mock import create_autospec, patch
16
17 import pytest
18
19 +import vllm.snapshot.runtime as snapshot_runtime
20 +import vllm.snapshot.server as snapshot_server
21 +from vllm.entrypoints.cli import snapshot as snapshot_cli
22 from vllm.entrypoints.cli.launch import (
23 LaunchSubcommand,
24 RenderSubcommand,
25 cmd_init,
26 )
27 +from vllm.snapshot.controller import (
28 + LocalSnapshotTools,
29 + SnapshotCreateError,
30 + SnapshotRestoreError,
31 + _child_engine_argv,
32 + create_snapshot,
33 + restore_snapshot,
34 +)
35 +from vllm.snapshot.manifest import (
36 + SnapshotCompatibilityError,
37 + SnapshotManifest,
38 + SnapshotRuntimeIdentity,
39 + SnapshotSecurityError,
40 + read_manifest,
41 + validate_artifact_root,
42 + validate_identity,
43 + write_manifest_atomic,
44 +)
45 +from vllm.snapshot.runtime import ProcessInventory, _decode_endpoint, _TcpSocketRecord
46 +from vllm.snapshot.server import SnapshotCanaryError
47 +from vllm.snapshot.types import Oracle, oracles_match
48 from vllm.utils.argparse_utils import FlexibleArgumentParser
49
50
@@ -109,3 +142,1011 @@ def test_launch_registered_in_main():
142 assert hasattr(launch_module, "cmd_init")
143 subcmds = launch_module.cmd_init()
144 assert any(s.name == "launch" for s in subcmds)
145 +
146 +
147 +_MODEL_REVISION = "c1899de289a04d12100db370d81485cdf75e47ca"
148 +
149 +
150 +def _run_cli(argv: list[str]) -> None:
151 + from vllm.entrypoints.cli import main as cli_main
152 +
153 + with patch.object(sys, "argv", ["vllm", *argv]):
154 + cli_main.main()
155 +
156 +
157 +def test_snapshot_environment_contract(
158 + tmp_path: Path,
159 + monkeypatch: pytest.MonkeyPatch,
160 + capsys: pytest.CaptureFixture[str],
161 +):
162 + from vllm.entrypoints.serve.utils import api_utils
163 +
164 + monkeypatch.setattr(api_utils, "cli_env_setup", pytest.fail)
165 + secret = "snapshot-secret"
166 + monkeypatch.setenv("VLLM_API_KEY", secret)
167 + monkeypatch.setenv("VLLM_USER_SETTING", secret)
168 + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", secret)
169 + environment = dict(LocalSnapshotTools()._environment_identity())
170 + assert (
171 + environment["VLLM_USER_SETTING"]
172 + == hashlib.sha256(b"VLLM_USER_SETTING\0snapshot-secret").hexdigest()
173 + )
174 + assert (
175 + environment["VLLM_USER_SETTING"] != environment["VLLM_WORKER_MULTIPROC_METHOD"]
176 + )
177 + assert secret not in str(environment)
178 + assert "VLLM_API_KEY" not in environment
179 +
180 + monkeypatch.delenv("VLLM_WORKER_MULTIPROC_METHOD")
181 + implicit_worker_default = dict(LocalSnapshotTools()._environment_identity())
182 + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
183 + assert dict(LocalSnapshotTools()._environment_identity()) == implicit_worker_default
184 +
185 + monkeypatch.setenv("VLLM_USER_SETTING", f"{secret}-changed")
186 + changed_environment = dict(LocalSnapshotTools()._environment_identity())
187 + assert changed_environment["VLLM_USER_SETTING"] != environment["VLLM_USER_SETTING"]
188 +
189 + artifact = tmp_path / "snapshot"
190 + artifact.mkdir(mode=0o700)
191 + write_manifest_atomic(artifact, _manifest(environment=tuple(environment.items())))
192 + before = os.environ.copy()
193 +
194 + _run_cli(["snapshot", "inspect", str(artifact)])
195 +
196 + assert os.environ == before
197 + output = capsys.readouterr().out
198 + assert secret not in output
199 + inspected = json.loads(output)
200 + inspected_environment = dict(inspected["environment"])
201 + assert inspected_environment == environment
202 +
203 + tools = LocalSnapshotTools()
204 + monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity())
205 + manifest = tools.make_manifest(
206 + argparse.Namespace(
207 + hf_token=secret,
208 + model_tag="Qwen/Qwen3-0.6B",
209 + revision=_MODEL_REVISION,
210 + served_model_name=None,
211 + tokenizer_revision=None,
212 + ),
213 + (
214 + "Qwen/Qwen3-0.6B",
215 + "--api-key",
216 + secret,
217 + f"{secret}-alternate",
218 + f"--api-key={secret}",
219 + "--hf-token",
220 + secret,
221 + f"--hf-token={secret}",
222 + ),
223 + ProcessInventory(100, (100, 101), (101,), "GPU-abc"),
224 + _oracle(),
225 + tmp_path,
226 + )
227 +
228 + assert manifest.engine_argv == (
229 + "Qwen/Qwen3-0.6B",
230 + "--api-key",
231 + "***",
232 + "***",
233 + "--api-key=***",
234 + "--hf-token",
235 + "***",
236 + "--hf-token=***",
237 + )
238 +
239 +
240 +@pytest.mark.parametrize(
241 + "case",
242 + [
243 + (("--hf_token", "SECRET"), ("--hf_token", "***")),
244 + (("--hf_token=SECRET",), ("--hf_token=***",)),
245 + (("--api_key", "SECRET"), ("--api_key", "***")),
246 + (("--api_key=SECRET",), ("--api_key=***",)),
247 + (("--hf-tok", "SECRET"), ("--hf-tok", "***")),
248 + (("--hf-tok=SECRET",), ("--hf-tok=***",)),
249 + (("--hf_tok", "SECRET"), ("--hf_tok", "***")),
250 + (("--api-k", "SECRET", "SECOND"), ("--api-k", "***", "***")),
251 + (("--api-k=SECRET",), ("--api-k=***",)),
252 + (("--hf-overrides", "{}"), ("--hf-overrides", "{}")),
253 + ],
254 +)
255 +def test_snapshot_manifest_redacts_accepted_secret_options(case):
256 + engine_argv, expected = case
257 + tools = create_autospec(LocalSnapshotTools, instance=True)
258 + tools.current_identity.return_value = _runtime_identity()
259 + tools._artifact_bytes.return_value = 0
260 + manifest = LocalSnapshotTools.make_manifest(
261 + tools,
262 + snapshot_server.parse_vllm_args(["model", *engine_argv]),
263 + engine_argv,
264 + ProcessInventory(100, (100, 101), (101,), "GPU-abc"),
265 + _oracle(),
266 + Path(),
267 + )
268 +
269 + assert manifest.engine_argv == expected
270 + if "SECRET" in " ".join(engine_argv):
271 + assert "***" in " ".join(manifest.engine_argv)
272 + assert "SECRET" not in " ".join(manifest.engine_argv)
273 +
274 +
275 +def parse_snapshot(*argv: str):
276 + parser = FlexibleArgumentParser()
277 + subparsers = parser.add_subparsers(dest="subparser", required=True)
278 + snapshot_cli.SnapshotSubcommand(
279 + create_requested=argv[:1] == ("create",)
280 + ).subparser_init(subparsers)
281 + return parser.parse_args(["snapshot", *argv])
282 +
283 +
284 +@pytest.mark.parametrize("option", ["--snapshot-dir", "--snapshot_dir", "--snapshot-d"])
285 +@pytest.mark.parametrize("joined", [False, True])
286 +def test_snapshot_create_forwards_only_engine_options(option, joined, monkeypatch):
287 + directory = [f"{option}=/tmp/snapshot"] if joined else [option, "/tmp/snapshot"]
288 + argv = ["create", "Qwen/Qwen3-0.6B", "--revision", _MODEL_REVISION, *directory]
289 + args = parse_snapshot(*argv)
290 + monkeypatch.setattr(sys, "argv", ["vllm", "snapshot", *argv])
291 +
292 + child = snapshot_server.parse_vllm_args(list(_child_engine_argv(args, None)))
293 +
294 + assert child.model == "Qwen/Qwen3-0.6B"
295 + assert child.revision == _MODEL_REVISION
296 + assert child.enable_sleep_mode
297 +
298 +
299 +@pytest.mark.parametrize("model_in_cli", [False, True])
300 +def test_snapshot_create_config_preserves_cli_precedence(
301 + tmp_path, monkeypatch, model_in_cli
302 +):
303 + config = tmp_path / "engine.yaml"
304 + config.write_text("model: config-model\nmax-model-len: 512\ndtype: float16\n")
305 + model = ["Qwen/Qwen3-0.6B"] if model_in_cli else []
306 + argv = [
307 + "create",
308 + *model,
309 + "--snapshot-dir",
310 + str(tmp_path / "snapshot"),
311 + "--config",
312 + str(config),
313 + "--revision",
314 + _MODEL_REVISION,
315 + "--max-model-len",
316 + "1024",
317 + ]
318 + args = parse_snapshot(*argv)
319 + monkeypatch.setattr(sys, "argv", ["vllm", "snapshot", *argv])
320 + child = snapshot_server.parse_vllm_args(list(_child_engine_argv(args, None)))
321 +
322 + assert (
323 + (args.model_tag or args.model)
324 + == child.model
325 + == ("Qwen/Qwen3-0.6B" if model_in_cli else "config-model")
326 + )
327 + assert args.max_model_len == child.max_model_len == 1024
328 + assert args.dtype == child.dtype == "float16"
329 +
330 +
331 +@pytest.mark.parametrize(
332 + ("probe", "status", "output"),
333 + [
334 + (None, 0, ""),
335 + ("nvidia-smi", 1, ""),
336 + ("find", 1, ""),
337 + ("ss", 1, ""),
338 + ("nvidia-smi", 0, "12345"),
339 + ("find", 0, "link_remap.fixture"),
340 + ("ss", 0, "LISTEN"),
341 + ],
342 +)
343 +def test_snapshot_e2e_cleanup_requires_successful_empty_probes(probe, status, output):
344 + script = (
345 + Path(__file__).resolve().parents[3]
346 + / ".buildkite/scripts/initialized-snapshot-e2e.sh"
347 + )
348 + if not script.exists():
349 + pytest.skip("snapshot E2E script is not packaged in this test image")
350 + source = script.read_text()
351 + # Exercise the real probes without entering Docker/CRIU setup.
352 + functions = source[source.index("link_remaps() {") : source.index("wait_clean() {")]
353 + shell = (
354 + """set -Eeuo pipefail
355 +GPU_UUID=fixture PORT_ONE=18001 PORT_TWO=18002 LINK_REMAP_BASELINE=''
356 +SNAPSHOT_PIDS=(999999999)
357 +nvidia-smi() { :; }
358 +find() { :; }
359 +ss() { :; }
360 +"""
361 + + functions
362 + )
363 + if probe:
364 + shell += f'\n{probe}() {{ printf %s "{output}"; return {status}; }}\n'
365 + shell += "\nif state_clean; then exit 0; else exit 1; fi\n"
366 +
367 + result = subprocess.run(["bash", "-c", shell], capture_output=True, timeout=10)
368 +
369 + assert result.returncode == (0 if probe is None else 1), result.stderr
370 +
371 +
372 +def test_snapshot_create_cli_accepts_only_pinned_compact_mode(
373 + monkeypatch: pytest.MonkeyPatch,
374 +):
375 + monkeypatch.setattr(snapshot_cli.platform, "system", lambda: "Linux")
376 + monkeypatch.setattr(snapshot_cli.platform, "machine", lambda: "x86_64")
377 + base = ("create", "Qwen/Qwen3-0.6B", "--snapshot-dir=/tmp/snapshot")
378 + compact = parse_snapshot(*base, "--revision", _MODEL_REVISION)
379 +
380 + snapshot_cli.validate_create_args(compact)
381 + assert compact.model_tag == "Qwen/Qwen3-0.6B"
382 +
383 + dispatched: list[argparse.Namespace] = []
384 + monkeypatch.setattr(snapshot_cli, "run_create", dispatched.append)
385 + _run_cli(["snapshot", *base, "--revision", _MODEL_REVISION])
386 + assert dispatched[0].revision == _MODEL_REVISION
387 +
388 +
389 +@pytest.mark.parametrize(
390 + ("attribute", "value", "message"),
391 + [
392 + ("prefill_context_parallel_size", 2, "prefill context parallel size 1"),
393 + ("model_tag", None, "immutable remote model"),
394 + ("logprobs_mode", "raw_logits", "log-probability mode"),
395 + ("speculative_config", {}, "does not support speculative decoding"),
396 + ("spec_method", "ngram", "does not support speculative decoding"),
397 + ("spec_model", "draft-model", "does not support speculative decoding"),
398 + ("spec_tokens", 3, "does not support speculative decoding"),
399 + ],
400 +)
401 +def test_snapshot_create_rejects_invalid_config(
402 + attribute: str,
403 + value: object,
404 + message: str,
405 + tmp_path: Path,
406 +):
407 + args = parse_snapshot(
408 + "create",
409 + "Qwen/Qwen3-0.6B",
410 + "--snapshot-dir=/tmp/snapshot",
411 + "--revision",
412 + _MODEL_REVISION,
413 + )
414 + if attribute == "model_tag":
415 + local_model = tmp_path / "model"
416 + local_model.mkdir()
417 + value = str(local_model)
418 + setattr(args, attribute, value)
419 + with pytest.raises(ValueError, match=message):
420 + snapshot_cli.validate_create_args(args)
421 +
422 +
423 +def _oracle(
424 + token_ids: tuple[int, ...] = (12095,),
425 + text: str = " Paris",
426 + logprob: float = -0.125,
427 +) -> Oracle:
428 + return Oracle(token_ids, text, logprob)
429 +
430 +
431 +def _manifest(**changes: object) -> SnapshotManifest:
432 + manifest = SnapshotManifest(
433 + schema_version=1,
434 + boundary="post-engine-init-reloadable-state-released",
435 + created_at="2026-08-06T00:00:00Z",
436 + artifact_bytes=8,
437 + vllm_version="0.1.dev1+snapshot",
438 + python_version="3.12.3",
439 + torch_version="2.9.0",
440 + cuda_runtime="12.9",
441 + driver_version="575.57.08",
442 + criu_version="4.1",
443 + cuda_checkpoint_sha256="a" * 64,
444 + kernel_release="6.8.0",
445 + host_id="host-a",
446 + gpu_name="NVIDIA A10",
447 + gpu_uuid="GPU-abc",
448 + model="Qwen/Qwen3-0.6B",
449 + served_model_name="Qwen/Qwen3-0.6B",
450 + model_revision="model-sha",
451 + tokenizer_revision="tokenizer-sha",
452 + engine_argv=("Qwen/Qwen3-0.6B",),
453 + environment=(("VLLM_USE_V1", "1"),),
454 + process_tree=(100, 101),
455 + cuda_holders=(101,),
456 + oracle_token_ids=(12095,),
457 + oracle_text=" Paris",
458 + oracle_sampled_token_logprob=-0.125,
459 + )
460 + return SnapshotManifest.model_validate(
461 + {**manifest.model_dump(mode="python"), **changes}, strict=True
462 + )
463 +
464 +
465 +def _runtime_identity(**changes: object) -> Any:
466 + fields = SnapshotRuntimeIdentity.model_fields.keys()
467 + payload = _manifest().model_dump(mode="python", include=fields)
468 + return SnapshotRuntimeIdentity.model_validate({**payload, **changes}, strict=True)
469 +
470 +
471 +def _fake_snapshot_tools() -> Any:
472 + tools = create_autospec(LocalSnapshotTools, instance=True)
473 + tools.launch_child.return_value = 100
474 + tools.wait_ready.return_value = _oracle()
475 + tools.inventory.return_value = ProcessInventory(100, (100, 101), (101,), "GPU-abc")
476 +
477 + tools.current_identity.side_effect = lambda _gpu_uuid: _runtime_identity()
478 + tools.restore.return_value = 100
479 + tools.request_oracle.return_value = _oracle()
480 + return tools
481 +
482 +
483 +@pytest.mark.parametrize("value", ["0", "-1", "nan", "inf"])
484 +def test_snapshot_timeout_must_be_positive_and_finite(
485 + value: str, monkeypatch: pytest.MonkeyPatch
486 +):
487 + monkeypatch.setenv("VLLM_SNAPSHOT_TIMEOUT_S", value)
488 + with pytest.raises(ValueError, match="positive finite number"):
489 + LocalSnapshotTools()
490 +
491 +
492 +def test_snapshot_create_abort_reaps_after_process_group_cleanup(
493 + tmp_path: Path, monkeypatch: pytest.MonkeyPatch
494 +):
495 + tools = LocalSnapshotTools()
496 + process = subprocess.Popen(
497 + [sys.executable, "-c", "raise SystemExit(1)"],
498 + stdin=subprocess.DEVNULL,
499 + stdout=subprocess.DEVNULL,
500 + stderr=subprocess.DEVNULL,
501 + start_new_session=True,
502 + )
503 + tools._children[process.pid] = process
504 + (tmp_path / "child.log").write_text("root failed")
505 + killed: list[tuple[int, signal.Signals]] = []
506 + monkeypatch.setattr(
507 + os,
508 + "killpg",
509 + lambda pid, sig: killed.append((pid, sig)),
510 + )
511 + monkeypatch.setattr(os, "waitid", lambda *_args: object(), raising=False)
512 +
513 + try:
514 + with pytest.raises(SnapshotCreateError, match="root failed"):
515 + tools.wait_ready(tmp_path, process.pid)
516 + assert process.returncode is None
517 + tools.abort_create(process.pid)
518 + finally:
519 + if process.pid in tools._children:
520 + tools.abort_create(process.pid)
521 +
522 + assert killed == [(process.pid, signal.SIGKILL)]
523 + assert process.pid not in tools._children
524 +
525 +
526 +def _supply_waitid(monkeypatch: pytest.MonkeyPatch) -> None:
527 + """`wait_ready` binds Linux-only `os.waitid` before it reads the marker."""
528 + monkeypatch.setattr(os, "waitid", lambda *_args: None, raising=False)
529 +
530 +
531 +def test_snapshot_wait_ready_reads_back_the_oracle_the_child_wrote(
532 + tmp_path: Path, monkeypatch: pytest.MonkeyPatch
533 +):
534 + _supply_waitid(monkeypatch)
535 + snapshot_server.write_ready_atomic(tmp_path / "ready.json", _oracle())
536 +
537 + assert LocalSnapshotTools().wait_ready(tmp_path, os.getpid()) == _oracle()
538 +
539 +
540 +@pytest.mark.parametrize(
541 + "payload",
542 + [
543 + {"token_ids": [12095], "sampled_token_logprob": -0.125},
544 + {"token_ids": ["12095"], "text": " Paris", "sampled_token_logprob": -0.125},
545 + ],
546 +)
547 +def test_snapshot_wait_ready_rejects_a_malformed_marker(
548 + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, payload: dict[str, object]
549 +):
550 + _supply_waitid(monkeypatch)
551 + (tmp_path / "ready.json").write_text(json.dumps(payload))
552 +
553 + with pytest.raises(SnapshotCreateError, match="ready marker invalid"):
554 + LocalSnapshotTools().wait_ready(tmp_path, os.getpid())
555 +
556 +
557 +def test_snapshot_create_abort_rejects_surviving_process_group(
558 + monkeypatch: pytest.MonkeyPatch,
559 +):
560 + tools = LocalSnapshotTools()
561 + process = create_autospec(subprocess.Popen, instance=True)
562 + process.wait.side_effect = subprocess.TimeoutExpired("snapshot child", 10)
563 + tools._children[100] = process
564 + monkeypatch.setattr(os, "killpg", lambda _pid, _sig: None)
565 +
566 + with pytest.raises(SnapshotCreateError, match="survived SIGKILL"):
567 + tools.abort_create(100)
568 +
569 +
570 +@pytest.mark.asyncio
571 +async def test_snapshot_child_rejects_mismatch_before_publication(
572 + tmp_path: Path, monkeypatch: pytest.MonkeyPatch
573 +):
574 + from vllm.entrypoints.launchers.api_server import entry as api_server
575 +
576 + canaries = iter((_oracle(), _oracle((42,), " Lyon")))
577 +
578 + @asynccontextmanager
579 + async def engine_context():
580 + yield object()
581 +
582 + async def next_canary(_engine: object) -> Oracle:
583 + return next(canaries)
584 +
585 + async def return_none(_engine: object) -> None:
586 + return None
587 +
588 + monkeypatch.setattr(
589 + api_server, "build_async_engine_client", lambda _args: engine_context()
590 + )
591 + monkeypatch.setattr(snapshot_server, "run_engine_canary", next_canary)
592 + monkeypatch.setattr(snapshot_server, "_release_reloadable_state", return_none)
593 + monkeypatch.setattr(snapshot_server, "_restore_reloadable_state", return_none)
594 + monkeypatch.setattr(snapshot_server, "detach_snapshot_streams", lambda: None)
595 + monkeypatch.setattr(
596 + snapshot_server,
597 + "write_ready_atomic",
598 + lambda *_args: pytest.fail("mismatched snapshot was published"),
599 + )
600 + control = snapshot_server.ControlArgs(
601 + ready_file=tmp_path / "ready.json",
602 + release_file=tmp_path / "release.json",
603 + release_timeout_s=1,
604 + )
605 + with pytest.raises(SnapshotCanaryError, match="rehearsal changed"):
606 + await snapshot_server.run_vllm_snapshot_child(control, argparse.Namespace())
607 +
608 +
609 +def test_snapshot_release_marker_ignores_unknown_keys(tmp_path: Path):
610 + marker = tmp_path / "release.json"
611 + marker.write_text(json.dumps({"release": True, "port": 8000, "unknown": 1}))
612 +
613 + listener = snapshot_server.read_release_marker(marker)
614 +
615 + assert listener == snapshot_server.ListenerConfig(host=None, port=8000)
616 +
617 +
618 +@pytest.mark.parametrize(
619 + "payload",
620 + [
621 + ["release"],
622 + {"port": 8000},
623 + {"release": "true", "port": 8000},
624 + {"release": 1, "port": 8000},
625 + {"release": True, "port": 0},
626 + {"release": True, "port": 70000},
627 + {"release": True, "port": True},
628 + {"release": True, "port": "8000"},
629 + {"release": True, "port": 8000, "host": 5},
630 + ],
631 +)
632 +def test_snapshot_release_marker_rejects_an_invalid_payload(
633 + tmp_path: Path, payload: object
634 +):
635 + marker = tmp_path / "release.json"
636 + marker.write_text(json.dumps(payload))
637 +
638 + with pytest.raises(snapshot_server.SnapshotBarrierError, match="release marker"):
639 + snapshot_server.read_release_marker(marker)
640 +
641 +
642 +def test_snapshot_create_rolls_back_failed_dump(tmp_path: Path):
643 + target = tmp_path / "snapshot"
644 + tools = _fake_snapshot_tools()
645 + tools.dump.side_effect = RuntimeError("dump failed")
646 + with pytest.raises(RuntimeError, match="dump failed"):
647 + create_snapshot(
648 + argparse.Namespace(snapshot_dir=str(target), model_tag="Qwen/Qwen3-0.6B"),
649 + tools=tools,
650 + )
651 + tools.abort_create.assert_called_once_with(100)
652 + assert not target.exists()
653 +
654 +
655 +def _restore_fixture(tmp_path: Path) -> argparse.Namespace:
656 + artifact = tmp_path / "snapshot"
657 + artifact.mkdir(mode=0o700)
658 + write_manifest_atomic(artifact, _manifest())
659 + return argparse.Namespace(snapshot_dir=str(artifact), host="127.0.0.1", port=9000)
660 +
661 +
662 +def test_snapshot_restore_rejects_identity_mismatch(tmp_path: Path):
663 + args = _restore_fixture(tmp_path)
664 +
665 + identity_mismatch = _fake_snapshot_tools()
666 + identity_mismatch.current_identity.side_effect = lambda _gpu_uuid: (
667 + _runtime_identity(vllm_version="different")
668 + )
669 + with pytest.raises(SnapshotCompatibilityError, match="vllm_version"):
670 + restore_snapshot(args, tools=identity_mismatch)
671 + identity_mismatch.restore.assert_not_called()
672 +
673 +
674 +def test_snapshot_identity_mismatch_names_every_differing_field():
675 + with pytest.raises(SnapshotCompatibilityError) as mismatch:
676 + validate_identity(
677 + _manifest(),
678 + _runtime_identity(vllm_version="different", gpu_uuid="GPU-other"),
679 + )
680 +
681 + assert str(mismatch.value) == "snapshot mismatch: vllm_version, gpu_uuid"
682 +
683 +
684 +def test_snapshot_restore_cleans_oracle_mismatch(tmp_path: Path):
685 + args = _restore_fixture(tmp_path)
686 + oracle_mismatch = _fake_snapshot_tools()
687 + oracle_mismatch.request_oracle.return_value = _oracle(logprob=-0.5)
688 + with pytest.raises(SnapshotRestoreError, match="oracle mismatch"):
689 + restore_snapshot(args, tools=oracle_mismatch)
690 + oracle_mismatch.cleanup.assert_called_once()
691 + oracle_mismatch.complete_restore.assert_not_called()
692 + assert oracles_match(_oracle(logprob=-100.0), _oracle(logprob=-100.0005))
693 + assert not oracles_match(_oracle(logprob=-100.0), _oracle(logprob=-100.002))
694 +
695 +
696 +@pytest.mark.parametrize(
697 + ("invalid_update", "diagnostic"),
698 + [
699 + ({"unexpected": True}, "unexpected"),
700 + ({"schema_version": True}, "schema_version"),
701 + ({"boundary": "unknown"}, "boundary"),
702 + ({"process_tree": [100, 100]}, "root"),
703 + ({"cuda_holders": [999]}, "root"),
704 + ({"oracle_token_ids": [1, 2]}, "oracle_token_ids"),
705 + ({"oracle_sampled_token_logprob": 0}, "oracle_sampled_token_logprob"),
706 + (
707 + {"oracle_sampled_token_logprob": float("nan")},
708 + "oracle_sampled_token_logprob",
709 + ),
710 + ],
711 +)
712 +def test_snapshot_manifest_validation(
713 + tmp_path: Path,
714 + invalid_update: dict[str, object],
715 + diagnostic: str,
716 +):
717 + artifact = tmp_path / "snapshot"
718 + artifact.mkdir(mode=0o700)
719 + write_manifest_atomic(artifact, _manifest())
720 + invalid = _manifest().model_dump(mode="json")
721 + invalid.update(invalid_update)
722 + (artifact / "manifest.json").write_text(json.dumps(invalid))
723 + with pytest.raises(
724 + SnapshotCompatibilityError,
725 + match=f"invalid snapshot manifest: {diagnostic}",
726 + ):
727 + read_manifest(artifact)
728 +
729 +
730 +def _local_restore_artifact(tmp_path: Path, name: str = "snapshot") -> Path:
731 + artifact = tmp_path / name
732 + artifact.mkdir(mode=0o700)
733 + (artifact / "child.log").write_bytes(b"startup log")
734 + (artifact / "child.log.snapshot-size").write_text("11\n")
735 + return artifact
736 +
737 +
738 +def _mark_snapshot_pids_free(monkeypatch: pytest.MonkeyPatch) -> list[int]:
739 + probed_pids: list[int] = []
740 +
741 + def missing(pid: int, _signal: int) -> None:
742 + probed_pids.append(pid)
743 + raise ProcessLookupError
744 +
745 + monkeypatch.setattr(snapshot_runtime.os, "kill", missing)
746 + return probed_pids
747 +
748 +
749 +@pytest.mark.parametrize(
750 + ("probe_error", "diagnostic"),
751 + [
752 + (None, "already occupied: 100"),
753 + (PermissionError(), "already occupied: 100"),
754 + (OSError("probe failed"), "availability probe failed: 100"),
755 + ],
756 +)
757 +def test_snapshot_restore_rejects_unavailable_pid_before_mutation(
758 + tmp_path: Path,
759 + monkeypatch: pytest.MonkeyPatch,
760 + probe_error: OSError | None,
761 + diagnostic: str,
762 +):
763 + args = _restore_fixture(tmp_path)
764 + artifact = Path(args.snapshot_dir)
765 + (artifact / "child.log").write_bytes(b"startup log")
766 + (artifact / "release.json").write_bytes(b"release sentinel")
767 + before = {path.name: path.read_bytes() for path in artifact.iterdir()}
768 + tools = LocalSnapshotTools()
769 +
770 + def fail_probe(_pid: int, _signal: int) -> None:
771 + if probe_error is not None:
772 + raise probe_error
773 +
774 + monkeypatch.setattr(snapshot_runtime.os, "kill", fail_probe)
775 + monkeypatch.setattr(tools, "preflight", lambda *_args: None)
776 + monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity())
777 + monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
778 +
779 + with pytest.raises(SnapshotRestoreError, match=diagnostic):
780 + restore_snapshot(args, tools=tools)
781 + assert {path.name: path.read_bytes() for path in artifact.iterdir()} == before
782 +
783 +
784 +@pytest.mark.parametrize("mode", ["pinned", "unpinned", "cleanup-fails"])
785 +def test_snapshot_restore_terminates_only_a_pinned_tree(
786 + tmp_path: Path,
787 + monkeypatch: pytest.MonkeyPatch,
788 + mode: str,
789 +):
790 + _mark_snapshot_pids_free(monkeypatch)
791 + artifact = _local_restore_artifact(tmp_path)
792 + manifest = _manifest()
793 + root_pid = manifest.process_tree[0]
794 + tools = LocalSnapshotTools()
795 + terminated: list[int] = []
796 +
797 + def terminate_and_wait(pid: int) -> None:
798 + terminated.append(pid)
799 + if mode == "cleanup-fails":
800 + raise SnapshotRestoreError("restored process cleanup is incomplete")
801 + tools._restored_processes.pop(pid)
802 +
803 + def pin(
804 + _artifact: Path, process_tree: tuple[int, ...], _holders: tuple[int, ...]
805 + ) -> None:
806 + if mode == "unpinned":
807 + raise SnapshotRestoreError(
808 + "restored session does not match the captured process tree"
809 + )
810 + tools._restored_processes[process_tree[0]] = ()
811 +
812 + monkeypatch.setattr(tools, "cleanup", terminate_and_wait)
813 + monkeypatch.setattr(tools, "_pin_restored_tree", pin)
814 + monkeypatch.setattr(tools, "_criu", lambda *_args: None)
815 + monkeypatch.setattr(
816 + tools,
817 + "_run",
818 + lambda command, **_kwargs: subprocess.CompletedProcess(command, 0, "333\n", ""),
819 + )
820 +
821 + with pytest.raises(SnapshotRestoreError, match="does not match") as failure:
822 + tools.restore(artifact, manifest)
823 +
824 + # A pinned tree is ours: it must be terminated and waited for, not left to
825 + # honour the abort marker. An unpinned PID is never signalled.
826 + assert terminated == ([] if mode == "unpinned" else [root_pid])
827 + if mode == "cleanup-fails":
828 + # The failure to terminate is reported beside the primary error, and a
829 + # tree that may still be running stays pinned.
830 + assert "restore cleanup failed" in str(failure.value)
831 + assert tools._restored_processes == {root_pid: ()}
832 + else:
833 + assert tools._restored_processes == {}
834 + assert json.loads((artifact / "release.json").read_text())["release"] is False
835 +
836 +
837 +def test_snapshot_private_path_and_link_remap_security(
838 + tmp_path: Path, monkeypatch: pytest.MonkeyPatch
839 +):
840 + _mark_snapshot_pids_free(monkeypatch)
841 + real = tmp_path / "real"
842 + real.mkdir(mode=0o700)
843 + linked = tmp_path / "linked"
844 + linked.symlink_to(real, target_is_directory=True)
845 + with pytest.raises(SnapshotSecurityError, match="symlink"):
846 + validate_artifact_root(linked, creating=False)
847 +
848 + unsafe = tmp_path / "unsafe"
849 + unsafe.mkdir()
850 + unsafe.chmod(0o777)
851 + with pytest.raises(SnapshotSecurityError, match="world-writable"):
852 + validate_artifact_root(unsafe / "snapshot", creating=True)
853 +
854 + artifact = _local_restore_artifact(tmp_path)
855 + saved_remaps = artifact / "link-remaps"
856 + saved_remaps.mkdir()
857 + (saved_remaps / "link_remap.270").write_bytes(b"semaphore state")
858 + shm_dir = tmp_path / "shm"
859 + shm_dir.mkdir()
860 + target = shm_dir / "link_remap.270"
861 + replacement = tmp_path / "replacement"
862 + replacement.write_bytes(b"semaphore state")
863 + target.symlink_to(replacement)
864 + tools = LocalSnapshotTools()
865 + tools.shm_dir = shm_dir
866 + monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
867 + with pytest.raises(SnapshotRestoreError, match="conflicting CRIU link remap"):
868 + tools.restore(artifact, _manifest())
869 +
870 +
871 +@pytest.mark.parametrize(
872 + ("blocked_state", "message"),
873 + [("tcp", "external established TCP"), ("io_uring", "kernel.io_uring_disabled=1")],
874 +)
875 +def test_snapshot_rejects_unsafe_process_state_before_criu(
876 + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, blocked_state: str, message: str
877 +):
878 + tools = LocalSnapshotTools()
879 + monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101))
880 + monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",))
881 + monkeypatch.setattr(
882 + tools,
883 + "_descriptor_inventory",
884 + lambda _tree: (
885 + (101,) if blocked_state == "io_uring" else (),
886 + {41: 101},
887 + ),
888 + )
889 + tcp_record = _TcpSocketRecord(
890 + family="AF_INET",
891 + local_raw="3C00000A:D6F6",
892 + remote_raw="01010101:01BB" if blocked_state == "tcp" else "00000000:0000",
893 + inode=41,
894 + )
895 + monkeypatch.setattr(tools, "_tcp_records", lambda: (tcp_record,))
896 + with pytest.raises(SnapshotCreateError, match=message) as excinfo:
897 + tools.inventory(100, tmp_path)
898 + if blocked_state == "tcp":
899 + assert "pid 101" in str(excinfo.value)
900 + assert "1.1.1.1:443" in str(excinfo.value)
901 +
902 +
903 +def test_decode_endpoint_families():
904 + assert _decode_endpoint("AF_INET", "3C00000A:D6F6") == "10.0.0.60:55030"
905 + assert (
906 + _decode_endpoint("AF_INET6", "00000000000000000000000001000000:01BB")
907 + == "[::1]:443"
908 + )
909 +
910 +
911 +def test_tcp_records_skips_missing_tables(tmp_path, monkeypatch):
912 + tcp, tcp6 = tmp_path / "tcp", tmp_path / "tcp6"
913 + contents = "header\n0: 0100007F:1F90 00000000:0000 01 0 0 0 0 0 41\n"
914 + tcp.write_text(contents)
915 + tables = (("AF_INET", tcp), ("AF_INET6", tcp6))
916 + monkeypatch.setattr(snapshot_runtime, "_TCP_TABLES", tables)
917 +
918 + tools = LocalSnapshotTools()
919 + assert tuple(record.family for record in tools._tcp_records()) == ("AF_INET",)
920 +
921 + tcp6.write_text(contents)
922 + families = tuple(record.family for record in tools._tcp_records())
923 + assert families == ("AF_INET", "AF_INET6")
924 +
925 +
926 +def test_snapshot_manifest_records_external_cache_files(tmp_path: Path):
927 + artifact = tmp_path / "snapshot"
928 + artifact.mkdir(mode=0o700)
929 + recorded = (
930 + (str(tmp_path / "kernel.so"), "b" * 64),
931 + (str(tmp_path / "launcher.lock"), "locked"),
932 + )
933 + write_manifest_atomic(artifact, _manifest(external_cache_files=recorded))
934 +
935 + assert read_manifest(artifact).external_cache_files == recorded
936 +
937 + # A manifest written before this field existed still loads.
938 + legacy = _manifest().model_dump(mode="json")
939 + del legacy["external_cache_files"]
940 + (artifact / "manifest.json").write_text(json.dumps(legacy))
941 +
942 + assert read_manifest(artifact).external_cache_files == ()
943 +
944 +
945 +def _pin_generated_cache_roots(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
946 + """Keep the resolver off whatever cache roots the host running this has."""
947 + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path / "vllm_cache"))
948 + monkeypatch.setenv("TRITON_CACHE_DIR", str(tmp_path / "triton_cache_default"))
949 + monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", str(tmp_path / "inductor_cache"))
950 +
951 +
952 +def _unreadable(_path: Path) -> str:
953 + raise OSError("cache file is unreadable")
954 +
955 +
956 +def test_snapshot_inventory_records_open_generated_cache_files(
957 + tmp_path: Path, monkeypatch: pytest.MonkeyPatch
958 +):
959 + cache_root = tmp_path / "triton_cache"
960 + cache_root.mkdir()
961 + kernel = cache_root / "kernel.so"
962 + kernel.write_bytes(b"compiled kernel")
963 + lock = cache_root / "launcher.lock"
964 + lock.write_bytes(b"held by pid 100")
965 + weights = tmp_path / "model.safetensors"
966 + weights.write_bytes(b"weights")
967 + artifact = tmp_path / "snapshot"
968 + artifact.mkdir(mode=0o700)
969 + (artifact / "child.log.lock").write_bytes(b"")
970 + _pin_generated_cache_roots(monkeypatch, tmp_path)
971 + monkeypatch.setenv("TRITON_CACHE_DIR", str(cache_root))
972 +
973 + targets = {
974 + 100: (str(kernel), str(weights), "socket:[41]"),
975 + 101: (str(lock), str(cache_root), str(artifact / "child.log.lock")),
976 + 102: (f"{cache_root / 'evicted.so'} (deleted)",),
977 + }
978 + tools = LocalSnapshotTools()
979 + monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101, 102))
980 + monkeypatch.setattr(tools, "_descriptor_targets", lambda pid: targets[pid])
981 + monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",))
982 + monkeypatch.setattr(tools, "_tcp_records", lambda: ())
983 + monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity())
984 +
985 + manifest = tools.make_manifest(
986 + argparse.Namespace(
987 + model_tag="Qwen/Qwen3-0.6B",
988 + revision=_MODEL_REVISION,
989 + served_model_name=None,
990 + tokenizer_revision=None,
991 + ),
992 + ("Qwen/Qwen3-0.6B",),
993 + tools.inventory(100, artifact),
994 + _oracle(),
995 + artifact,
996 + )
997 +
998 + # Weights, sockets, the cache directory itself, files inside the artifact,
999 + # and deleted targets that CRIU remaps are all left out.
1000 + assert manifest.external_cache_files == (
1001 + (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()),
1002 + (str(lock), "locked"),
1003 + )
1004 +
1005 +
1006 +def test_snapshot_inventory_marks_a_cache_file_it_cannot_read(
1007 + tmp_path: Path, monkeypatch: pytest.MonkeyPatch
1008 +):
1009 + cache_root = tmp_path / "triton_cache"
1010 + cache_root.mkdir()
1011 + kernel = cache_root / "kernel.so"
1012 + kernel.write_bytes(b"compiled kernel")
1013 + lock = cache_root / "launcher.lock"
1014 + lock.write_bytes(b"held by pid 100")
1015 + artifact = tmp_path / "snapshot"
1016 + artifact.mkdir(mode=0o700)
1017 + _pin_generated_cache_roots(monkeypatch, tmp_path)
1018 + monkeypatch.setenv("TRITON_CACHE_DIR", str(cache_root))
1019 +
1020 + tools = LocalSnapshotTools()
1021 + monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100,))
1022 + monkeypatch.setattr(
1023 + tools, "_descriptor_targets", lambda _pid: (str(kernel), str(lock))
1024 + )
1025 + monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("100, GPU-abc",))
1026 + monkeypatch.setattr(tools, "_tcp_records", lambda: ())
1027 + monkeypatch.setattr(tools, "_sha256", _unreadable)
1028 +
1029 + # A file that cannot be digested keeps its own marker, so the manifest says
1030 + # which entries restore can check for existence only.
1031 + assert tools.inventory(100, artifact).external_cache_files == (
1032 + (str(kernel), "unreadable"),
1033 + (str(lock), "locked"),
1034 + )
1035 +
1036 +
1037 +def test_snapshot_inventory_matches_resolved_cache_roots_not_siblings(
1038 + tmp_path: Path, monkeypatch: pytest.MonkeyPatch
1039 +):
1040 + # Descriptor targets come back kernel-canonical, so the fixture must be too.
1041 + base = Path(os.path.realpath(tmp_path))
1042 + real_cache = base / "mnt_cache"
1043 + real_cache.mkdir()
1044 + kernel = real_cache / "kernel.so"
1045 + kernel.write_bytes(b"compiled kernel")
1046 + linked_cache = base / "linked_cache"
1047 + linked_cache.symlink_to(real_cache, target_is_directory=True)
1048 + sibling = base / "vllm_cache-backup"
1049 + sibling.mkdir()
1050 + stale = sibling / "stale.so"
1051 + stale.write_bytes(b"stale kernel")
1052 + artifact = base / "snapshot"
1053 + artifact.mkdir(mode=0o700)
1054 +
1055 + _pin_generated_cache_roots(monkeypatch, base)
1056 + monkeypatch.setenv("TRITON_CACHE_DIR", str(linked_cache))
1057 +
1058 + tools = LocalSnapshotTools()
1059 + monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100,))
1060 + monkeypatch.setattr(
1061 + tools, "_descriptor_targets", lambda _pid: (str(kernel), str(stale))
1062 + )
1063 + monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("100, GPU-abc",))
1064 + monkeypatch.setattr(tools, "_tcp_records", lambda: ())
1065 +
1066 + # A cache root reached through a symlink still covers the file the kernel
1067 + # reports by its real path, and a directory that merely shares the root's
1068 + # name prefix is not part of the cache at all.
1069 + assert tools.inventory(100, artifact).external_cache_files == (
1070 + (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()),
1071 + )
1072 +
1073 +
1074 +def test_snapshot_restore_rejects_rotated_cache_file_before_mutation(
1075 + tmp_path: Path, monkeypatch: pytest.MonkeyPatch
1076 +):
1077 + _mark_snapshot_pids_free(monkeypatch)
1078 + kernel = tmp_path / "kernel.so"
1079 + kernel.write_bytes(b"compiled kernel")
1080 + lock = tmp_path / "launcher.lock"
1081 + lock.write_bytes(b"held by pid 100")
1082 + manifest = _manifest(
1083 + external_cache_files=(
1084 + (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()),
1085 + (str(lock), "locked"),
1086 + )
1087 + )
1088 + artifact = _local_restore_artifact(tmp_path)
1089 + before = {path.name: path.read_bytes() for path in artifact.iterdir()}
1090 + tools = LocalSnapshotTools()
1091 + monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
1092 +
1093 + kernel.write_bytes(b"recompiled kernel")
1094 + with pytest.raises(SnapshotRestoreError, match="changed since capture") as changed:
1095 + tools.restore(artifact, manifest)
1096 + assert str(kernel) in str(changed.value)
1097 +
1098 + # Restoring the digest gets past the first entry, so the lock file is only
1099 + # checked for existence: the second failure proves the first one passed.
1100 + kernel.write_bytes(b"compiled kernel")
1101 + lock.unlink()
1102 + with pytest.raises(SnapshotRestoreError, match="vanished since capture") as gone:
1103 + tools.restore(artifact, manifest)
1104 + assert str(lock) in str(gone.value)
1105 +
1106 + assert {path.name: path.read_bytes() for path in artifact.iterdir()} == before
1107 +
1108 +
1109 +def test_snapshot_restore_rejects_a_cache_file_it_cannot_read(
1110 + tmp_path: Path, monkeypatch: pytest.MonkeyPatch
1111 +):
1112 + _mark_snapshot_pids_free(monkeypatch)
1113 + kernel = tmp_path / "kernel.so"
1114 + kernel.write_bytes(b"compiled kernel")
1115 + manifest = _manifest(
1116 + external_cache_files=(
1117 + (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()),
1118 + )
1119 + )
1120 + artifact = _local_restore_artifact(tmp_path)
1121 + tools = LocalSnapshotTools()
1122 + monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
1123 + monkeypatch.setattr(tools, "_sha256", _unreadable)
1124 +
1125 + # A recorded digest that cannot be read again is a failed check, not a file
1126 + # to wave through on existence.
1127 + with pytest.raises(
1128 + SnapshotRestoreError, match="unreadable since capture"
1129 + ) as blocked:
1130 + tools.restore(artifact, manifest)
1131 + assert str(kernel) in str(blocked.value)
1132 +
1133 +
1134 +def test_snapshot_runtime_installer_logs_the_failing_exit_code():
1135 + installer = Path(__file__).parents[3] / "tools" / "install_snapshot_runtime.sh"
1136 + if not installer.is_file():
1137 + pytest.skip("installer script is not shipped in this test image")
1138 + timing = [
1139 + line
1140 + for line in installer.read_text().splitlines()
1141 + if line.startswith(("_snapshot_install_started=", "trap "))
1142 + ]
1143 + assert len(timing) == 2
1144 +
1145 + result = subprocess.run(
1146 + ["bash", "-c", "\n".join([*timing, "exit 7"])],
1147 + capture_output=True,
1148 + text=True,
1149 + )
1150 +
1151 + assert result.returncode == 7
1152 + assert "(exit 7)" in result.stdout
tests/models/quantization/test_mimo_v2_qkv_shard.py
+8
@@ -19,6 +19,7 @@ import torch
19
20 from vllm.model_executor.models.mimo_v2 import _shard_fp8_qkv_proj
21 from vllm.utils.math_utils import cdiv
22 +from vllm.utils.torch_utils import set_default_torch_num_threads
23
24 pytestmark = pytest.mark.cpu_test
25
@@ -39,6 +40,13 @@ GEOMETRIES = {
40 }
41
42
43 +@pytest.fixture(scope="module", autouse=True)
44 +def single_threaded_cpu_ops():
45 + # Do not prime the parent's thread pool before later engine forks.
46 + with set_default_torch_num_threads(1):
47 + yield
48 +
49 +
50 def _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp):
51 return (
52 num_heads // ckpt_tp * head_dim
tools/install_snapshot_runtime.sh new
+109
@@ -0,0 +1,109 @@
1 +#!/usr/bin/env bash
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +
5 +set -euxo pipefail
6 +
7 +# Surface the isolated snapshot-runtime install cost in the CI build log. The
8 +# trap captures the status first: the date substitution below would reset it.
9 +_snapshot_install_started=$(date +%s)
10 +code=0
11 +trap 'code=$?; echo "[TIMING] snapshot runtime install: $(( $(date +%s) - _snapshot_install_started )) s (exit $code)"' EXIT
12 +
13 +readonly CRIU_VERSION="4.2.1"
14 +readonly CRIU_SOURCE_SHA256="feffdf4638125ebb12d2434754f80a1d7bbba85a3e6bee98c216f88fb99a5d96"
15 +readonly DYNAMO_CHECKPOINT_HELPER_COMMIT="fdf25efad60f696a73393caf341c527a93d84190"
16 +readonly DYNAMO_CHECKPOINT_HELPER_SHA256="7b8ea21baecb1da729202c32e0fecbbbafb789d65e5c1c0649ea253ad09b8f82"
17 +readonly DYNAMO_LICENSE_SHA256="aeeca3d74a13b91d17ab9211c407b8d94a77bf8d9e5750c467818d3066b8adcf"
18 +
19 +# Keep the compiler and NL headers required by runtime JIT and RDMA.
20 +snapshot_shared_build_deps=(
21 + build-essential
22 + libnl-3-dev
23 + libnl-route-3-dev
24 +)
25 +snapshot_build_deps=(
26 + pkg-config
27 + libbsd-dev
28 + libcap-dev
29 + libnet1-dev
30 + libprotobuf-dev
31 + libprotobuf-c-dev
32 + protobuf-c-compiler
33 + protobuf-compiler
34 + python3-protobuf
35 + libgnutls28-dev
36 + libnftables-dev
37 + uuid-dev
38 +)
39 +snapshot_runtime_deps=(
40 + libbsd0
41 + libcap2
42 + libnet1
43 + libnl-3-200
44 + libnl-route-3-200
45 + libprotobuf-c1
46 + libgnutls30t64
47 + libnftables1
48 + iproute2
49 + iptables
50 + procps
51 + uuid-runtime
52 + util-linux
53 +)
54 +
55 +apt-get update -y
56 +apt-get install -y --no-install-recommends \
57 + "${snapshot_shared_build_deps[@]}" \
58 + "${snapshot_build_deps[@]}" \
59 + "${snapshot_runtime_deps[@]}"
60 +
61 +snapshot_build_dir="$(mktemp -d)"
62 +curl -fsSL \
63 + "https://github.com/checkpoint-restore/criu/archive/refs/tags/v${CRIU_VERSION}.tar.gz" \
64 + -o "${snapshot_build_dir}/criu.tar.gz"
65 +echo "${CRIU_SOURCE_SHA256} ${snapshot_build_dir}/criu.tar.gz" | sha256sum -c -
66 +mkdir "${snapshot_build_dir}/criu"
67 +tar -xzf "${snapshot_build_dir}/criu.tar.gz" \
68 + --strip-components=1 -C "${snapshot_build_dir}/criu"
69 +make -C "${snapshot_build_dir}/criu" -j"${MAX_JOBS:-$(nproc)}" criu cuda_plugin
70 +install -D -m 0755 "${snapshot_build_dir}/criu/criu/criu" \
71 + /usr/local/sbin/criu
72 +install -D -m 0755 \
73 + "${snapshot_build_dir}/criu/plugins/cuda/cuda_plugin.so" \
74 + /usr/local/lib/criu/cuda_plugin.so
75 +
76 +helper_url="https://raw.githubusercontent.com/ai-dynamo/dynamo/${DYNAMO_CHECKPOINT_HELPER_COMMIT}/deploy/snapshot/cmd/cuda-checkpoint-helper/main.c"
77 +curl -fsSL "${helper_url}" -o "${snapshot_build_dir}/cuda-checkpoint.c"
78 +echo "${DYNAMO_CHECKPOINT_HELPER_SHA256} ${snapshot_build_dir}/cuda-checkpoint.c" \
79 + | sha256sum -c -
80 +gcc -O2 -Wall -Wextra \
81 + -I/usr/local/cuda/include \
82 + -L/usr/local/cuda/lib64/stubs \
83 + "${snapshot_build_dir}/cuda-checkpoint.c" -lcuda \
84 + -o /usr/local/sbin/cuda-checkpoint
85 +
86 +license_url="https://raw.githubusercontent.com/ai-dynamo/dynamo/${DYNAMO_CHECKPOINT_HELPER_COMMIT}/LICENSE"
87 +curl -fsSL "${license_url}" -o "${snapshot_build_dir}/DYNAMO-LICENSE"
88 +echo "${DYNAMO_LICENSE_SHA256} ${snapshot_build_dir}/DYNAMO-LICENSE" \
89 + | sha256sum -c -
90 +
91 +install -D -m 0644 "${snapshot_build_dir}/criu.tar.gz" \
92 + /usr/share/doc/criu/source/criu.tar.gz
93 +install -D -m 0644 "${snapshot_build_dir}/criu/COPYING" \
94 + /usr/share/doc/criu/COPYING
95 +install -D -m 0644 "${snapshot_build_dir}/cuda-checkpoint.c" \
96 + /usr/share/doc/cuda-checkpoint-helper/main.c
97 +install -D -m 0644 "${snapshot_build_dir}/DYNAMO-LICENSE" \
98 + /usr/share/doc/cuda-checkpoint-helper/LICENSE
99 +
100 +criu --version
101 +ln -s /usr/local/cuda/lib64/stubs/libcuda.so \
102 + "${snapshot_build_dir}/libcuda.so.1"
103 +LD_LIBRARY_PATH="${snapshot_build_dir}" \
104 + cuda-checkpoint --help | grep -q -- "--action"
105 +test -x /usr/local/lib/criu/cuda_plugin.so
106 +
107 +apt-get purge -y --auto-remove "${snapshot_build_deps[@]}"
108 +rm -rf "${snapshot_build_dir}"
109 +rm -rf /var/lib/apt/lists/*
vllm/entrypoints/cli/main.py
+10 -2
@@ -46,6 +46,7 @@ def main():
46 import vllm.entrypoints.cli.openai
47 import vllm.entrypoints.cli.run_batch
48 import vllm.entrypoints.cli.serve
49 + import vllm.entrypoints.cli.snapshot
50 from vllm.entrypoints.serve.utils.api_utils import (
51 VLLM_SUBCMD_PARSER_EPILOG,
52 cli_env_setup,
@@ -59,9 +60,11 @@ def main():
60 vllm.entrypoints.cli.benchmark.main,
61 vllm.entrypoints.cli.collect_env,
62 vllm.entrypoints.cli.run_batch,
63 + vllm.entrypoints.cli.snapshot,
64 ]
65
64 - cli_env_setup()
66 + if sys.argv[1:2] != ["snapshot"]:
67 + cli_env_setup()
68
69 vllm.entrypoints.cli.benchmark.main.maybe_exec_rust_bench()
70
@@ -94,7 +97,12 @@ def main():
97 subparsers = parser.add_subparsers(required=False, dest="subparser")
98 cmds = {}
99 for cmd_module in CMD_MODULES:
97 - new_cmds = cmd_module.cmd_init()
100 + if cmd_module is vllm.entrypoints.cli.snapshot:
101 + new_cmds = cmd_module.cmd_init(
102 + create_requested=sys.argv[1:3] == ["snapshot", "create"]
103 + )
104 + else:
105 + new_cmds = cmd_module.cmd_init()
106 for cmd in new_cmds:
107 cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
108 cmds[cmd.name] = cmd
vllm/entrypoints/cli/snapshot.py new
+143
@@ -0,0 +1,143 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import argparse
5 +import json
6 +import os
7 +import platform
8 +import typing
9 +from pathlib import Path
10 +
11 +from vllm.entrypoints.cli.types import CLISubcommand
12 +
13 +if typing.TYPE_CHECKING:
14 + from vllm.utils.argparse_utils import FlexibleArgumentParser
15 +else:
16 + FlexibleArgumentParser = argparse.ArgumentParser
17 +
18 +
19 +def validate_create_args(args: argparse.Namespace) -> None:
20 + model = str(getattr(args, "model_tag", None) or args.model)
21 + revision = str(getattr(args, "revision", None) or "")
22 + if (
23 + Path(model).exists()
24 + or len(revision) != 40
25 + or any(character not in "0123456789abcdefABCDEF" for character in revision)
26 + ):
27 + raise ValueError(
28 + "snapshot create requires an immutable remote model and "
29 + "40-character --revision"
30 + )
31 + parallel_sizes = (
32 + ("tensor parallel", args.tensor_parallel_size),
33 + ("pipeline parallel", args.pipeline_parallel_size),
34 + ("data parallel", args.data_parallel_size),
35 + ("prefill context parallel", args.prefill_context_parallel_size),
36 + )
37 + for name, size in parallel_sizes:
38 + if size != 1:
39 + raise ValueError(f"snapshot create currently requires {name} size 1")
40 + unsupported_frontend = (
41 + args.api_server_count not in (None, 1)
42 + or args.grpc
43 + or args.headless
44 + or args.api_key
45 + or os.environ.get("VLLM_API_KEY")
46 + or args.uds
47 + or args.middleware
48 + )
49 + if unsupported_frontend or any(
50 + (
51 + args.ssl_keyfile,
52 + args.ssl_certfile,
53 + args.ssl_ca_certs,
54 + args.enable_ssl_refresh,
55 + args.ssl_ciphers,
56 + )
57 + ):
58 + raise ValueError(
59 + "snapshot create requires one unauthenticated plaintext HTTP server"
60 + )
61 + if args.logprobs_mode in {"raw_logits", "processed_logits"}:
62 + raise ValueError("snapshot canary requires a log-probability mode")
63 + if any(
64 + getattr(args, name, None) is not None
65 + for name in ("speculative_config", "spec_method", "spec_model", "spec_tokens")
66 + ):
67 + raise ValueError("snapshot create does not support speculative decoding")
68 + if platform.system() != "Linux" or platform.machine() != "x86_64":
69 + raise ValueError("snapshot create requires Linux x86_64")
70 +
71 +
72 +def run_create(args: argparse.Namespace) -> None:
73 + validate_create_args(args)
74 + from vllm.snapshot.controller import create_snapshot
75 +
76 + create_snapshot(args)
77 +
78 +
79 +def run_inspect(args: argparse.Namespace) -> None:
80 + from vllm.snapshot.manifest import inspect_snapshot
81 +
82 + print(json.dumps(inspect_snapshot(Path(args.snapshot_dir)), indent=2))
83 +
84 +
85 +def run_restore(args: argparse.Namespace) -> None:
86 + from vllm.snapshot.controller import restore_snapshot
87 +
88 + restore_snapshot(args)
89 +
90 +
91 +class SnapshotSubcommand(CLISubcommand):
92 + """The `snapshot` subcommand for the vLLM CLI."""
93 +
94 + name = "snapshot"
95 +
96 + def __init__(self, *, create_requested: bool = False) -> None:
97 + self.create_requested = create_requested
98 +
99 + @staticmethod
100 + def cmd(args: argparse.Namespace) -> None:
101 + args.snapshot_dispatch(args)
102 +
103 + def subparser_init(
104 + self, subparsers: argparse._SubParsersAction
105 + ) -> FlexibleArgumentParser:
106 + parser = subparsers.add_parser(
107 + self.name,
108 + help="Create, inspect, or restore an initialized vLLM snapshot.",
109 + usage="vllm snapshot <create|inspect|restore> [options]",
110 + )
111 + actions = parser.add_subparsers(required=True, dest="snapshot_action")
112 +
113 + create_parser = actions.add_parser(
114 + "create", help="Create a snapshot from an initialized TP1 engine."
115 + )
116 + if self.create_requested:
117 + from vllm.entrypoints.launchers.cli_args import make_arg_parser
118 +
119 + create_parser = make_arg_parser(create_parser)
120 + else:
121 + create_parser.add_argument("model_tag", nargs="?")
122 + create_parser.add_argument("--snapshot-dir", required=True)
123 + create_parser.set_defaults(snapshot_dispatch=run_create)
124 +
125 + inspect_parser = actions.add_parser(
126 + "inspect", help="Inspect a snapshot without restoring it."
127 + )
128 + inspect_parser.add_argument("snapshot_dir")
129 + inspect_parser.set_defaults(snapshot_dispatch=run_inspect)
130 +
131 + restore_parser = actions.add_parser(
132 + "restore", help="Restore a same-host TP1 snapshot."
133 + )
134 + restore_parser.add_argument("snapshot_dir")
135 + restore_parser.add_argument("--host", default=None)
136 + restore_parser.add_argument("--port", type=int, default=8000)
137 + restore_parser.set_defaults(snapshot_dispatch=run_restore)
138 +
139 + return parser
140 +
141 +
142 +def cmd_init(*, create_requested: bool = False) -> list[CLISubcommand]:
143 + return [SnapshotSubcommand(create_requested=create_requested)]
vllm/snapshot/__init__.py new
+4
@@ -0,0 +1,4 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Create and restore explicitly compatible vLLM engine snapshots."""
vllm/snapshot/controller.py new
+132
@@ -0,0 +1,132 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Create and restore transaction orchestration for local snapshots."""
4 +
5 +import argparse
6 +import shutil
7 +import sys
8 +from pathlib import Path
9 +
10 +from vllm.snapshot.manifest import (
11 + read_manifest,
12 + validate_artifact_root,
13 + validate_identity,
14 +)
15 +from vllm.snapshot.runtime import (
16 + LocalSnapshotTools,
17 + SnapshotCreateError,
18 + SnapshotRestoreError,
19 + _error_detail,
20 +)
21 +from vllm.snapshot.types import Oracle, oracles_match
22 +
23 +
24 +def _child_engine_argv(
25 + args: argparse.Namespace, engine_argv: tuple[str, ...] | None
26 +) -> tuple[str, ...]:
27 + """Build the child argv from the first argv source that is present.
28 +
29 + In order: an explicit ``engine_argv`` from the caller, the process argv
30 + tail when invoked as ``vllm snapshot create ...``, then the bare model from
31 + the parsed args. Whichever source wins, snapshot options are stripped and
32 + sleep mode is appended.
33 + """
34 + if not engine_argv:
35 + process_argv = tuple(sys.argv[1:])
36 + if process_argv[:2] == ("snapshot", "create"):
37 + engine_argv = process_argv[2:]
38 + else:
39 + model = getattr(args, "model_tag", None) or getattr(args, "model", None)
40 + if not model:
41 + raise SnapshotCreateError("snapshot create requires a model argument")
42 + engine_argv = (str(model),)
43 + child_argv: list[str] = []
44 + iterator = iter(engine_argv)
45 + for item in iterator:
46 + option, separator, _ = item.partition("=")
47 + option = option.replace("_", "-")
48 + if len(option) > 2 and "--snapshot-dir".startswith(option):
49 + if not separator:
50 + next(iterator, None)
51 + else:
52 + child_argv.append(item)
53 + if "--enable-sleep-mode" not in child_argv:
54 + child_argv.append("--enable-sleep-mode")
55 + return tuple(child_argv)
56 +
57 +
58 +def create_snapshot(
59 + args: argparse.Namespace,
60 + *,
61 + engine_argv: tuple[str, ...] | None = None,
62 + tools: LocalSnapshotTools | None = None,
63 +) -> None:
64 + target = Path(args.snapshot_dir).absolute()
65 + toolset = tools or LocalSnapshotTools()
66 + toolset.preflight("create", target)
67 + if target.exists() or target.is_symlink():
68 + raise SnapshotCreateError(f"snapshot target already exists: {target}")
69 + validate_artifact_root(target, creating=True)
70 + target.parent.mkdir(mode=0o700, parents=False, exist_ok=True)
71 + validate_artifact_root(target.parent, creating=False)
72 + target.mkdir(mode=0o700)
73 + validate_artifact_root(target, creating=False)
74 + published = False
75 + root_pid: int | None = None
76 + try:
77 + child_argv = _child_engine_argv(args, engine_argv)
78 + root_pid = toolset.launch_child(target, child_argv)
79 + oracle = toolset.wait_ready(target, root_pid)
80 + inventory = toolset.inventory(root_pid, target)
81 + toolset.dump(target, inventory)
82 + toolset.verify_dead(inventory)
83 + manifest = toolset.make_manifest(args, child_argv, inventory, oracle, target)
84 + toolset.publish(target, manifest)
85 + published = True
86 + except BaseException:
87 + if root_pid is not None:
88 + toolset.abort_create(root_pid)
89 + raise
90 + finally:
91 + # Runs even when abort_create raised, so no partial artifact is left.
92 + if not published and target.exists():
93 + shutil.rmtree(target)
94 +
95 +
96 +def restore_snapshot(
97 + args: argparse.Namespace, *, tools: LocalSnapshotTools | None = None
98 +) -> None:
99 + artifact = Path(args.snapshot_dir).absolute()
100 + toolset = tools or LocalSnapshotTools()
101 + toolset.preflight("restore", artifact)
102 + manifest = read_manifest(artifact)
103 + validate_identity(manifest, toolset.current_identity(manifest.gpu_uuid))
104 +
105 + root_pid: int | None = None
106 + try:
107 + root_pid = toolset.restore(artifact, manifest)
108 + toolset.release(artifact, args.host, args.port)
109 + toolset.wait_listener(root_pid, args.host, args.port)
110 + actual = toolset.request_oracle(args.host, args.port, manifest)
111 + expected = Oracle(
112 + token_ids=manifest.oracle_token_ids,
113 + text=manifest.oracle_text,
114 + sampled_token_logprob=manifest.oracle_sampled_token_logprob,
115 + )
116 + if not oracles_match(expected, actual):
117 + raise SnapshotRestoreError(
118 + f"snapshot oracle mismatch: expected={expected!r}, actual={actual!r}"
119 + )
120 + toolset.complete_restore(root_pid)
121 + except BaseException as error:
122 + if root_pid is not None:
123 + try:
124 + toolset.cleanup(root_pid)
125 + except BaseException as cleanup_error:
126 + raise SnapshotRestoreError(
127 + f"snapshot restore failed: {_error_detail(error)}; "
128 + f"cleanup failed: {_error_detail(cleanup_error)}"
129 + ) from error
130 + if isinstance(error, SnapshotRestoreError):
131 + raise
132 + raise SnapshotRestoreError(str(error)) from error
vllm/snapshot/manifest.py new
+277
@@ -0,0 +1,277 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import json
5 +import math
6 +import os
7 +import stat
8 +from pathlib import Path
9 +from typing import Annotated, Any, Literal
10 +
11 +from pydantic import (
12 + BaseModel,
13 + ConfigDict,
14 + Field,
15 + StrictFloat,
16 + StrictInt,
17 + StrictStr,
18 + ValidationError,
19 + field_validator,
20 + model_validator,
21 +)
22 +from typing_extensions import Self
23 +
24 +
25 +class SnapshotCompatibilityError(RuntimeError):
26 + """The snapshot identity does not match the requested runtime."""
27 +
28 +
29 +class SnapshotSecurityError(RuntimeError):
30 + """The snapshot artifact does not meet the private-path contract."""
31 +
32 +
33 +_MAX_PID = 2**31 - 1
34 +_PositivePid = Annotated[int, Field(gt=0, le=_MAX_PID)]
35 +_OracleTokenId = Annotated[int, Field(ge=0)]
36 +
37 +
38 +class SnapshotRuntimeIdentity(BaseModel):
39 + model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
40 +
41 + vllm_version: str
42 + python_version: str
43 + torch_version: str
44 + cuda_runtime: str
45 + driver_version: str
46 + criu_version: str
47 + cuda_checkpoint_sha256: str
48 + kernel_release: str
49 + host_id: str
50 + gpu_name: str
51 + gpu_uuid: str
52 + environment: tuple[tuple[str, str], ...]
53 +
54 +
55 +class SnapshotManifest(SnapshotRuntimeIdentity):
56 + schema_version: Literal[1]
57 + boundary: Literal["post-engine-init-reloadable-state-released"]
58 + created_at: str
59 + artifact_bytes: Annotated[int, Field(ge=0)]
60 + model: str
61 + served_model_name: str
62 + model_revision: str
63 + tokenizer_revision: str
64 + engine_argv: tuple[str, ...]
65 + process_tree: Annotated[tuple[_PositivePid, ...], Field(min_length=1)]
66 + cuda_holders: Annotated[tuple[_PositivePid, ...], Field(min_length=1)]
67 + external_cache_files: tuple[tuple[str, str], ...] = ()
68 + oracle_token_ids: Annotated[
69 + tuple[_OracleTokenId, ...], Field(min_length=1, max_length=1)
70 + ]
71 + oracle_text: str
72 + oracle_sampled_token_logprob: float
73 +
74 + @field_validator("schema_version", mode="before")
75 + @classmethod
76 + def _require_schema_version_one(cls, value: Any) -> Any:
77 + if type(value) is not int or value != 1:
78 + raise ValueError("schema version must be integer 1")
79 + return value
80 +
81 + @field_validator("oracle_sampled_token_logprob", mode="before")
82 + @classmethod
83 + def _require_finite_float(cls, value: Any) -> Any:
84 + if type(value) is not float or not math.isfinite(value):
85 + raise ValueError("sampled token log probability must be a finite float")
86 + return value
87 +
88 + @model_validator(mode="after")
89 + def _validate_process_relationships(self) -> Self:
90 + if len(set(self.process_tree)) != len(self.process_tree):
91 + raise ValueError("process_tree contains duplicate PIDs")
92 + if len(set(self.cuda_holders)) != len(self.cuda_holders):
93 + raise ValueError("cuda_holders contains duplicate PIDs")
94 + if not set(self.cuda_holders).issubset(self.process_tree):
95 + raise ValueError("cuda_holders must be a subset of process_tree")
96 + return self
97 +
98 +
99 +class ReleaseMarker(BaseModel):
100 + """Release marker the controller writes; unknown keys are ignored."""
101 +
102 + release: Literal[True]
103 + host: StrictStr | None = None
104 + port: Annotated[StrictInt, Field(ge=1, le=65535)]
105 +
106 + @field_validator("release", mode="before")
107 + @classmethod
108 + def _require_json_true(cls, value: Any) -> Any:
109 + # A literal accepts 1 as True; the barrier requires JSON true itself.
110 + if value is not True:
111 + raise ValueError("release must be JSON true")
112 + return value
113 +
114 +
115 +class ReadyMarker(BaseModel):
116 + """Ready marker the child writes; unknown keys are ignored."""
117 +
118 + token_ids: tuple[StrictInt, ...]
119 + text: StrictStr
120 + sampled_token_logprob: StrictFloat | StrictInt
121 +
122 +
123 +def validate_identity(
124 + expected: SnapshotRuntimeIdentity, actual: SnapshotRuntimeIdentity
125 +) -> None:
126 + """Require an exact match for every field that can affect compatibility."""
127 + # ``expected`` may be a manifest, so compare the identity fields only.
128 + fields = set(SnapshotRuntimeIdentity.model_fields)
129 + expected_values = expected.model_dump(include=fields)
130 + actual_values = actual.model_dump(include=fields)
131 + differing = [
132 + name
133 + for name in SnapshotRuntimeIdentity.model_fields
134 + if expected_values[name] != actual_values[name]
135 + ]
136 + if differing:
137 + raise SnapshotCompatibilityError(f"snapshot mismatch: {', '.join(differing)}")
138 +
139 +
140 +def _existing_path_chain(path: Path):
141 + current = path
142 + while True:
143 + if current.exists() or current.is_symlink():
144 + yield current
145 + if current.parent == current:
146 + break
147 + current = current.parent
148 +
149 +
150 +def _is_trusted_sticky_directory(path_stat: os.stat_result) -> bool:
151 + return (
152 + path_stat.st_uid == 0
153 + and stat.S_ISDIR(path_stat.st_mode)
154 + and bool(path_stat.st_mode & stat.S_ISVTX)
155 + )
156 +
157 +
158 +def validate_artifact_root(path: Path, *, creating: bool) -> None:
159 + """Validate that an artifact path cannot be replaced by another user."""
160 + path = Path(os.path.abspath(path))
161 + if path.is_symlink():
162 + raise SnapshotSecurityError(f"snapshot path is a symlink: {path}")
163 + if not creating and not path.exists():
164 + raise SnapshotSecurityError(f"snapshot path does not exist: {path}")
165 +
166 + effective_uid = os.geteuid()
167 + for component in _existing_path_chain(path):
168 + component_stat = component.lstat()
169 + if stat.S_ISLNK(component_stat.st_mode):
170 + raise SnapshotSecurityError(
171 + f"snapshot path contains a symlink: {component}"
172 + )
173 + if component_stat.st_uid not in (0, effective_uid):
174 + raise SnapshotSecurityError(
175 + f"snapshot path has an untrusted owner: {component}"
176 + )
177 + if component_stat.st_mode & (
178 + stat.S_IWGRP | stat.S_IWOTH
179 + ) and not _is_trusted_sticky_directory(component_stat):
180 + raise SnapshotSecurityError(
181 + f"snapshot path has a group- or world-writable ancestor: {component}"
182 + )
183 +
184 + if path.exists():
185 + path_stat = path.lstat()
186 + if not stat.S_ISDIR(path_stat.st_mode):
187 + raise SnapshotSecurityError(f"snapshot path is not a directory: {path}")
188 + if path_stat.st_uid != effective_uid:
189 + raise SnapshotSecurityError(
190 + f"snapshot directory is not owned by the current user: {path}"
191 + )
192 + if stat.S_IMODE(path_stat.st_mode) & 0o077:
193 + raise SnapshotSecurityError(
194 + f"snapshot directory must have mode 0700: {path}"
195 + )
196 +
197 +
198 +def _validation_path(error: ValidationError) -> str:
199 + detail = error.errors(include_url=False)[0]
200 + location = detail["loc"]
201 + if not location:
202 + return "JSON" if detail["type"] == "json_invalid" else "root"
203 + return ".".join(map(str, location))
204 +
205 +
206 +def _fsync_directory(path: Path) -> None:
207 + directory_fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
208 + try:
209 + os.fsync(directory_fd)
210 + finally:
211 + os.close(directory_fd)
212 +
213 +
214 +def _write_json_atomic(path: Path, payload: object, *, overwrite: bool = False) -> None:
215 + temporary = path.with_name(f"{path.name}.tmp")
216 + if overwrite:
217 + temporary.unlink(missing_ok=True)
218 + elif path.exists() or path.is_symlink():
219 + raise FileExistsError(path)
220 + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
221 + descriptor = os.open(temporary, flags, 0o600)
222 + try:
223 + with os.fdopen(descriptor, "wb") as output:
224 + output.write(
225 + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
226 + )
227 + output.flush()
228 + os.fsync(output.fileno())
229 + os.replace(temporary, path)
230 + _fsync_directory(path.parent)
231 + except BaseException:
232 + temporary.unlink(missing_ok=True)
233 + raise
234 +
235 +
236 +def write_manifest_atomic(path: Path, manifest: SnapshotManifest) -> None:
237 + """Write a new private manifest and make its directory entry durable."""
238 + path = Path(path)
239 + validate_artifact_root(path, creating=False)
240 + destination = path / "manifest.json"
241 + try:
242 + _write_json_atomic(destination, manifest.model_dump(mode="json"))
243 + except FileExistsError as error:
244 + raise SnapshotSecurityError(
245 + f"manifest already exists: {destination}"
246 + ) from error
247 +
248 +
249 +def read_manifest(path: Path) -> SnapshotManifest:
250 + """Read a snapshot published by atomically installing its manifest."""
251 + path = Path(path)
252 + validate_artifact_root(path, creating=False)
253 + manifest_path = path / "manifest.json"
254 + if manifest_path.is_symlink():
255 + raise SnapshotSecurityError(f"manifest is a symlink: {manifest_path}")
256 + flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
257 + try:
258 + file_descriptor = os.open(manifest_path, flags)
259 + except FileNotFoundError as error:
260 + raise SnapshotSecurityError(
261 + f"manifest does not exist: {manifest_path}"
262 + ) from error
263 + try:
264 + with os.fdopen(file_descriptor, encoding="utf-8") as manifest_file:
265 + payload = manifest_file.read()
266 + except UnicodeDecodeError as error:
267 + raise SnapshotCompatibilityError("invalid snapshot manifest: JSON") from error
268 + try:
269 + return SnapshotManifest.model_validate_json(payload, strict=True)
270 + except ValidationError as error:
271 + raise SnapshotCompatibilityError(
272 + f"invalid snapshot manifest: {_validation_path(error)}"
273 + ) from error
274 +
275 +
276 +def inspect_snapshot(path: Path) -> dict[str, object]:
277 + return read_manifest(path).model_dump(mode="json")
vllm/snapshot/runtime.py new
+1246
@@ -0,0 +1,1246 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Host runtime operations for local snapshots."""
4 +
5 +import argparse
6 +import ast
7 +import hashlib
8 +import importlib.metadata
9 +import importlib.util
10 +import json
11 +import os
12 +import platform
13 +import select
14 +import shutil
15 +import signal
16 +import socket
17 +import stat
18 +import struct
19 +import subprocess
20 +import sys
21 +import tempfile
22 +import time
23 +import urllib.request
24 +from collections.abc import Iterable
25 +from contextlib import suppress
26 +from dataclasses import dataclass
27 +from datetime import datetime, timezone
28 +from pathlib import Path
29 +from typing import NoReturn
30 +
31 +import regex as re
32 +from pydantic import ValidationError
33 +
34 +from vllm.snapshot.manifest import (
35 + ReadyMarker,
36 + SnapshotManifest,
37 + SnapshotRuntimeIdentity,
38 + _fsync_directory,
39 + _validation_path,
40 + _write_json_atomic,
41 + validate_artifact_root,
42 + write_manifest_atomic,
43 +)
44 +from vllm.snapshot.types import Oracle
45 +
46 +
47 +class SnapshotCreateError(RuntimeError):
48 + """Snapshot creation did not produce a complete artifact."""
49 +
50 +
51 +class SnapshotRestoreError(RuntimeError):
52 + """Snapshot restore failed."""
53 +
54 +
55 +_COMMON_CRIU_OPTIONS = (
56 + "--shell-job",
57 + "--ext-unix-sk",
58 + "--tcp-established",
59 + "--link-remap",
60 + "--file-locks",
61 +)
62 +_TCP_TABLES = (
63 + ("AF_INET", Path("/proc/net/tcp")),
64 + ("AF_INET6", Path("/proc/net/tcp6")),
65 +)
66 +
67 +_CACHE_REMEDY = "recreate the snapshot or restore with the same cache state"
68 +# Recorded instead of a digest for files whose content is volatile by design.
69 +_CACHE_LOCKED = "locked"
70 +# Recorded instead of a digest for a file capture found but could not read.
71 +_CACHE_UNREADABLE = "unreadable"
72 +
73 +
74 +def _generated_cache_prefixes() -> tuple[str, ...]:
75 + """Path prefixes of the generated caches an engine process keeps open.
76 +
77 + vLLM redirects Inductor and Triton under ``VLLM_CACHE_ROOT`` when it
78 + compiles, but a process can still hold their upstream defaults open.
79 + A cache root is resolved and separator-terminated so that it still matches
80 + the kernel-canonical descriptor targets under a symlinked root, and so that
81 + a sibling directory sharing its name does not. Inductor's default location
82 + is a bare prefix instead, because its directory carries a per-user suffix.
83 + """
84 + import vllm.envs as envs
85 +
86 + inductor_root = os.environ.get("TORCHINDUCTOR_CACHE_DIR")
87 + roots = tuple(
88 + f"{os.path.realpath(os.path.expanduser(root)).rstrip(os.sep)}{os.sep}"
89 + for root in (
90 + envs.VLLM_CACHE_ROOT,
91 + os.environ.get("TRITON_CACHE_DIR") or "~/.triton/cache",
92 + "~/.cache/flashinfer",
93 + *((inductor_root,) if inductor_root else ()),
94 + )
95 + )
96 + if inductor_root:
97 + return roots
98 + return roots + (
99 + os.path.join(os.path.realpath(tempfile.gettempdir()), "torchinductor_"),
100 + )
101 +
102 +
103 +def _error_detail(error: BaseException) -> str:
104 + message = str(error)
105 + return f"{type(error).__name__}: {message}" if message else type(error).__name__
106 +
107 +
108 +def _close_pidfds(handles: Iterable[tuple[int, int]]) -> None:
109 + for _pid, pidfd in handles:
110 + with suppress(OSError):
111 + os.close(pidfd)
112 +
113 +
114 +@dataclass(frozen=True)
115 +class ProcessInventory:
116 + root_pid: int
117 + process_tree: tuple[int, ...]
118 + cuda_holders: tuple[int, ...]
119 + gpu_uuid: str
120 + external_cache_files: tuple[tuple[str, str], ...] = ()
121 +
122 +
123 +@dataclass(frozen=True)
124 +class _TcpSocketRecord:
125 + family: str
126 + local_raw: str
127 + remote_raw: str
128 + inode: int
129 +
130 +
131 +def _process_comm(pid: int) -> str:
132 + try:
133 + return (Path("/proc") / str(pid) / "comm").read_text().strip()
134 + except OSError:
135 + return "unknown"
136 +
137 +
138 +def _decode_endpoint(family: str, raw: str) -> str:
139 + """Render a ``/proc/net/tcp`` ``HEXADDR:HEXPORT`` field as ``ip:port``.
140 +
141 + The kernel prints the address as little-endian 32-bit words: one for IPv4,
142 + four in order for IPv6.
143 + """
144 + hex_address, _, hex_port = raw.partition(":")
145 + port = int(hex_port, 16)
146 + if family == "AF_INET":
147 + packed = struct.pack("<I", int(hex_address, 16))
148 + return f"{socket.inet_ntop(socket.AF_INET, packed)}:{port}"
149 + packed = struct.pack(
150 + "<4I", *(int(hex_address[index : index + 8], 16) for index in (0, 8, 16, 24))
151 + )
152 + return f"[{socket.inet_ntop(socket.AF_INET6, packed)}]:{port}"
153 +
154 +
155 +def _validate_tcp_connections(
156 + records: tuple[_TcpSocketRecord, ...], socket_owners: dict[int, int]
157 +) -> None:
158 + owned = tuple(record for record in records if record.inode in socket_owners)
159 + endpoints = {
160 + (record.family, record.local_raw, record.remote_raw) for record in owned
161 + }
162 + offenders = tuple(
163 + record
164 + for record in owned
165 + if (record.family, record.remote_raw, record.local_raw) not in endpoints
166 + )
167 + if not offenders:
168 + return
169 + named = "; ".join(
170 + f"pid {socket_owners[record.inode]} "
171 + f"({_process_comm(socket_owners[record.inode])}) "
172 + f"{_decode_endpoint(record.family, record.local_raw)} -> "
173 + f"{_decode_endpoint(record.family, record.remote_raw)}"
174 + for record in offenders[:5]
175 + )
176 + remaining = f" (+{len(offenders) - 5} more)" if len(offenders) > 5 else ""
177 + raise SnapshotCreateError(
178 + f"snapshot tree has an external established TCP connection: {named}{remaining}"
179 + )
180 +
181 +
182 +class LocalSnapshotTools:
183 + """Host-owned CRIU and cuda-checkpoint implementation."""
184 +
185 + def __init__(self) -> None:
186 + self.criu = shutil.which("criu") or "criu"
187 + self.cuda_checkpoint = shutil.which("cuda-checkpoint") or "cuda-checkpoint"
188 + self.nvidia_smi = shutil.which("nvidia-smi") or "nvidia-smi"
189 + plugin_dir = os.environ.get("CRIU_CUDA_PLUGIN_DIR", "")
190 + self.plugin_dir = Path(plugin_dir) if plugin_dir else None
191 + self.shm_dir = Path("/dev/shm")
192 + self.timeout_s = float(os.environ.get("VLLM_SNAPSHOT_TIMEOUT_S", "900"))
193 + if not 0 < self.timeout_s < float("inf"):
194 + raise ValueError("VLLM_SNAPSHOT_TIMEOUT_S must be a positive finite number")
195 + self._children: dict[int, subprocess.Popen[bytes]] = {}
196 + self._restored_processes: dict[int, tuple[tuple[int, int], ...]] = {}
197 +
198 + def _privileged(self) -> list[str]:
199 + return [] if os.geteuid() == 0 else ["sudo", "-n"]
200 +
201 + def _run(
202 + self,
203 + command: list[str],
204 + *,
205 + timeout: float | None = None,
206 + ) -> subprocess.CompletedProcess[str]:
207 + return subprocess.run(
208 + command,
209 + check=True,
210 + capture_output=True,
211 + text=True,
212 + timeout=timeout or self.timeout_s,
213 + )
214 +
215 + def preflight(self, action: str, artifact: Path) -> None:
216 + if platform.system() != "Linux" or platform.machine() != "x86_64":
217 + raise RuntimeError("snapshot requires Linux x86_64")
218 + if action == "restore" and (
219 + not callable(getattr(os, "pidfd_open", None))
220 + or not callable(getattr(signal, "pidfd_send_signal", None))
221 + ):
222 + raise RuntimeError("snapshot requires Linux pidfd support")
223 + for command in (self.criu, self.cuda_checkpoint, self.nvidia_smi):
224 + if shutil.which(command) is None:
225 + raise RuntimeError(f"snapshot dependency not found: {command}")
226 + if self.plugin_dir is None:
227 + raise RuntimeError("CRIU_CUDA_PLUGIN_DIR is required")
228 + if not (self.plugin_dir / "cuda_plugin.so").is_file():
229 + raise RuntimeError(f"CRIU CUDA plugin not found in {self.plugin_dir}")
230 + if os.geteuid() != 0:
231 + self._run(["sudo", "-n", "true"], timeout=5)
232 + if action == "restore":
233 + validate_artifact_root(artifact, creating=False)
234 + else:
235 + validate_artifact_root(artifact, creating=True)
236 +
237 + def launch_child(
238 + self,
239 + workdir: Path,
240 + engine_argv: tuple[str, ...],
241 + ) -> int:
242 + command = [
243 + sys.executable,
244 + "-m",
245 + "vllm.snapshot.server",
246 + "--ready-file",
247 + str(workdir / "ready.json"),
248 + "--release-file",
249 + str(workdir / "release.json"),
250 + "--release-timeout-s",
251 + str(self.timeout_s),
252 + "--",
253 + *engine_argv,
254 + ]
255 + child_environment = os.environ.copy()
256 + child_environment.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
257 + with (workdir / "child.log").open("wb") as log_file:
258 + process = subprocess.Popen(
259 + command,
260 + env=child_environment,
261 + stdin=subprocess.DEVNULL,
262 + stdout=log_file,
263 + stderr=subprocess.STDOUT,
264 + start_new_session=True,
265 + )
266 + self._children[process.pid] = process
267 + return process.pid
268 +
269 + def wait_ready(self, workdir: Path, root_pid: int) -> Oracle:
270 + ready_file = workdir / "ready.json"
271 + deadline = time.monotonic() + self.timeout_s
272 + waitid = os.waitid # type: ignore[attr-defined]
273 + while time.monotonic() < deadline:
274 + if ready_file.is_file() and ready_file.stat().st_size:
275 + try:
276 + marker = ReadyMarker.model_validate_json(ready_file.read_text())
277 + except ValidationError as error:
278 + raise SnapshotCreateError(
279 + f"snapshot ready marker invalid: {_validation_path(error)}"
280 + ) from error
281 + return Oracle(
282 + token_ids=marker.token_ids,
283 + text=marker.text,
284 + sampled_token_logprob=marker.sampled_token_logprob,
285 + )
286 + process = self._children.get(root_pid)
287 + if process is not None and waitid(
288 + os.P_PID,
289 + root_pid,
290 + os.WEXITED | os.WNOHANG | os.WNOWAIT,
291 + ):
292 + log = (workdir / "child.log").read_text(errors="replace")
293 + raise SnapshotCreateError(f"snapshot child exited before ready:\n{log}")
294 + try:
295 + os.kill(root_pid, 0)
296 + except ProcessLookupError as error:
297 + log = (workdir / "child.log").read_text(errors="replace")
298 + raise SnapshotCreateError(
299 + f"snapshot child exited before ready:\n{log}"
300 + ) from error
301 + time.sleep(0.05)
302 + raise SnapshotCreateError("snapshot child did not become ready before timeout")
303 +
304 + def _tree_pids(self, root_pid: int) -> tuple[int, ...]:
305 + rows = self._run(["ps", "-eo", "pid=,ppid="]).stdout.splitlines()
306 + children: dict[int, list[int]] = {}
307 + for row in rows:
308 + pid, parent_pid = map(int, row.split())
309 + children.setdefault(parent_pid, []).append(pid)
310 + stack = [root_pid]
311 + found: list[int] = []
312 + while stack:
313 + pid = stack.pop()
314 + if pid in found:
315 + continue
316 + found.append(pid)
317 + stack.extend(children.get(pid, []))
318 + return tuple(found)
319 +
320 + def _cuda_process_rows(self) -> tuple[str, ...]:
321 + output = self._run(
322 + [
323 + self.nvidia_smi,
324 + "--query-compute-apps=pid,gpu_uuid",
325 + "--format=csv,noheader,nounits",
326 + ]
327 + ).stdout
328 + return tuple(line for line in output.splitlines() if line.strip())
329 +
330 + def _cuda_pids(self, rows: tuple[str, ...]) -> tuple[int, ...]:
331 + return tuple(sorted({int(line.split(",", 1)[0].strip()) for line in rows}))
332 +
333 + def _descriptor_targets(self, pid: int) -> tuple[str, ...]:
334 + descriptor_dir = Path("/proc") / str(pid) / "fd"
335 + try:
336 + descriptors = tuple(descriptor_dir.iterdir())
337 + except FileNotFoundError as error:
338 + raise SnapshotCreateError(
339 + f"snapshot process exited during descriptor inventory: {pid}"
340 + ) from error
341 + targets: list[str] = []
342 + for descriptor in descriptors:
343 + try:
344 + targets.append(os.readlink(descriptor))
345 + except FileNotFoundError:
346 + continue
347 + return tuple(targets)
348 +
349 + def _descriptor_inventory(
350 + self, process_tree: tuple[int, ...]
351 + ) -> tuple[tuple[int, ...], dict[int, int]]:
352 + """Collect io_uring holders and each socket inode's owning pid."""
353 + io_uring_pids: list[int] = []
354 + socket_owners: dict[int, int] = {}
355 + for pid in process_tree:
356 + targets = self._descriptor_targets(pid)
357 + if "anon_inode:[io_uring]" in targets:
358 + io_uring_pids.append(pid)
359 + for target in targets:
360 + if target.startswith("socket:[") and target.endswith("]"):
361 + # Shared descriptors name one holder; any of them locates
362 + # the connection for the operator.
363 + socket_owners.setdefault(int(target[len("socket:[") : -1]), pid)
364 + return tuple(io_uring_pids), socket_owners
365 +
366 + def _tcp_records(self) -> tuple[_TcpSocketRecord, ...]:
367 + records: list[_TcpSocketRecord] = []
368 + for family, table in _TCP_TABLES:
369 + if not table.exists():
370 + continue
371 + for row in table.read_text().splitlines()[1:]:
372 + fields = row.split()
373 + if fields[3] != "01":
374 + continue
375 + records.append(
376 + _TcpSocketRecord(
377 + family=family,
378 + local_raw=fields[1],
379 + remote_raw=fields[2],
380 + inode=int(fields[9]),
381 + )
382 + )
383 + return tuple(records)
384 +
385 + def _cache_fingerprint(self, path: Path) -> str | None:
386 + """Digest a cache file, or mark the ones whose content is volatile."""
387 + try:
388 + if not stat.S_ISREG(path.stat().st_mode):
389 + return None
390 + except OSError:
391 + return None
392 + if path.name.endswith(".lock"):
393 + return _CACHE_LOCKED
394 + try:
395 + return self._sha256(path)
396 + except OSError:
397 + # Restore can only check that this one still exists, so record the
398 + # downgrade rather than passing it off as a volatile file.
399 + return _CACHE_UNREADABLE
400 +
401 + def _external_cache_files(
402 + self, process_tree: tuple[int, ...], artifact: Path
403 + ) -> tuple[tuple[str, str], ...]:
404 + """Fingerprint the generated-cache files the captured tree holds open.
405 +
406 + CRIU reopens these by path on restore, so a JIT cache that rotates
407 + between capture and restore fails the restore permanently. Recording
408 + them lets restore fail early and name the file that changed. Files
409 + inside the artifact are excluded: they belong to the snapshot itself.
410 + Deleted targets are excluded too, because CRIU remaps those links.
411 + """
412 + prefixes = _generated_cache_prefixes()
413 + artifact_prefix = f"{os.path.abspath(artifact)}{os.sep}"
414 + candidates: set[str] = set()
415 + for pid in process_tree:
416 + for target in self._descriptor_targets(pid):
417 + if target.endswith(" (deleted)") or target.startswith(artifact_prefix):
418 + continue
419 + if target.endswith(".lock") or target.startswith(prefixes):
420 + candidates.add(target)
421 + recorded = (
422 + (path, self._cache_fingerprint(Path(path))) for path in sorted(candidates)
423 + )
424 + return tuple(
425 + (path, fingerprint)
426 + for path, fingerprint in recorded
427 + if fingerprint is not None
428 + )
429 +
430 + def inventory(self, root_pid: int, artifact: Path) -> ProcessInventory:
431 + process_tree = self._tree_pids(root_pid)
432 + cuda_rows = self._cuda_process_rows()
433 + cuda_holders = tuple(
434 + pid for pid in self._cuda_pids(cuda_rows) if pid in process_tree
435 + )
436 + if not cuda_holders:
437 + raise SnapshotCreateError("snapshot tree has no CUDA-holding process")
438 + io_uring_pids, socket_owners = self._descriptor_inventory(process_tree)
439 + if io_uring_pids:
440 + raise SnapshotCreateError(
441 + "snapshot tree owns io_uring state that CRIU cannot dump; set "
442 + "kernel.io_uring_disabled=1 before starting an unprivileged "
443 + "vLLM process, or use =2 to disable it host-wide "
444 + f"(pids: {', '.join(map(str, io_uring_pids))})"
445 + )
446 + _validate_tcp_connections(self._tcp_records(), socket_owners)
447 + gpu_uuid = self._gpu_uuid_for_pids(cuda_holders, cuda_rows)
448 + return ProcessInventory(
449 + root_pid=root_pid,
450 + process_tree=process_tree,
451 + cuda_holders=cuda_holders,
452 + gpu_uuid=gpu_uuid,
453 + external_cache_files=self._external_cache_files(process_tree, artifact),
454 + )
455 +
456 + def _gpu_uuid_for_pids(self, pids: tuple[int, ...], rows: tuple[str, ...]) -> str:
457 + wanted = set(pids)
458 + matches: set[str] = set()
459 + try:
460 + for line in rows:
461 + pid, gpu_uuid = (item.strip() for item in line.split(",", 1))
462 + if int(pid) in wanted:
463 + matches.add(gpu_uuid)
464 + except (TypeError, ValueError) as error:
465 + raise SnapshotCreateError(
466 + "could not identify the GPU used by the snapshot process tree"
467 + ) from error
468 + if len(matches) != 1:
469 + raise SnapshotCreateError(
470 + "snapshot process tree must use exactly one identifiable GPU"
471 + )
472 + return matches.pop()
473 +
474 + def _criu(
475 + self,
476 + action: str,
477 + artifact: Path,
478 + arguments: list[str],
479 + ) -> None:
480 + if self.plugin_dir is None:
481 + raise RuntimeError("CRIU_CUDA_PLUGIN_DIR is required")
482 + env = os.environ.copy()
483 + env["PATH"] = f"{Path(self.cuda_checkpoint).parent}:{env['PATH']}"
484 + env["VLLM_SNAPSHOT_IGNORED_FD_MAP"] = str(artifact / "ignored-nvidia-fds.tsv")
485 + env["RIG_CRIU_IGNORED_FD_MAP"] = env["VLLM_SNAPSHOT_IGNORED_FD_MAP"]
486 + command = [
487 + *self._privileged(),
488 + "env",
489 + f"PATH={env['PATH']}",
490 + f"VLLM_SNAPSHOT_IGNORED_FD_MAP={env['VLLM_SNAPSHOT_IGNORED_FD_MAP']}",
491 + f"RIG_CRIU_IGNORED_FD_MAP={env['RIG_CRIU_IGNORED_FD_MAP']}",
492 + self.criu,
493 + action,
494 + "-L",
495 + str(self.plugin_dir),
496 + *arguments,
497 + ]
498 + self._run(command)
499 +
500 + def _record_child_log_size(self, artifact: Path) -> None:
501 + child_log = artifact / "child.log"
502 + size_path = artifact / "child.log.snapshot-size"
503 + payload = f"{child_log.stat().st_size}\n".encode()
504 + descriptor = os.open(
505 + size_path,
506 + os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0),
507 + 0o600,
508 + )
509 + with os.fdopen(descriptor, "wb") as output:
510 + output.write(payload)
511 + output.flush()
512 + os.fsync(output.fileno())
513 +
514 + def _link_remap_names(self) -> set[str]:
515 + return {
516 + path.name
517 + for path in self.shm_dir.glob("link_remap.*")
518 + if path.name.removeprefix("link_remap.").isdigit()
519 + }
520 +
521 + def _capture_link_remaps(self, artifact: Path, names: set[str]) -> None:
522 + if not names:
523 + return
524 + destination_dir = artifact / "link-remaps"
525 + destination_dir.mkdir(mode=0o700)
526 + for name in sorted(names):
527 + source = self.shm_dir / name
528 + source_fd = os.open(
529 + source,
530 + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0),
531 + )
532 + try:
533 + metadata = os.fstat(source_fd)
534 + if not stat.S_ISREG(metadata.st_mode) or metadata.st_uid != os.getuid():
535 + raise SnapshotCreateError(f"invalid CRIU link remap: {source}")
536 + destination_fd = os.open(
537 + destination_dir / name,
538 + os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0),
539 + 0o600,
540 + )
541 + with os.fdopen(destination_fd, "wb") as destination:
542 + with os.fdopen(os.dup(source_fd), "rb") as contents:
543 + shutil.copyfileobj(contents, destination)
544 + destination.flush()
545 + os.fsync(destination.fileno())
546 + finally:
547 + os.close(source_fd)
548 + source.unlink()
549 +
550 + def _rollback_link_remaps(self, created: dict[str, tuple[int, int]]) -> None:
551 + for name, identity in created.items():
552 + target = self.shm_dir / name
553 + try:
554 + metadata = target.lstat()
555 + except FileNotFoundError:
556 + continue
557 + if (metadata.st_dev, metadata.st_ino) == identity:
558 + target.unlink()
559 +
560 + def _restore_link_remaps(self, artifact: Path) -> dict[str, tuple[int, int]]:
561 + source_dir = artifact / "link-remaps"
562 + if not source_dir.exists():
563 + return {}
564 + created: dict[str, tuple[int, int]] = {}
565 + try:
566 + for source in sorted(source_dir.iterdir()):
567 + if (
568 + not source.name.removeprefix("link_remap.").isdigit()
569 + or source.is_symlink()
570 + or not source.is_file()
571 + ):
572 + raise SnapshotRestoreError(f"invalid saved link remap: {source}")
573 + payload = source.read_bytes()
574 + target = self.shm_dir / source.name
575 + try:
576 + descriptor = os.open(
577 + target,
578 + os.O_WRONLY
579 + | os.O_CREAT
580 + | os.O_EXCL
581 + | getattr(os, "O_NOFOLLOW", 0),
582 + 0o600,
583 + )
584 + except FileExistsError as error:
585 + try:
586 + descriptor = os.open(
587 + target,
588 + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0),
589 + )
590 + except OSError:
591 + raise SnapshotRestoreError(
592 + f"conflicting CRIU link remap exists: {target}"
593 + ) from error
594 + with os.fdopen(descriptor, "rb") as existing:
595 + metadata = os.fstat(existing.fileno())
596 + if (
597 + not stat.S_ISREG(metadata.st_mode)
598 + or metadata.st_uid != os.getuid()
599 + or existing.read() != payload
600 + ):
601 + raise SnapshotRestoreError(
602 + f"conflicting CRIU link remap exists: {target}"
603 + ) from error
604 + continue
605 + metadata = os.fstat(descriptor)
606 + created[source.name] = (metadata.st_dev, metadata.st_ino)
607 + with os.fdopen(descriptor, "wb") as output:
608 + output.write(payload)
609 + output.flush()
610 + os.fsync(output.fileno())
611 + except BaseException:
612 + self._rollback_link_remaps(created)
613 + raise
614 + return created
615 +
616 + def _reset_child_log(self, artifact: Path) -> None:
617 + size_path = artifact / "child.log.snapshot-size"
618 + descriptor = os.open(
619 + size_path,
620 + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0),
621 + )
622 + with os.fdopen(descriptor, "r") as source:
623 + size = int(source.read().strip())
624 + if size < 0:
625 + raise SnapshotRestoreError("captured child log size is invalid")
626 +
627 + child_log = artifact / "child.log"
628 + descriptor = os.open(
629 + child_log,
630 + os.O_WRONLY | getattr(os, "O_NOFOLLOW", 0),
631 + )
632 + try:
633 + metadata = os.fstat(descriptor)
634 + if not stat.S_ISREG(metadata.st_mode) or metadata.st_size < size:
635 + raise SnapshotRestoreError("captured child log is invalid")
636 + os.ftruncate(descriptor, size)
637 + os.fsync(descriptor)
638 + finally:
639 + os.close(descriptor)
640 +
641 + def dump(self, workdir: Path, inventory: ProcessInventory) -> None:
642 + images = workdir / "images"
643 + images.mkdir(mode=0o700)
644 + fd_map = workdir / "ignored-nvidia-fds.tsv"
645 + fd_map.touch(mode=0o600)
646 + remaps_before = self._link_remap_names()
647 + try:
648 + self._criu(
649 + "dump",
650 + workdir,
651 + [
652 + "--tree",
653 + str(inventory.root_pid),
654 + "--images-dir",
655 + str(images),
656 + "--log-file",
657 + "dump.log",
658 + *_COMMON_CRIU_OPTIONS,
659 + ],
660 + )
661 + self._capture_link_remaps(
662 + workdir,
663 + self._link_remap_names() - remaps_before,
664 + )
665 + except BaseException:
666 + for name in self._link_remap_names() - remaps_before:
667 + (self.shm_dir / name).unlink(missing_ok=True)
668 + raise
669 + self._record_child_log_size(workdir)
670 +
671 + def verify_dead(self, inventory: ProcessInventory) -> None:
672 + process = self._children.get(inventory.root_pid)
673 + if process is not None:
674 + process.wait(timeout=10)
675 + for pid in inventory.process_tree:
676 + try:
677 + os.kill(pid, 0)
678 + except ProcessLookupError:
679 + continue
680 + raise SnapshotCreateError(f"process survived CRIU dump: {pid}")
681 + self._children.pop(inventory.root_pid, None)
682 +
683 + def _version(self, command: list[str]) -> str:
684 + try:
685 + result = self._run(command, timeout=10)
686 + except subprocess.CalledProcessError as error:
687 + return (error.stdout or error.stderr).splitlines()[0]
688 + return result.stdout.splitlines()[0] if result.stdout else "unknown"
689 +
690 + def _sha256(self, path: Path) -> str:
691 + digest = hashlib.sha256()
692 + with path.open("rb") as source:
693 + for chunk in iter(lambda: source.read(1024 * 1024), b""):
694 + digest.update(chunk)
695 + return digest.hexdigest()
696 +
697 + def _plugin_identity(self) -> tuple[tuple[str, str], ...]:
698 + if self.plugin_dir is None:
699 + return ()
700 + return tuple(
701 + (f"snapshot_plugin_sha256:{path.name}", self._sha256(path))
702 + for path in sorted(self.plugin_dir.glob("*.so"))
703 + )
704 +
705 + def _torch_identity(self) -> tuple[str, str]:
706 + torch_version = importlib.metadata.version("torch")
707 + torch_spec = importlib.util.find_spec("torch")
708 + locations = torch_spec and torch_spec.submodule_search_locations
709 + if not locations:
710 + raise RuntimeError("installed torch package could not be located")
711 + version_path = Path(next(iter(locations))) / "version.py"
712 + # Parse the version file rather than importing torch: snapshot create's
713 + # CLI path stays import-light (vllm/entrypoints/cli/snapshot.py).
714 + syntax = ast.parse(version_path.read_text(), filename=str(version_path))
715 + cuda_runtime = ""
716 + for statement in syntax.body:
717 + target = None
718 + value = None
719 + if isinstance(statement, ast.Assign) and len(statement.targets) == 1:
720 + target = statement.targets[0]
721 + value = statement.value
722 + elif isinstance(statement, ast.AnnAssign):
723 + target = statement.target
724 + value = statement.value
725 + if isinstance(target, ast.Name) and target.id == "cuda" and value:
726 + parsed = ast.literal_eval(value)
727 + cuda_runtime = str(parsed or "")
728 + break
729 + return torch_version, cuda_runtime
730 +
731 + def _gpu_identity(self, gpu_uuid: str) -> tuple[str, str, str]:
732 + output = self._run(
733 + [
734 + self.nvidia_smi,
735 + "--query-gpu=name,uuid,driver_version",
736 + "--format=csv,noheader,nounits",
737 + f"--id={gpu_uuid}",
738 + ]
739 + ).stdout.strip()
740 + name, uuid, driver = (item.strip() for item in output.split(",", 2))
741 + return name, uuid, driver
742 +
743 + def _artifact_bytes(self, workdir: Path) -> int:
744 + return sum(path.stat().st_size for path in workdir.rglob("*") if path.is_file())
745 +
746 + def _environment_identity(self) -> tuple[tuple[str, str], ...]:
747 + prefixes = ("VLLM_", "CUDA_", "NCCL_", "TORCH_", "TRITON_")
748 + environment = os.environ.copy()
749 + environment.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
750 + selected = tuple(
751 + sorted(
752 + (key, hashlib.sha256(key.encode() + b"\0" + value.encode()).hexdigest())
753 + for key, value in environment.items()
754 + if key.startswith(prefixes)
755 + and key not in {"VLLM_API_KEY", "VLLM_SNAPSHOT_TIMEOUT_S"}
756 + )
757 + )
758 + return selected + self._plugin_identity()
759 +
760 + def make_manifest(
761 + self,
762 + args: argparse.Namespace,
763 + engine_argv: tuple[str, ...],
764 + inventory: ProcessInventory,
765 + oracle: Oracle,
766 + workdir: Path,
767 + ) -> SnapshotManifest:
768 + identity = self.current_identity(inventory.gpu_uuid)
769 + redacted_argv = list(engine_argv)
770 + redact_api_keys = False
771 + for index, argument in enumerate(redacted_argv):
772 + option, separator, _ = argument.partition("=")
773 + normalized_option = option.replace("_", "-")
774 + # Parsed argv can contain unambiguous long-option abbreviations.
775 + if len(normalized_option) > 2 and "--api-key".startswith(normalized_option):
776 + redact_api_keys = not separator
777 + if separator:
778 + redacted_argv[index] = f"{option}=***"
779 + elif len(normalized_option) > 2 and "--hf-token".startswith(
780 + normalized_option
781 + ):
782 + if separator:
783 + redacted_argv[index] = f"{option}=***"
784 + elif index + 1 < len(redacted_argv):
785 + redacted_argv[index + 1] = "***"
786 + redact_api_keys = False
787 + elif argument.startswith("-"):
788 + redact_api_keys = False
789 + elif redact_api_keys:
790 + redacted_argv[index] = "***"
791 + engine_argv = tuple(redacted_argv)
792 + revision = str(getattr(args, "revision", None) or "")
793 + tokenizer_revision = str(getattr(args, "tokenizer_revision", None) or revision)
794 + source_model = str(getattr(args, "model_tag", None) or args.model)
795 + from vllm.config.model import get_served_model_name
796 +
797 + served_model_name = get_served_model_name(
798 + source_model, getattr(args, "served_model_name", None)
799 + )
800 + return SnapshotManifest(
801 + schema_version=1,
802 + boundary="post-engine-init-reloadable-state-released",
803 + created_at=datetime.now(timezone.utc).isoformat(),
804 + artifact_bytes=self._artifact_bytes(workdir),
805 + **identity.model_dump(),
806 + model=source_model,
807 + served_model_name=str(served_model_name),
808 + model_revision=revision,
809 + tokenizer_revision=tokenizer_revision,
810 + engine_argv=engine_argv,
811 + process_tree=inventory.process_tree,
812 + cuda_holders=inventory.cuda_holders,
813 + external_cache_files=inventory.external_cache_files,
814 + oracle_token_ids=oracle.token_ids,
815 + oracle_text=oracle.text,
816 + oracle_sampled_token_logprob=oracle.sampled_token_logprob,
817 + )
818 +
819 + def publish(self, workdir: Path, manifest: SnapshotManifest) -> None:
820 + write_manifest_atomic(workdir, manifest)
821 + _fsync_directory(workdir.parent)
822 +
823 + def current_identity(self, gpu_uuid: str) -> SnapshotRuntimeIdentity:
824 + gpu_name, gpu_uuid, driver_version = self._gpu_identity(gpu_uuid)
825 + torch_version, cuda_runtime = self._torch_identity()
826 + host_id = Path("/etc/machine-id").read_text().strip()
827 + return SnapshotRuntimeIdentity(
828 + vllm_version=importlib.metadata.version("vllm"),
829 + python_version=platform.python_version(),
830 + torch_version=torch_version,
831 + cuda_runtime=cuda_runtime,
832 + driver_version=driver_version,
833 + criu_version=self._version([self.criu, "--version"]),
834 + cuda_checkpoint_sha256=self._sha256(Path(self.cuda_checkpoint)),
835 + kernel_release=platform.release(),
836 + host_id=host_id,
837 + gpu_name=gpu_name,
838 + gpu_uuid=gpu_uuid,
839 + environment=self._environment_identity(),
840 + )
841 +
842 + def _read_restored_pid(self, pidfile: Path) -> int:
843 + try:
844 + payload = self._run([*self._privileged(), "cat", str(pidfile)]).stdout
845 + except Exception as error:
846 + raise SnapshotRestoreError("restored PID file is missing") from error
847 + match = re.fullmatch(r"([1-9][0-9]{0,9})\n?", payload)
848 + if match is None:
849 + raise SnapshotRestoreError("restored PID file is invalid")
850 + restored_pid = int(match.group(1))
851 + if restored_pid > 2**31 - 1:
852 + raise SnapshotRestoreError("restored PID file is invalid")
853 + return restored_pid
854 +
855 + def _process_state(self, pid: int) -> tuple[int, int, int, int]:
856 + payload = (Path("/proc") / str(pid) / "stat").read_text()
857 + command_end = payload.rfind(")")
858 + fields = payload[command_end + 2 :].split()
859 + if command_end < 0 or len(fields) < 20:
860 + raise ValueError(f"invalid process state for PID {pid}")
861 + return (int(fields[1]), int(fields[2]), int(fields[3]), int(fields[19]))
862 +
863 + def _process_states(self) -> dict[int, tuple[int, int, int, int]]:
864 + states: dict[int, tuple[int, int, int, int]] = {}
865 + for entry in Path("/proc").iterdir():
866 + if not entry.name.isdigit():
867 + continue
868 + pid = int(entry.name)
869 + try:
870 + states[pid] = self._process_state(pid)
871 + except (OSError, ValueError):
872 + continue
873 + return states
874 +
875 + def _process_command(self, pid: int) -> tuple[str, ...]:
876 + payload = (Path("/proc") / str(pid) / "cmdline").read_bytes()
877 + return tuple(
878 + os.fsdecode(argument) for argument in payload.split(b"\0") if argument
879 + )
880 +
881 + def _pidfd_open(self, pid: int) -> int:
882 + pidfd_open = getattr(os, "pidfd_open", None)
883 + if not callable(pidfd_open):
884 + raise SnapshotRestoreError("snapshot requires Linux pidfd support")
885 + return pidfd_open(pid)
886 +
887 + def _pidfd_exited(self, pidfd: int) -> bool:
888 + poller = select.poll()
889 + poller.register(pidfd, select.POLLIN | select.POLLHUP | select.POLLERR)
890 + return any(self._pidfd_event_is_exit(event) for _fd, event in poller.poll(0))
891 +
892 + def _pidfd_event_is_exit(self, event: int) -> bool:
893 + if event & (select.POLLERR | select.POLLNVAL):
894 + raise SnapshotRestoreError("restored process pidfd poll failed")
895 + return bool(event & (select.POLLIN | select.POLLHUP))
896 +
897 + def _pin_restored_tree(
898 + self,
899 + artifact: Path,
900 + process_tree: tuple[int, ...],
901 + cuda_holders: tuple[int, ...],
902 + ) -> None:
903 + root_pid = process_tree[0]
904 + expected_pids = set(process_tree)
905 + expected_cuda_holders = set(cuda_holders)
906 + if not expected_cuda_holders or not expected_cuda_holders.issubset(
907 + expected_pids
908 + ):
909 + raise SnapshotRestoreError(
910 + "captured CUDA holders are outside the captured process tree"
911 + )
912 + if root_pid in self._restored_processes:
913 + raise SnapshotRestoreError("restored PID already has a pinned transaction")
914 + handles: list[tuple[int, int]] = []
915 + try:
916 + for pid in process_tree:
917 + handles.append((pid, self._pidfd_open(pid)))
918 + handles_by_pid = dict(handles)
919 + if any(self._pidfd_exited(pidfd) for _pid, pidfd in handles):
920 + raise SnapshotRestoreError(
921 + "restored process exited during verification"
922 + )
923 +
924 + root_state = self._process_state(root_pid)
925 + if root_state[1:3] != (root_pid, root_pid):
926 + raise SnapshotRestoreError(
927 + "restored root is not the session and process-group leader"
928 + )
929 + command = self._process_command(root_pid)
930 + release_file = str(artifact / "release.json")
931 + if not any(
932 + command[index : index + 2] == ("-m", "vllm.snapshot.server")
933 + for index in range(len(command) - 1)
934 + ) or not any(
935 + command[index : index + 2] == ("--release-file", release_file)
936 + for index in range(len(command) - 1)
937 + ):
938 + raise SnapshotRestoreError(
939 + "restored root command does not match the snapshot barrier"
940 + )
941 + if self._process_state(root_pid) != root_state:
942 + raise SnapshotRestoreError("restored PID changed during verification")
943 + if self._pidfd_exited(handles[0][1]):
944 + raise SnapshotRestoreError("restored root exited during verification")
945 +
946 + first_states = self._process_states()
947 + live_session = {
948 + pid for pid, state in first_states.items() if state[2] == root_pid
949 + }
950 + if not expected_cuda_holders.issubset(live_session):
951 + raise SnapshotRestoreError(
952 + "restored session is missing captured CUDA holders"
953 + )
954 + if live_session != expected_pids:
955 + raise SnapshotRestoreError(
956 + "restored session does not match the captured process tree"
957 + )
958 + if first_states[root_pid] != root_state:
959 + raise SnapshotRestoreError("restored PID changed during verification")
960 + for pid in process_tree[1:]:
961 + state = first_states[pid]
962 + pidfd = handles_by_pid[pid]
963 + if self._process_state(pid) != state or self._pidfd_exited(pidfd):
964 + raise SnapshotRestoreError(
965 + "restored process tree changed during verification"
966 + )
967 +
968 + second_states = self._process_states()
969 + second_session = {
970 + pid for pid, state in second_states.items() if state[2] == root_pid
971 + }
972 + if second_session != expected_pids or any(
973 + second_states[pid] != first_states[pid] for pid in process_tree
974 + ):
975 + raise SnapshotRestoreError(
976 + "restored process tree changed during verification"
977 + )
978 + if any(self._pidfd_exited(pidfd) for _pid, pidfd in handles):
979 + raise SnapshotRestoreError(
980 + "restored process exited during verification"
981 + )
982 + except BaseException as error:
983 + _close_pidfds(handles)
984 + if isinstance(error, SnapshotRestoreError):
985 + raise
986 + if not isinstance(error, Exception):
987 + raise
988 + raise SnapshotRestoreError(
989 + "restored process tree could not be pinned"
990 + ) from error
991 + self._restored_processes[root_pid] = tuple(handles)
992 +
993 + def _abort_failed_restore(
994 + self,
995 + artifact: Path,
996 + root_pid: int,
997 + staged_remaps: dict[str, tuple[int, int]],
998 + primary: BaseException,
999 + ) -> NoReturn:
1000 + secondary: list[str] = []
1001 + try:
1002 + _write_json_atomic(
1003 + artifact / "release.json", {"release": False}, overwrite=True
1004 + )
1005 + except BaseException as error:
1006 + secondary.append(f"abort marker failed: {_error_detail(error)}")
1007 + # The marker above is only a request. Terminate and wait for the tree we
1008 + # own, before releasing the shared state it may still be using. A tree
1009 + # that could not be pinned is not ours to signal.
1010 + if root_pid in self._restored_processes:
1011 + try:
1012 + self.cleanup(root_pid)
1013 + except BaseException as error:
1014 + secondary.append(f"restore cleanup failed: {_error_detail(error)}")
1015 + try:
1016 + self._rollback_link_remaps(staged_remaps)
1017 + except BaseException as error:
1018 + secondary.append(f"remap rollback failed: {_error_detail(error)}")
1019 + if secondary:
1020 + raise SnapshotRestoreError(
1021 + f"snapshot restore failed: {_error_detail(primary)}; "
1022 + + "; ".join(secondary)
1023 + ) from primary
1024 + raise primary
1025 +
1026 + def _verify_external_cache_files(self, manifest: SnapshotManifest) -> None:
1027 + """Reject a rotated generated cache before CRIU reopens it by path.
1028 +
1029 + Like the PID probe above this reports a stale host early rather than
1030 + guarding it: CRIU stays authoritative for a cache replaced afterwards.
1031 + """
1032 + for path, fingerprint in manifest.external_cache_files:
1033 + target = Path(path)
1034 + if not target.is_file():
1035 + raise SnapshotRestoreError(
1036 + f"external cache file vanished since capture: {path}; "
1037 + f"{_CACHE_REMEDY}"
1038 + )
1039 + if fingerprint in (_CACHE_LOCKED, _CACHE_UNREADABLE):
1040 + continue
1041 + try:
1042 + current = self._sha256(target)
1043 + except OSError as error:
1044 + raise SnapshotRestoreError(
1045 + f"external cache file is unreadable since capture: {path}; "
1046 + f"{_CACHE_REMEDY}"
1047 + ) from error
1048 + if current != fingerprint:
1049 + raise SnapshotRestoreError(
1050 + f"external cache file changed since capture: {path}; "
1051 + f"{_CACHE_REMEDY}"
1052 + )
1053 +
1054 + def restore(self, artifact: Path, manifest: SnapshotManifest) -> int:
1055 + # This is necessarily a best-effort TOCTOU check. CRIU restore and the
1056 + # rollback below remain authoritative if a PID is claimed after it.
1057 + for pid in manifest.process_tree:
1058 + try:
1059 + os.kill(pid, 0)
1060 + except ProcessLookupError:
1061 + continue
1062 + except PermissionError:
1063 + pass
1064 + except Exception as error:
1065 + raise SnapshotRestoreError(
1066 + f"captured PID availability probe failed: {pid}"
1067 + ) from error
1068 + raise SnapshotRestoreError(f"captured PID is already occupied: {pid}")
1069 + self._verify_external_cache_files(manifest)
1070 + expected_root_pid = manifest.process_tree[0]
1071 + (artifact / "release.json").unlink(missing_ok=True)
1072 + pidfile = artifact / "restored.pid"
1073 + pidfile.unlink(missing_ok=True)
1074 + self._reset_child_log(artifact)
1075 + staged_remaps = self._restore_link_remaps(artifact)
1076 + try:
1077 + self._criu(
1078 + "restore",
1079 + artifact,
1080 + [
1081 + "--images-dir",
1082 + str(artifact / "images"),
1083 + "--log-file",
1084 + "restore.log",
1085 + *_COMMON_CRIU_OPTIONS,
1086 + "--restore-detached",
1087 + "--pidfile",
1088 + str(pidfile),
1089 + ],
1090 + )
1091 + # Pin before reading the PID file so that a failure between restore
1092 + # and the read cannot orphan the tree.
1093 + self._pin_restored_tree(
1094 + artifact, manifest.process_tree, manifest.cuda_holders
1095 + )
1096 + restored_pid = self._read_restored_pid(pidfile)
1097 + if restored_pid != expected_root_pid:
1098 + raise SnapshotRestoreError(
1099 + "restored PID does not match the captured process tree"
1100 + )
1101 + except BaseException as error:
1102 + self._abort_failed_restore(
1103 + artifact, expected_root_pid, staged_remaps, error
1104 + )
1105 + return restored_pid
1106 +
1107 + def release(self, artifact: Path, host: str | None, port: int) -> None:
1108 + if not 1 <= port <= 65535:
1109 + raise SnapshotRestoreError(f"snapshot restore port is invalid: {port}")
1110 + family = socket.AF_INET6 if host and ":" in host else socket.AF_INET
1111 + with socket.socket(family=family, type=socket.SOCK_STREAM) as listener_probe:
1112 + listener_probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
1113 + try:
1114 + listener_probe.bind((host or "", port))
1115 + except OSError as error:
1116 + address = f"{host or '0.0.0.0'}:{port}"
1117 + raise SnapshotRestoreError(
1118 + f"snapshot restore address is already in use: {address}"
1119 + ) from error
1120 +
1121 + _write_json_atomic(
1122 + artifact / "release.json",
1123 + {"release": True, "host": host, "port": port},
1124 + overwrite=True,
1125 + )
1126 +
1127 + def _connect_host(self, host: str | None) -> str:
1128 + if not host or host == "0.0.0.0":
1129 + return "127.0.0.1"
1130 + if host == "::":
1131 + return "::1"
1132 + return host
1133 +
1134 + def wait_listener(self, root_pid: int, host: str | None, port: int) -> None:
1135 + deadline = time.monotonic() + self.timeout_s
1136 + connect_host = self._connect_host(host)
1137 + while time.monotonic() < deadline:
1138 + try:
1139 + with socket.create_connection((connect_host, port), timeout=1):
1140 + return
1141 + except OSError as error:
1142 + handles = self._restored_processes.get(root_pid)
1143 + if handles is None or self._pidfd_exited(handles[0][1]):
1144 + raise SnapshotRestoreError(
1145 + "restored process exited before binding HTTP"
1146 + ) from error
1147 + time.sleep(0.1)
1148 + raise SnapshotRestoreError("restored HTTP listener did not become ready")
1149 +
1150 + def request_oracle(
1151 + self, host: str | None, port: int, manifest: SnapshotManifest
1152 + ) -> Oracle:
1153 + connect_host = self._connect_host(host)
1154 + url_host = f"[{connect_host}]" if ":" in connect_host else connect_host
1155 + payload = json.dumps(
1156 + {
1157 + "model": manifest.served_model_name,
1158 + "prompt": "The capital of France is",
1159 + "min_tokens": 1,
1160 + "max_tokens": 1,
1161 + "temperature": 0,
1162 + "seed": 0,
1163 + "logprobs": 0,
1164 + "return_token_ids": True,
1165 + }
1166 + ).encode()
1167 + request = urllib.request.Request(
1168 + f"http://{url_host}:{port}/v1/completions",
1169 + data=payload,
1170 + headers={"Content-Type": "application/json"},
1171 + )
1172 + with urllib.request.urlopen(request, timeout=120) as response:
1173 + body = json.load(response)
1174 + choice = body["choices"][0]
1175 + try:
1176 + return Oracle(
1177 + token_ids=tuple(choice["token_ids"]),
1178 + text=choice["text"],
1179 + sampled_token_logprob=choice["logprobs"]["token_logprobs"][0],
1180 + )
1181 + except ValueError as error:
1182 + raise SnapshotRestoreError(
1183 + "snapshot HTTP canary must return exactly one finite token logprob"
1184 + ) from error
1185 + except (IndexError, KeyError, TypeError) as error:
1186 + raise SnapshotRestoreError(
1187 + "snapshot HTTP canary response is missing sampled token logprob"
1188 + ) from error
1189 +
1190 + def cleanup(self, root_pid: int) -> None:
1191 + handles = self._restored_processes.get(root_pid)
1192 + if handles is None:
1193 + raise SnapshotRestoreError("no pinned restore transaction for cleanup")
1194 + pidfd_send_signal = getattr(signal, "pidfd_send_signal", None)
1195 + if not callable(pidfd_send_signal):
1196 + raise SnapshotRestoreError("restored process cleanup is incomplete")
1197 + try:
1198 + for _pid, pidfd in handles:
1199 + with suppress(ProcessLookupError):
1200 + pidfd_send_signal(pidfd, signal.SIGTERM, None, 0)
1201 + time.sleep(min(0.5, self.timeout_s))
1202 + for _pid, pidfd in handles:
1203 + with suppress(ProcessLookupError):
1204 + pidfd_send_signal(pidfd, signal.SIGKILL, None, 0)
1205 + except BaseException as error:
1206 + raise SnapshotRestoreError(
1207 + f"restored process cleanup is incomplete: {_error_detail(error)}"
1208 + ) from error
1209 +
1210 + poller = select.poll()
1211 + for _pid, pidfd in handles:
1212 + poller.register(pidfd, select.POLLIN | select.POLLHUP | select.POLLERR)
1213 + pending = {pidfd for _pid, pidfd in handles}
1214 + deadline = time.monotonic() + self.timeout_s
1215 + while pending:
1216 + remaining = deadline - time.monotonic()
1217 + if remaining <= 0:
1218 + raise SnapshotRestoreError("restored process cleanup is incomplete")
1219 + for pidfd, event in poller.poll(max(1, int(remaining * 1000))):
1220 + if not self._pidfd_event_is_exit(event):
1221 + continue
1222 + pending.discard(pidfd)
1223 + with suppress(KeyError, OSError):
1224 + poller.unregister(pidfd)
1225 + self._restored_processes.pop(root_pid)
1226 + _close_pidfds(handles)
1227 +
1228 + def complete_restore(self, root_pid: int) -> None:
1229 + handles = self._restored_processes.pop(root_pid, None)
1230 + if handles is None:
1231 + raise SnapshotRestoreError("no pinned restore transaction to complete")
1232 + _close_pidfds(handles)
1233 +
1234 + def abort_create(self, root_pid: int) -> None:
1235 + process = self._children.get(root_pid)
1236 + if process is None:
1237 + return
1238 + with suppress(ProcessLookupError):
1239 + os.killpg(root_pid, signal.SIGKILL)
1240 + try:
1241 + process.wait(timeout=10)
1242 + except subprocess.TimeoutExpired as error:
1243 + raise SnapshotCreateError(
1244 + f"snapshot child process group {root_pid} survived SIGKILL"
1245 + ) from error
1246 + self._children.pop(root_pid, None)
vllm/snapshot/server.py new
+252
@@ -0,0 +1,252 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import argparse
5 +import asyncio
6 +import json
7 +import os
8 +import stat
9 +import sys
10 +from dataclasses import dataclass
11 +from pathlib import Path
12 +from typing import Any
13 +
14 +from pydantic import ValidationError
15 +
16 +from vllm.snapshot.manifest import ReleaseMarker, _validation_path, _write_json_atomic
17 +from vllm.snapshot.types import Oracle, oracles_match
18 +
19 +_CANARY_PROMPT = "The capital of France is"
20 +
21 +
22 +class SnapshotCanaryError(RuntimeError):
23 + """The initialized engine did not produce a valid snapshot oracle."""
24 +
25 +
26 +class SnapshotBarrierError(RuntimeError):
27 + """The controller supplied an invalid snapshot release marker."""
28 +
29 +
30 +@dataclass(frozen=True)
31 +class ListenerConfig:
32 + host: str | None
33 + port: int
34 +
35 +
36 +@dataclass(frozen=True)
37 +class ControlArgs:
38 + ready_file: Path
39 + release_file: Path
40 + release_timeout_s: float
41 +
42 +
43 +def write_ready_atomic(path: Path, oracle: Oracle) -> None:
44 + _write_json_atomic(
45 + path,
46 + {
47 + "sampled_token_logprob": oracle.sampled_token_logprob,
48 + "token_ids": oracle.token_ids,
49 + "text": oracle.text,
50 + },
51 + )
52 +
53 +
54 +def read_release_marker(path: Path) -> ListenerConfig:
55 + path = Path(path)
56 + path_stat = path.lstat()
57 + if stat.S_ISLNK(path_stat.st_mode):
58 + raise SnapshotBarrierError("release marker must not be a symlink")
59 + try:
60 + payload = json.loads(path.read_text())
61 + except (OSError, json.JSONDecodeError) as error:
62 + raise SnapshotBarrierError("release marker is not valid JSON") from error
63 + try:
64 + marker = ReleaseMarker.model_validate(payload)
65 + except ValidationError as error:
66 + raise SnapshotBarrierError(
67 + f"release marker is invalid: {_validation_path(error)}"
68 + ) from error
69 + return ListenerConfig(host=marker.host, port=marker.port)
70 +
71 +
72 +def parse_control_args(argv: list[str]) -> tuple[ControlArgs, list[str]]:
73 + parser = argparse.ArgumentParser(add_help=False)
74 + parser.add_argument("--ready-file", type=Path, required=True)
75 + parser.add_argument("--release-file", type=Path, required=True)
76 + parser.add_argument("--release-timeout-s", type=float, default=900.0)
77 + control_args, remaining = parser.parse_known_args(argv)
78 + if remaining and remaining[0] == "--":
79 + remaining = remaining[1:]
80 + if control_args.release_timeout_s <= 0:
81 + raise ValueError("release timeout must be positive")
82 + return (
83 + ControlArgs(
84 + ready_file=control_args.ready_file,
85 + release_file=control_args.release_file,
86 + release_timeout_s=control_args.release_timeout_s,
87 + ),
88 + remaining,
89 + )
90 +
91 +
92 +async def _release_reloadable_state(engine: Any) -> None:
93 + """Discard model and KV state before the process image is captured."""
94 + await engine.sleep(level=2)
95 +
96 +
97 +async def _restore_reloadable_state(engine: Any) -> None:
98 + """Rebuild state discarded by ``_release_reloadable_state``."""
99 + await engine.wake_up(tags=["weights"])
100 + await engine.collective_rpc("reload_weights")
101 + await engine.wake_up(tags=["kv_cache"])
102 +
103 +
104 +def oracle_from_request_output(request_output: Any) -> Oracle:
105 + try:
106 + candidate = request_output.outputs[0]
107 + token_ids = tuple(candidate.token_ids)
108 + (sampled_token_id,) = token_ids
109 + return Oracle(
110 + token_ids=token_ids,
111 + text=candidate.text,
112 + sampled_token_logprob=candidate.logprobs[0][sampled_token_id].logprob,
113 + )
114 + except ValueError as error:
115 + raise SnapshotCanaryError(
116 + "snapshot canary must produce exactly one finite token logprob"
117 + ) from error
118 + except (AttributeError, IndexError, KeyError, TypeError) as error:
119 + raise SnapshotCanaryError(
120 + "snapshot canary did not return sampled token logprob"
121 + ) from error
122 +
123 +
124 +async def run_engine_canary(engine: Any) -> Oracle:
125 + from vllm import SamplingParams
126 +
127 + final_output = None
128 + sampling_params = SamplingParams(
129 + temperature=0,
130 + min_tokens=1,
131 + max_tokens=1,
132 + seed=0,
133 + logprobs=0,
134 + )
135 + async for output in engine.generate(
136 + _CANARY_PROMPT,
137 + sampling_params,
138 + request_id="vllm-snapshot-canary",
139 + ):
140 + final_output = output
141 + if final_output is None:
142 + raise SnapshotCanaryError("canary generation returned no output")
143 + return oracle_from_request_output(final_output)
144 +
145 +
146 +def detach_snapshot_streams() -> None:
147 + """Remove launch-log descriptors from the reusable process image."""
148 + sys.stdout.flush()
149 + sys.stderr.flush()
150 + sink = os.open(os.devnull, os.O_WRONLY)
151 + try:
152 + os.dup2(sink, sys.stdout.fileno())
153 + os.dup2(sink, sys.stderr.fileno())
154 + finally:
155 + os.close(sink)
156 +
157 +
158 +async def wait_for_release_marker(
159 + path: Path,
160 + *,
161 + timeout_s: float,
162 + poll_interval_s: float = 0.05,
163 +) -> ListenerConfig:
164 + if timeout_s <= 0:
165 + raise ValueError("release timeout must be positive")
166 + if poll_interval_s <= 0:
167 + raise ValueError("release poll interval must be positive")
168 + remaining_s = timeout_s
169 + while not path.exists():
170 + if remaining_s <= 0:
171 + raise TimeoutError(f"release marker not found before timeout: {path}")
172 + delay_s = min(poll_interval_s, remaining_s)
173 + await asyncio.sleep(delay_s)
174 + remaining_s -= delay_s
175 + return read_release_marker(path)
176 +
177 +
178 +def parse_vllm_args(argv: list[str]) -> Any:
179 + from vllm.entrypoints.launchers.cli_args import (
180 + make_arg_parser,
181 + validate_parsed_serve_args,
182 + )
183 + from vllm.utils.argparse_utils import FlexibleArgumentParser
184 +
185 + parser = FlexibleArgumentParser(prog="python -m vllm.snapshot.server")
186 + args = make_arg_parser(parser).parse_args(argv)
187 + if args.model_tag is not None:
188 + args.model = args.model_tag
189 + if args.grpc or args.headless:
190 + raise ValueError("snapshot server requires the HTTP frontend")
191 + if args.api_server_count not in (None, 1):
192 + raise ValueError("snapshot server supports one API server")
193 + args.api_server_count = None
194 + validate_parsed_serve_args(args)
195 + return args
196 +
197 +
198 +async def run_vllm_snapshot_child(control: ControlArgs, args: Any) -> None:
199 + from vllm.entrypoints.launchers.api_server.entry import (
200 + build_and_serve,
201 + build_async_engine_client,
202 + )
203 + from vllm.entrypoints.launchers.launcher import setup_server
204 +
205 + args.enable_sleep_mode = True
206 +
207 + async with build_async_engine_client(args) as engine:
208 + oracle = await run_engine_canary(engine)
209 + try:
210 + await _release_reloadable_state(engine)
211 + await _restore_reloadable_state(engine)
212 + rehearsal_oracle = await run_engine_canary(engine)
213 + if not oracles_match(oracle, rehearsal_oracle):
214 + raise SnapshotCanaryError("snapshot rehearsal changed canary output")
215 + await _release_reloadable_state(engine)
216 + except SnapshotCanaryError:
217 + raise
218 + except Exception as error:
219 + raise SnapshotCanaryError("snapshot rehearsal failed") from error
220 +
221 + detach_snapshot_streams()
222 + write_ready_atomic(control.ready_file, oracle)
223 + listener = await wait_for_release_marker(
224 + control.release_file,
225 + timeout_s=control.release_timeout_s,
226 + )
227 + args.host = listener.host
228 + args.port = listener.port
229 + await _restore_reloadable_state(engine)
230 + listen_address, sock = setup_server(args, reuse_port=False)
231 + try:
232 + shutdown_task = await build_and_serve(
233 + engine,
234 + listen_address,
235 + sock,
236 + args,
237 + )
238 + await shutdown_task
239 + finally:
240 + sock.close()
241 +
242 +
243 +def main(argv: list[str] | None = None) -> None:
244 + control, remaining = parse_control_args(sys.argv[1:] if argv is None else argv)
245 + args = parse_vllm_args(remaining)
246 + import uvloop
247 +
248 + uvloop.run(run_vllm_snapshot_child(control, args))
249 +
250 +
251 +if __name__ == "__main__":
252 + main()
vllm/snapshot/types.py new
+37
@@ -0,0 +1,37 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import math
5 +from dataclasses import dataclass
6 +
7 +
8 +@dataclass(frozen=True)
9 +class Oracle:
10 + token_ids: tuple[int, ...]
11 + text: str
12 + sampled_token_logprob: float
13 +
14 + def __post_init__(self) -> None:
15 + if len(self.token_ids) != 1:
16 + raise ValueError("snapshot oracle must contain exactly one token")
17 + value = self.sampled_token_logprob
18 + if (
19 + isinstance(value, bool)
20 + or not isinstance(value, (int, float))
21 + or not math.isfinite(value)
22 + ):
23 + raise ValueError("sampled token logprob must be finite")
24 + object.__setattr__(self, "sampled_token_logprob", float(value))
25 +
26 +
27 +def oracles_match(expected: Oracle, actual: Oracle) -> bool:
28 + return (
29 + expected.token_ids == actual.token_ids
30 + and expected.text == actual.text
31 + and math.isclose(
32 + expected.sampled_token_logprob,
33 + actual.sampled_token_logprob,
34 + rel_tol=0.0,
35 + abs_tol=1e-3,
36 + )
37 + )
vllm/utils/argparse_utils.py
+14 -7
@@ -519,7 +519,7 @@ class FlexibleArgumentParser(ArgumentParser):
519 config_args = self.load_config_file(file_path)
520
521 # 0th index might be the sub command {serve,chat,complete,...}
522 - # optionally followed by model_tag (only for serve)
522 + # optionally followed by model_tag (serve or snapshot create)
523 # followed by config args
524 # followed by rest of cli args.
525 # maintaining this order will enforce the precedence
@@ -527,8 +527,11 @@ class FlexibleArgumentParser(ArgumentParser):
527 if args[0].startswith("-"):
528 # No sub command (e.g., api_server entry point)
529 args = config_args + args[0:index] + args[index + 2 :]
530 - elif args[0] == "serve":
531 - model_in_cli = len(args) > 1 and not args[1].startswith("-")
530 + elif args[0] == "serve" or args[:2] == ["snapshot", "create"]:
531 + model_index = 1 if args[0] == "serve" else 2
532 + model_in_cli = len(args) > model_index and not args[model_index].startswith(
533 + "-"
534 + )
535 model_in_config = any(arg == "--model" for arg in config_args)
536
537 if not model_in_cli and not model_in_config:
@@ -540,15 +543,19 @@ class FlexibleArgumentParser(ArgumentParser):
543 if model_in_cli:
544 # Model specified as positional arg, keep CLI version
545 args = (
543 - [args[0]]
544 - + [args[1]]
546 + args[: model_index + 1]
547 + config_args
546 - + args[2:index]
548 + + args[model_index + 1 : index]
549 + args[index + 2 :]
550 )
551 else:
552 # No model in CLI, use config if available
551 - args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
553 + args = (
554 + args[:model_index]
555 + + config_args
556 + + args[model_index:index]
557 + + args[index + 2 :]
558 + )
559 else:
560 args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
561