[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>
This commit is contained in:
co-authored by
Taneem Ibrahim
elehayym
Andreas Karatzas
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent
8bf051cd51
commit
13058faa49
+311
@@ -0,0 +1,311 @@
|
||||
#!/usr/bin/env bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
set -Eeuo pipefail
|
||||
readonly MODEL="Qwen/Qwen3-0.6B"
|
||||
readonly MODEL_REVISION="c1899de289a04d12100db370d81485cdf75e47ca"
|
||||
readonly EXPECTED_TOKEN_ID="12095"
|
||||
readonly EXPECTED_TEXT=" Paris"
|
||||
readonly CLEANUP_RESERVE_S=300
|
||||
RUN_LIMIT_S="${SNAPSHOT_E2E_RUN_LIMIT_S:-3300}"
|
||||
[[ "$RUN_LIMIT_S" =~ ^[1-9][0-9]{2,3}$ ]] \
|
||||
|| { echo "SNAPSHOT_E2E_RUN_LIMIT_S must be a bounded decimal integer" >&2; exit 1; }
|
||||
(( RUN_LIMIT_S >= CLEANUP_RESERVE_S + 60 && RUN_LIMIT_S <= 3300 )) \
|
||||
|| { echo "SNAPSHOT_E2E_RUN_LIMIT_S must be between 360 and 3300" >&2; exit 1; }
|
||||
readonly RUN_LIMIT_S
|
||||
readonly HOST_HF_CACHE="${SNAPSHOT_E2E_HOST_HF_CACHE:-}"
|
||||
HF_CACHE_MOUNT=()
|
||||
if [[ -n "$HOST_HF_CACHE" ]]; then
|
||||
[[ "$HOST_HF_CACHE" == /* && -d "$HOST_HF_CACHE" ]] \
|
||||
|| { echo "SNAPSHOT_E2E_HOST_HF_CACHE must be an existing absolute directory" >&2; exit 1; }
|
||||
HF_CACHE_MOUNT=(--mount "type=bind,source=$HOST_HF_CACHE,target=/e2e/hf,readonly")
|
||||
PREPARE_ARTIFACT_ROOT='install -d -m 0700 /e2e'
|
||||
else
|
||||
PREPARE_ARTIFACT_ROOT='mkdir -m 0700 /e2e && mkdir -m 0700 /e2e/hf'
|
||||
fi
|
||||
readonly PREPARE_ARTIFACT_ROOT
|
||||
if (( $# != 2 )); then
|
||||
echo "usage: $0 IMAGE_TAG BUILDKITE_COMMIT" >&2
|
||||
exit 2
|
||||
fi
|
||||
readonly IMAGE_TAG="$1"
|
||||
readonly EXPECTED_COMMIT="$2"
|
||||
readonly DEADLINE_EPOCH="$(( $(date +%s) + RUN_LIMIT_S ))"
|
||||
RUN_ID="" CONTAINER_NAME="" CONTAINER_ID=""
|
||||
GPU_UUID="" PORT_ONE="" PORT_TWO=""
|
||||
LINK_REMAP_BASELINE=""
|
||||
SNAPSHOT_PIDS=()
|
||||
die() {
|
||||
echo "initialized snapshot E2E: $*" >&2
|
||||
exit 1
|
||||
}
|
||||
|
||||
run() {
|
||||
local label="$1" maximum_s="$2" available limit
|
||||
shift 2
|
||||
available="$((DEADLINE_EPOCH - $(date +%s) - CLEANUP_RESERVE_S))"
|
||||
(( available > 0 )) || die "no time remains before the cleanup reserve"
|
||||
(( available < maximum_s )) && limit="$available" || limit="$maximum_s"
|
||||
echo "--- ${label} (timeout ${limit}s)" >&2
|
||||
timeout --signal=TERM --kill-after=30s "${limit}s" "$@"
|
||||
}
|
||||
|
||||
link_remaps() {
|
||||
find /dev/shm -maxdepth 1 -name 'link_remap.*' \
|
||||
-printf '%p|%D|%i|%m|%s\n' | sort
|
||||
}
|
||||
|
||||
gpu_pids() {
|
||||
nvidia-smi --id="$GPU_UUID" --query-compute-apps=pid \
|
||||
--format=csv,noheader,nounits | sed '/^[[:space:]]*$/d'
|
||||
}
|
||||
|
||||
port_listeners() {
|
||||
ss -H -ltn "sport = :$1"
|
||||
}
|
||||
|
||||
state_clean() {
|
||||
local pid port output
|
||||
for pid in "${SNAPSHOT_PIDS[@]}"; do
|
||||
[[ ! -e "/proc/$pid" ]] || return 1
|
||||
done
|
||||
output="$(gpu_pids)" || return 1
|
||||
[[ -z "$output" ]] || return 1
|
||||
output="$(link_remaps)" || return 1
|
||||
[[ "$output" == "$LINK_REMAP_BASELINE" ]] || return 1
|
||||
for port in "$PORT_ONE" "$PORT_TWO"; do
|
||||
[[ -n "$port" ]] || continue
|
||||
output="$(port_listeners "$port")" || return 1
|
||||
[[ -z "$output" ]] || return 1
|
||||
done
|
||||
}
|
||||
|
||||
wait_clean() {
|
||||
local _attempt
|
||||
for ((_attempt = 0; _attempt < 60; _attempt++)); do
|
||||
state_clean && return 0
|
||||
sleep 1
|
||||
done
|
||||
return 1
|
||||
}
|
||||
|
||||
container_names() {
|
||||
timeout 15s docker container ls --all --format '{{.Names}}'
|
||||
}
|
||||
|
||||
cleanup() {
|
||||
local failed=0 names="" actual="" expected=""
|
||||
if [[ -n "$CONTAINER_NAME" ]]; then
|
||||
names="$(container_names 2>/dev/null)" || return 1
|
||||
if grep -Fxq "$CONTAINER_NAME" <<< "$names"; then
|
||||
actual="$(docker container inspect \
|
||||
--format '{{.Id}}|{{index .Config.Labels "ai.vllm.snapshot.e2e.run"}}' \
|
||||
"$CONTAINER_NAME" 2>/dev/null)" || return 1
|
||||
expected="${CONTAINER_ID:-${actual%%|*}}|$RUN_ID"
|
||||
if [[ "$actual" != "$expected" ]]; then
|
||||
echo "refusing to remove a container with mismatched identity" >&2
|
||||
failed=1
|
||||
else
|
||||
timeout 60s docker rm --force "$CONTAINER_NAME" >/dev/null \
|
||||
|| failed=1
|
||||
fi
|
||||
fi
|
||||
names="$(container_names 2>/dev/null)" || return 1
|
||||
grep -Fxq "$CONTAINER_NAME" <<< "$names" && failed=1
|
||||
fi
|
||||
if [[ -n "$PORT_TWO" ]] && ! wait_clean; then
|
||||
echo "process, GPU, port, or link-remap residue remained" >&2
|
||||
failed=1
|
||||
fi
|
||||
return "$failed"
|
||||
}
|
||||
|
||||
on_exit() {
|
||||
local status=$?
|
||||
trap - EXIT INT TERM
|
||||
set +e
|
||||
cleanup || status=1
|
||||
exit "$status"
|
||||
}
|
||||
|
||||
trap on_exit EXIT
|
||||
trap 'exit 130' INT
|
||||
trap 'exit 143' TERM
|
||||
|
||||
parse_inspect() {
|
||||
timeout 30s python3 -c '
|
||||
import json, math, sys
|
||||
data = json.load(sys.stdin)
|
||||
model, revision, gpu_uuid, token_id, text = sys.argv[1:]
|
||||
expected = {
|
||||
"boundary": "post-engine-init-reloadable-state-released",
|
||||
"model": model, "model_revision": revision,
|
||||
"gpu_uuid": gpu_uuid, "oracle_token_ids": [int(token_id)],
|
||||
"oracle_text": text,
|
||||
}
|
||||
bad = next((field for field, value in expected.items() if data.get(field) != value), None)
|
||||
if bad:
|
||||
raise SystemExit(f"inspect field {bad} mismatch: {data.get(bad)!r}")
|
||||
logprob = data.get("oracle_sampled_token_logprob")
|
||||
if type(logprob) is not float or not math.isfinite(logprob):
|
||||
raise SystemExit("inspect logprob is not finite")
|
||||
pids = data.get("process_tree")
|
||||
if not pids:
|
||||
raise SystemExit("inspect process_tree is empty")
|
||||
for pid in pids:
|
||||
print(f"pid={pid}")
|
||||
print(f"oracle_logprob={logprob}")
|
||||
' "$MODEL" "$MODEL_REVISION" "$GPU_UUID" "$EXPECTED_TOKEN_ID" \
|
||||
"$EXPECTED_TEXT"
|
||||
}
|
||||
|
||||
assert_oracle() {
|
||||
local port="$1" ordinal="$2" response
|
||||
response="$(run "restore ${ordinal} public completion" 120 curl \
|
||||
--fail --silent --show-error --max-time 120 \
|
||||
--header 'Content-Type: application/json' \
|
||||
--data "{\"model\":\"${MODEL}\",\"prompt\":\"The capital of France is\",\"temperature\":0,\"seed\":0,\"min_tokens\":1,\"max_tokens\":1,\"logprobs\":0,\"return_token_ids\":true}" \
|
||||
"http://127.0.0.1:${port}/v1/completions")"
|
||||
printf '%s' "$response" | timeout 30s python3 -c '
|
||||
import json, math, sys
|
||||
choice = json.load(sys.stdin)["choices"][0]
|
||||
token_id, text, expected, ordinal = sys.argv[1:]
|
||||
values = choice["logprobs"]["token_logprobs"]
|
||||
actual = float(values[0]) if len(values) == 1 else math.nan
|
||||
if (choice["token_ids"] != [int(token_id)] or choice["text"] != text
|
||||
or not math.isfinite(actual) or abs(actual - float(expected)) > 1e-3):
|
||||
raise SystemExit(f"public oracle mismatch: {choice!r}")
|
||||
print(f"restore={ordinal} token_id={token_id} text={text!r} logprob={actual}")
|
||||
' "$EXPECTED_TOKEN_ID" "$EXPECTED_TEXT" "$INSPECT_LOGPROB" "$ordinal"
|
||||
}
|
||||
|
||||
[[ "$(uname -s)" == "Linux" && "$(uname -m)" == "x86_64" ]] \
|
||||
|| die "host must be Linux x86_64"
|
||||
[[ "$EXPECTED_COMMIT" =~ ^[0-9a-f]{40}$ ]] \
|
||||
|| die "BUILDKITE_COMMIT must be a full lowercase commit hash"
|
||||
[[ "$(< /proc/sys/kernel/io_uring_disabled)" == "2" ]] \
|
||||
|| die "kernel.io_uring_disabled must already equal 2"
|
||||
run "Docker preflight" 60 docker info >/dev/null
|
||||
|
||||
VISIBLE_GPU="${NVIDIA_VISIBLE_DEVICES:-}"
|
||||
if [[ -z "$VISIBLE_GPU" || "$VISIBLE_GPU" == "all" || "$VISIBLE_GPU" == "void" ]]; then
|
||||
VISIBLE_GPU="${CUDA_VISIBLE_DEVICES:-}"
|
||||
fi
|
||||
if [[ -z "$VISIBLE_GPU" || "$VISIBLE_GPU" == "all" || "$VISIBLE_GPU" == "void" ]]; then
|
||||
CDI_GPU="${BUILDKITE_PLUGIN_DOCKER_DEVICES_0:-}"
|
||||
if [[ "$CDI_GPU" == nvidia.com/gpu=* ]]; then
|
||||
VISIBLE_GPU="${CDI_GPU#nvidia.com/gpu=}"
|
||||
else
|
||||
VISIBLE_GPU=""
|
||||
fi
|
||||
fi
|
||||
[[ -n "$VISIBLE_GPU" && "$VISIBLE_GPU" != all && "$VISIBLE_GPU" != void \
|
||||
&& "$VISIBLE_GPU" != *","* ]] \
|
||||
|| die "the runner must expose exactly one selected GPU"
|
||||
GPU_ROW="$(nvidia-smi --id="$VISIBLE_GPU" \
|
||||
--query-gpu=name,uuid,mig.mode.current,memory.total \
|
||||
--format=csv,noheader,nounits)"
|
||||
IFS=',' read -r GPU_NAME GPU_UUID MIG_MODE GPU_MEMORY_MIB <<< "$GPU_ROW"
|
||||
for variable in GPU_NAME GPU_UUID MIG_MODE GPU_MEMORY_MIB; do
|
||||
printf -v "$variable" '%s' "${!variable#"${!variable%%[![:space:]]*}"}"
|
||||
done
|
||||
[[ "$GPU_NAME" == *H200* && "$GPU_UUID" == GPU-* && "$MIG_MODE" == Disabled ]] \
|
||||
|| die "selected device is not a full non-MIG H200: $GPU_ROW"
|
||||
(( GPU_MEMORY_MIB >= 130000 )) || die "selected H200 does not expose full memory"
|
||||
GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU"
|
||||
[[ -z "$GPU_PIDS" ]] || die "selected H200 is not idle"
|
||||
|
||||
RUN_ID="snapshot-e2e-$(date +%s)-$$-$RANDOM"
|
||||
CONTAINER_NAME="vllm-$RUN_ID"
|
||||
read -r PORT_ONE PORT_TWO < <(python3 -c 'import socket; s=[socket.socket(),socket.socket()]; [x.bind(("127.0.0.1",0)) for x in s]; print(*(x.getsockname()[1] for x in s))')
|
||||
LINK_REMAP_BASELINE="$(link_remaps)"
|
||||
|
||||
run "pull exact candidate image" 900 docker pull --platform linux/amd64 \
|
||||
"$IMAGE_TAG" >/dev/null
|
||||
IFS='|' read -r IMAGE_ID IMAGE_PLATFORM OCI_COMMIT VLLM_COMMIT < <(
|
||||
docker image inspect \
|
||||
--format '{{.Id}}|{{.Os}}/{{.Architecture}}|{{index .Config.Labels "org.opencontainers.image.revision"}}|{{index .Config.Labels "ai.vllm.build.commit"}}' \
|
||||
"$IMAGE_TAG"
|
||||
)
|
||||
[[ "$IMAGE_ID" =~ ^sha256:[0-9a-f]{64}$ && "$IMAGE_PLATFORM" == linux/amd64 ]] \
|
||||
|| die "candidate image identity is invalid"
|
||||
[[ "$OCI_COMMIT" == "$EXPECTED_COMMIT" && "$VLLM_COMMIT" == "$EXPECTED_COMMIT" ]] \
|
||||
|| die "candidate image commit labels do not match"
|
||||
CONTAINER_ID="$(run "start exact candidate container" 90 docker run --detach \
|
||||
--name "$CONTAINER_NAME" --label "ai.vllm.snapshot.e2e.run=$RUN_ID" \
|
||||
--gpus "device=$GPU_UUID" --user 0 --privileged --pid=host --ipc=host \
|
||||
--network=host --env "CUDA_VISIBLE_DEVICES=$GPU_UUID" --env HF_HUB_DISABLE_TELEMETRY=1 \
|
||||
--env VLLM_NO_USAGE_STATS=1 --env VLLM_USE_V2_MODEL_RUNNER=1 \
|
||||
${HF_CACHE_MOUNT[@]+"${HF_CACHE_MOUNT[@]}"} \
|
||||
--entrypoint sleep "$IMAGE_ID" infinity)"
|
||||
[[ "$CONTAINER_ID" =~ ^[0-9a-f]{64}$ ]] || die "invalid container ID"
|
||||
run "prepare private artifact root" 60 docker exec "$CONTAINER_NAME" sh -c \
|
||||
"$PREPARE_ARTIFACT_ROOT && df -Pk / | awk 'NR == 2 && \$4 >= 10485760 {ok=1} END {exit !ok}'"
|
||||
run "verify snapshot runtime" 60 docker exec "$CONTAINER_NAME" sh -c \
|
||||
'criu --version; test -x /usr/local/sbin/cuda-checkpoint; test -f /usr/local/lib/criu/cuda_plugin.so'
|
||||
run "verify selected CUDA device" 60 docker exec "$CONTAINER_NAME" python3 -c '
|
||||
import sys, torch
|
||||
if torch.cuda.device_count() != 1:
|
||||
raise SystemExit("expected exactly one visible CUDA device")
|
||||
actual = str(torch.cuda.get_device_properties(0).uuid).removeprefix("GPU-")
|
||||
if actual != sys.argv[1].removeprefix("GPU-"):
|
||||
raise SystemExit(f"CUDA device {actual} does not match {sys.argv[1]}")
|
||||
print(f"selected_cuda_uuid=GPU-{actual}")
|
||||
' "$GPU_UUID"
|
||||
OFFLINE_EXEC=(docker exec --env HF_HOME=/e2e/hf --env HF_HUB_OFFLINE=1 \
|
||||
--env TRANSFORMERS_OFFLINE=1 --env VLLM_SNAPSHOT_TIMEOUT_S=900 \
|
||||
"$CONTAINER_NAME")
|
||||
if [[ -n "$HOST_HF_CACHE" ]]; then
|
||||
run "validate pinned model in read-only host cache" 120 \
|
||||
"${OFFLINE_EXEC[@]}" python3 -c \
|
||||
'from pathlib import Path; from huggingface_hub import snapshot_download; import sys; p=snapshot_download(repo_id=sys.argv[1], revision=sys.argv[2], local_files_only=True); assert Path(p).name == sys.argv[2]' \
|
||||
"$MODEL" "$MODEL_REVISION"
|
||||
else
|
||||
run "prefetch pinned public model" 900 docker exec --env HF_HOME=/e2e/hf \
|
||||
"$CONTAINER_NAME" python3 -c \
|
||||
'from pathlib import Path; from huggingface_hub import snapshot_download; import sys; p=snapshot_download(repo_id=sys.argv[1], revision=sys.argv[2]); assert Path(p).name == sys.argv[2]' \
|
||||
"$MODEL" "$MODEL_REVISION"
|
||||
fi
|
||||
|
||||
run "create compact initialized snapshot" 1200 "${OFFLINE_EXEC[@]}" \
|
||||
vllm snapshot create "$MODEL" --revision "$MODEL_REVISION" \
|
||||
--tokenizer-revision "$MODEL_REVISION" --snapshot-dir /e2e/artifact \
|
||||
--dtype float16 --max-model-len 512 --gpu-memory-utilization 0.50
|
||||
wait_clean || die "create left process, GPU, port, or link-remap residue"
|
||||
|
||||
INSPECT_JSON="$(run "inspect compact initialized snapshot" 120 \
|
||||
"${OFFLINE_EXEC[@]}" vllm snapshot inspect /e2e/artifact)"
|
||||
INSPECT_RECEIPT="$(printf '%s' "$INSPECT_JSON" | parse_inspect)"
|
||||
mapfile -t SNAPSHOT_PIDS < <(sed -n 's/^pid=//p' <<< "$INSPECT_RECEIPT")
|
||||
INSPECT_LOGPROB="$(sed -n 's/^oracle_logprob=//p' <<< "$INSPECT_RECEIPT")"
|
||||
(( ${#SNAPSHOT_PIDS[@]} > 0 )) || die "inspect returned no PIDs"
|
||||
state_clean || die "create left recorded state alive"
|
||||
echo "$INSPECT_RECEIPT"
|
||||
|
||||
run "restore compact snapshot 1" 900 "${OFFLINE_EXEC[@]}" \
|
||||
vllm snapshot restore /e2e/artifact --host 127.0.0.1 --port "$PORT_ONE"
|
||||
PORT_LISTENERS="$(port_listeners "$PORT_ONE")" || die "failed to inspect restore 1 port"
|
||||
[[ -n "$PORT_LISTENERS" ]] || die "restore 1 did not bind its port"
|
||||
GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU"
|
||||
[[ -n "$GPU_PIDS" ]] || die "restore 1 did not engage the selected GPU"
|
||||
assert_oracle "$PORT_ONE" 1
|
||||
run "stop container after restore 1" 60 docker stop --time 20 "$CONTAINER_NAME" >/dev/null
|
||||
[[ "$(docker container inspect --format '{{.Id}}' "$CONTAINER_NAME")" == "$CONTAINER_ID" ]] \
|
||||
|| die "docker stop replaced the container"
|
||||
wait_clean || die "restore 1 teardown left residue"
|
||||
|
||||
run "start the same container" 60 docker start "$CONTAINER_NAME" >/dev/null
|
||||
[[ "$(docker container inspect --format '{{.Id}}' "$CONTAINER_NAME")" == "$CONTAINER_ID" ]] \
|
||||
|| die "docker start replaced the container"
|
||||
run "restore compact snapshot 2" 900 "${OFFLINE_EXEC[@]}" \
|
||||
vllm snapshot restore /e2e/artifact --host 127.0.0.1 --port "$PORT_TWO"
|
||||
PORT_LISTENERS="$(port_listeners "$PORT_TWO")" || die "failed to inspect restore 2 port"
|
||||
[[ -n "$PORT_LISTENERS" ]] || die "restore 2 did not bind its port"
|
||||
GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU"
|
||||
[[ -n "$GPU_PIDS" ]] || die "restore 2 did not engage the selected GPU"
|
||||
assert_oracle "$PORT_TWO" 2
|
||||
run "stop container after restore 2" 60 docker stop --time 20 "$CONTAINER_NAME" >/dev/null
|
||||
wait_clean || die "restore 2 teardown left residue"
|
||||
cleanup || die "exact cleanup failed"
|
||||
echo "container_id=$CONTAINER_ID restore_1_port=$PORT_ONE restore_2_port=$PORT_TWO"
|
||||
echo "result=PASS"
|
||||
@@ -2,6 +2,26 @@ group: Entrypoints
|
||||
depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: ":nvidia: (H200) Initialized Snapshot E2E"
|
||||
key: initialized-snapshot-e2e
|
||||
device: h200
|
||||
num_devices: 1
|
||||
no_plugin: true
|
||||
optional: true
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "."
|
||||
env:
|
||||
SNAPSHOT_E2E_RUN_LIMIT_S: "1200"
|
||||
source_file_dependencies:
|
||||
- .buildkite/scripts/initialized-snapshot-e2e.sh
|
||||
- .buildkite/test_areas/entrypoints.yaml
|
||||
- docker/Dockerfile
|
||||
- tools/install_snapshot_runtime.sh
|
||||
- vllm/entrypoints/cli/
|
||||
- vllm/snapshot/
|
||||
commands:
|
||||
- bash .buildkite/scripts/initialized-snapshot-e2e.sh "$IMAGE_TAG" "$BUILDKITE_COMMIT"
|
||||
|
||||
- label: ":nvidia: (H200 MIG 35GB) Entrypoints Unit"
|
||||
device: h200_35gb
|
||||
key: entrypoints-unit-tests
|
||||
|
||||
@@ -735,6 +735,8 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
FROM ${FINAL_BASE_IMAGE} AS vllm-runtime-base
|
||||
|
||||
ARG INSTALL_RUBIN_PRERELEASE
|
||||
ARG TARGETPLATFORM
|
||||
ARG max_jobs=2
|
||||
ARG CUDA_VERSION
|
||||
ARG PYTHON_VERSION
|
||||
ARG NCCL_VERSION
|
||||
@@ -854,6 +856,20 @@ RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
|
||||
&& ldconfig; \
|
||||
fi
|
||||
|
||||
# Install the initialized-snapshot runtime on the architecture supported by
|
||||
# `vllm snapshot`. CRIU and its CUDA plugin are built from the same pinned
|
||||
# source. The `cuda-checkpoint`-compatible helper is built from NVIDIA Dynamo's
|
||||
# Apache-2.0 source against the public CUDA Driver checkpoint APIs instead of
|
||||
# redistributing NVIDIA's proprietary cuda-checkpoint binary.
|
||||
COPY tools/install_snapshot_runtime.sh /usr/local/bin/install_snapshot_runtime.sh
|
||||
RUN if [ "${TARGETPLATFORM}" = "linux/amd64" ]; then \
|
||||
MAX_JOBS="${max_jobs}" /usr/local/bin/install_snapshot_runtime.sh; \
|
||||
else \
|
||||
echo "Initialized snapshots are not installed for ${TARGETPLATFORM}"; \
|
||||
fi
|
||||
|
||||
ENV CRIU_CUDA_PLUGIN_DIR=/usr/local/lib/criu
|
||||
|
||||
# Install uv for faster pip installs
|
||||
RUN python3 -m pip install uv
|
||||
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
# Initialized engine snapshots
|
||||
|
||||
Initialized engine snapshots are an experimental way to trade local disk space
|
||||
and host privileges for a faster vLLM activation. Snapshot creation initializes
|
||||
the engine and records deterministic generation output before CRIU captures the
|
||||
process tree and CUDA state. Restore validates the saved environment, restores
|
||||
the engine, binds the HTTP server, and reproduces the recorded token and sampled
|
||||
token log probability before returning.
|
||||
|
||||
This path is intended for repeatedly activating the same model and engine
|
||||
configuration on the same machine. It is not a portable model artifact.
|
||||
|
||||
## Requirements
|
||||
|
||||
Snapshots currently require:
|
||||
|
||||
- Linux on x86-64 with one NVIDIA GPU.
|
||||
- TP1 with one unauthenticated plaintext HTTP server. Other parallel sizes,
|
||||
TLS, middleware, Unix sockets, and speculative decoding are unsupported.
|
||||
- Because snapshot mode requires an unauthenticated plaintext HTTP server, use
|
||||
a trusted host or network boundary, or external controls. See the vLLM
|
||||
[Security guide](../usage/security.md).
|
||||
- [CRIU](https://github.com/checkpoint-restore/criu), its CUDA plugin, a
|
||||
`cuda-checkpoint`-compatible helper, and `nvidia-smi` on `PATH`.
|
||||
- Root or passwordless `sudo` for CRIU.
|
||||
- `io_uring` disabled before launch because CRIU cannot dump it. Use
|
||||
`kernel.io_uring_disabled=1` for an unprivileged process. A process running
|
||||
as root, including the `docker exec` flow below, bypasses `=1`, so use `=2`
|
||||
host-wide.
|
||||
- No established TCP connection to a peer outside the captured process tree.
|
||||
Current Hugging Face hub clients hold their connections for the process
|
||||
lifetime, so download the model in a separate step and run create with
|
||||
`HF_HUB_OFFLINE=1`, as the quickstart below does.
|
||||
- A remote model ID and an immutable 40-character `--revision`. Local model
|
||||
directories and mutable revisions are not supported.
|
||||
- Enough disk for the artifact, with the same installed vLLM package, model
|
||||
files, container filesystem, and generated-cache paths available at restore.
|
||||
Generated-cache files are not copied into the artifact. The manifest
|
||||
fingerprints the ones the captured tree holds open, so restore fails early
|
||||
and names the file when one is removed or replaced. Only open descriptors
|
||||
are recorded. A library the tree mapped and then closed is not fingerprinted,
|
||||
and CRIU still reopens it by path, so replacing that file permanently
|
||||
invalidates the artifact without an early error.
|
||||
|
||||
The official `vllm/vllm-openai` Linux x86-64 image includes the snapshot
|
||||
runtime. It still requires a compatible host driver, kernel, and privileges.
|
||||
Arm64 images omit it. Source installs must set `CRIU_CUDA_PLUGIN_DIR` to the
|
||||
directory containing `cuda_plugin.so`.
|
||||
|
||||
Run snapshot commands with `docker exec` inside a long-lived container. Restore
|
||||
hands the API server off as a detached process, so a one-shot container would
|
||||
stop that server when its PID 1 exits. Snapshot preflight requires every
|
||||
component of the artifact path to be owned by the invoking user or root, with
|
||||
the directory itself at mode 0700, and the commands below run as root inside
|
||||
the official image, so the bind-mounted host directory is created with `sudo`
|
||||
and root ownership. The model downloads in its own step so the captured tree
|
||||
holds no hub connection, and create runs offline. This example also keeps the
|
||||
container filesystem and `/dev/shm` namespace stable for the lifetime of the
|
||||
artifact:
|
||||
|
||||
```bash
|
||||
sudo sysctl kernel.io_uring_disabled=2
|
||||
|
||||
snapshot_root="$(pwd)/vllm-snapshots"
|
||||
sudo install -d -m 0700 -o root -g root "${snapshot_root}"
|
||||
|
||||
docker run --detach --name vllm-snapshot \
|
||||
--gpus all \
|
||||
--privileged \
|
||||
--pid=host \
|
||||
--ipc=host \
|
||||
--network=host \
|
||||
--mount "type=bind,source=${snapshot_root},target=/snapshots" \
|
||||
--entrypoint sleep \
|
||||
vllm/vllm-openai:latest infinity
|
||||
|
||||
docker exec vllm-snapshot hf download Qwen/Qwen3-0.6B \
|
||||
--revision c1899de289a04d12100db370d81485cdf75e47ca
|
||||
|
||||
docker exec -e HF_HUB_OFFLINE=1 vllm-snapshot vllm snapshot create Qwen/Qwen3-0.6B \
|
||||
--snapshot-dir /snapshots/qwen3-0.6b \
|
||||
--revision c1899de289a04d12100db370d81485cdf75e47ca \
|
||||
--dtype float16 \
|
||||
--max-model-len 512
|
||||
|
||||
docker exec vllm-snapshot vllm snapshot inspect /snapshots/qwen3-0.6b
|
||||
|
||||
docker exec vllm-snapshot vllm snapshot restore \
|
||||
/snapshots/qwen3-0.6b --host 0.0.0.0 --port 8000
|
||||
```
|
||||
|
||||
Keep that container and its mounts available while the snapshot is in use.
|
||||
Stop and remove it only after the restored API server is no longer needed.
|
||||
|
||||
Create initializes the engine, records a one-token canary, releases and reloads
|
||||
weights and KV cache to rehearse restore, then releases them again for capture.
|
||||
The manifest is published only after CRIU completes and the source tree stops.
|
||||
Creation is offline preparation and is not part of restore latency.
|
||||
|
||||
The private `0700` artifact contains process memory, CUDA state, engine
|
||||
arguments, compatibility identity, and canary output; its manifest is `0600`.
|
||||
Literal API keys and Hugging Face tokens are redacted from the manifest's engine
|
||||
arguments. The manifest records selected environment names plus deterministic
|
||||
name-bound fingerprints, not their values. The protected CRIU artifact can
|
||||
still contain process secrets, so treat it as sensitive data. Restore reloads
|
||||
model files and KV cache before binding HTTP, then reproduces the canary or
|
||||
tears down the restored tree. The inspect command prints that identity and
|
||||
canary without executing the saved process.
|
||||
|
||||
## Restore behavior
|
||||
|
||||
Restore fails before CRIU runs if the saved identity does not match the current
|
||||
host. It does not silently fall back to ordinary startup. After CRIU restores
|
||||
the process tree, vLLM releases the saved engine to bind the requested HTTP
|
||||
address and checks the first generated token and sampled-token log probability
|
||||
against the snapshot canary. The command returns only after that check passes.
|
||||
The restored API server continues to run as a detached process.
|
||||
|
||||
The pre-release port probe is best-effort only. It neither reserves the port
|
||||
nor authenticates the listener that appears afterward.
|
||||
|
||||
Rollback terminates and waits for the restored tree after process identity
|
||||
verification succeeds. If verification fails, vLLM writes an abort marker for
|
||||
the snapshot server instead of signaling unverified PIDs. If that server has
|
||||
already exited, surviving engine processes can require operator cleanup.
|
||||
Verify that the previous tree has stopped before retrying; do not kill a process
|
||||
solely because its PID appears in the manifest, since PIDs can be reused.
|
||||
|
||||
The artifact is reusable after its previous restored tree stops. Only one
|
||||
snapshot or external CRIU operation may use a shared `/dev/shm` mount at a time.
|
||||
|
||||
## Tradeoffs and limitations
|
||||
|
||||
- Creation has its own latency and briefly requires the full engine. Artifact
|
||||
size can approach the captured process and GPU memory.
|
||||
- Artifacts are not guaranteed to survive a power loss. The manifest is
|
||||
fsynced, but publication does not explicitly flush all captured payload files
|
||||
and directories. Recreate the artifact after an unclean host shutdown.
|
||||
- Restore currently requires the same host, GPU, driver, kernel, Python,
|
||||
PyTorch, installed vLLM version, model revision, engine arguments,
|
||||
selected environment variables, and CRIU plugin binaries.
|
||||
- Only dense float16 TP1 has been validated. Other model formats depend on their
|
||||
existing sleep level 2 reload support; NCCL state is not restored.
|
||||
- CRIU support varies by kernel and driver. Preserve package, library, model,
|
||||
and generated-cache paths for the artifact lifetime.
|
||||
- A snapshot can include application secrets or request state present in the
|
||||
process. Build it before traffic and protect it like process memory.
|
||||
- Any feature that opens an external connection must close it before capture.
|
||||
|
||||
For a lower-complexity option that retains a live process, see
|
||||
[Sleep mode](sleep_mode.md). Sleep mode and initialized snapshots retain
|
||||
different amounts of state and have different idle resource costs.
|
||||
File diff suppressed because it is too large
Load Diff
@@ -19,6 +19,7 @@ import torch
|
||||
|
||||
from vllm.model_executor.models.mimo_v2 import _shard_fp8_qkv_proj
|
||||
from vllm.utils.math_utils import cdiv
|
||||
from vllm.utils.torch_utils import set_default_torch_num_threads
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
@@ -39,6 +40,13 @@ GEOMETRIES = {
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", autouse=True)
|
||||
def single_threaded_cpu_ops():
|
||||
# Do not prime the parent's thread pool before later engine forks.
|
||||
with set_default_torch_num_threads(1):
|
||||
yield
|
||||
|
||||
|
||||
def _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp):
|
||||
return (
|
||||
num_heads // ckpt_tp * head_dim
|
||||
|
||||
Executable
+109
@@ -0,0 +1,109 @@
|
||||
#!/usr/bin/env bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
set -euxo pipefail
|
||||
|
||||
# Surface the isolated snapshot-runtime install cost in the CI build log. The
|
||||
# trap captures the status first: the date substitution below would reset it.
|
||||
_snapshot_install_started=$(date +%s)
|
||||
code=0
|
||||
trap 'code=$?; echo "[TIMING] snapshot runtime install: $(( $(date +%s) - _snapshot_install_started )) s (exit $code)"' EXIT
|
||||
|
||||
readonly CRIU_VERSION="4.2.1"
|
||||
readonly CRIU_SOURCE_SHA256="feffdf4638125ebb12d2434754f80a1d7bbba85a3e6bee98c216f88fb99a5d96"
|
||||
readonly DYNAMO_CHECKPOINT_HELPER_COMMIT="fdf25efad60f696a73393caf341c527a93d84190"
|
||||
readonly DYNAMO_CHECKPOINT_HELPER_SHA256="7b8ea21baecb1da729202c32e0fecbbbafb789d65e5c1c0649ea253ad09b8f82"
|
||||
readonly DYNAMO_LICENSE_SHA256="aeeca3d74a13b91d17ab9211c407b8d94a77bf8d9e5750c467818d3066b8adcf"
|
||||
|
||||
# Keep the compiler and NL headers required by runtime JIT and RDMA.
|
||||
snapshot_shared_build_deps=(
|
||||
build-essential
|
||||
libnl-3-dev
|
||||
libnl-route-3-dev
|
||||
)
|
||||
snapshot_build_deps=(
|
||||
pkg-config
|
||||
libbsd-dev
|
||||
libcap-dev
|
||||
libnet1-dev
|
||||
libprotobuf-dev
|
||||
libprotobuf-c-dev
|
||||
protobuf-c-compiler
|
||||
protobuf-compiler
|
||||
python3-protobuf
|
||||
libgnutls28-dev
|
||||
libnftables-dev
|
||||
uuid-dev
|
||||
)
|
||||
snapshot_runtime_deps=(
|
||||
libbsd0
|
||||
libcap2
|
||||
libnet1
|
||||
libnl-3-200
|
||||
libnl-route-3-200
|
||||
libprotobuf-c1
|
||||
libgnutls30t64
|
||||
libnftables1
|
||||
iproute2
|
||||
iptables
|
||||
procps
|
||||
uuid-runtime
|
||||
util-linux
|
||||
)
|
||||
|
||||
apt-get update -y
|
||||
apt-get install -y --no-install-recommends \
|
||||
"${snapshot_shared_build_deps[@]}" \
|
||||
"${snapshot_build_deps[@]}" \
|
||||
"${snapshot_runtime_deps[@]}"
|
||||
|
||||
snapshot_build_dir="$(mktemp -d)"
|
||||
curl -fsSL \
|
||||
"https://github.com/checkpoint-restore/criu/archive/refs/tags/v${CRIU_VERSION}.tar.gz" \
|
||||
-o "${snapshot_build_dir}/criu.tar.gz"
|
||||
echo "${CRIU_SOURCE_SHA256} ${snapshot_build_dir}/criu.tar.gz" | sha256sum -c -
|
||||
mkdir "${snapshot_build_dir}/criu"
|
||||
tar -xzf "${snapshot_build_dir}/criu.tar.gz" \
|
||||
--strip-components=1 -C "${snapshot_build_dir}/criu"
|
||||
make -C "${snapshot_build_dir}/criu" -j"${MAX_JOBS:-$(nproc)}" criu cuda_plugin
|
||||
install -D -m 0755 "${snapshot_build_dir}/criu/criu/criu" \
|
||||
/usr/local/sbin/criu
|
||||
install -D -m 0755 \
|
||||
"${snapshot_build_dir}/criu/plugins/cuda/cuda_plugin.so" \
|
||||
/usr/local/lib/criu/cuda_plugin.so
|
||||
|
||||
helper_url="https://raw.githubusercontent.com/ai-dynamo/dynamo/${DYNAMO_CHECKPOINT_HELPER_COMMIT}/deploy/snapshot/cmd/cuda-checkpoint-helper/main.c"
|
||||
curl -fsSL "${helper_url}" -o "${snapshot_build_dir}/cuda-checkpoint.c"
|
||||
echo "${DYNAMO_CHECKPOINT_HELPER_SHA256} ${snapshot_build_dir}/cuda-checkpoint.c" \
|
||||
| sha256sum -c -
|
||||
gcc -O2 -Wall -Wextra \
|
||||
-I/usr/local/cuda/include \
|
||||
-L/usr/local/cuda/lib64/stubs \
|
||||
"${snapshot_build_dir}/cuda-checkpoint.c" -lcuda \
|
||||
-o /usr/local/sbin/cuda-checkpoint
|
||||
|
||||
license_url="https://raw.githubusercontent.com/ai-dynamo/dynamo/${DYNAMO_CHECKPOINT_HELPER_COMMIT}/LICENSE"
|
||||
curl -fsSL "${license_url}" -o "${snapshot_build_dir}/DYNAMO-LICENSE"
|
||||
echo "${DYNAMO_LICENSE_SHA256} ${snapshot_build_dir}/DYNAMO-LICENSE" \
|
||||
| sha256sum -c -
|
||||
|
||||
install -D -m 0644 "${snapshot_build_dir}/criu.tar.gz" \
|
||||
/usr/share/doc/criu/source/criu.tar.gz
|
||||
install -D -m 0644 "${snapshot_build_dir}/criu/COPYING" \
|
||||
/usr/share/doc/criu/COPYING
|
||||
install -D -m 0644 "${snapshot_build_dir}/cuda-checkpoint.c" \
|
||||
/usr/share/doc/cuda-checkpoint-helper/main.c
|
||||
install -D -m 0644 "${snapshot_build_dir}/DYNAMO-LICENSE" \
|
||||
/usr/share/doc/cuda-checkpoint-helper/LICENSE
|
||||
|
||||
criu --version
|
||||
ln -s /usr/local/cuda/lib64/stubs/libcuda.so \
|
||||
"${snapshot_build_dir}/libcuda.so.1"
|
||||
LD_LIBRARY_PATH="${snapshot_build_dir}" \
|
||||
cuda-checkpoint --help | grep -q -- "--action"
|
||||
test -x /usr/local/lib/criu/cuda_plugin.so
|
||||
|
||||
apt-get purge -y --auto-remove "${snapshot_build_deps[@]}"
|
||||
rm -rf "${snapshot_build_dir}"
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
@@ -46,6 +46,7 @@ def main():
|
||||
import vllm.entrypoints.cli.openai
|
||||
import vllm.entrypoints.cli.run_batch
|
||||
import vllm.entrypoints.cli.serve
|
||||
import vllm.entrypoints.cli.snapshot
|
||||
from vllm.entrypoints.serve.utils.api_utils import (
|
||||
VLLM_SUBCMD_PARSER_EPILOG,
|
||||
cli_env_setup,
|
||||
@@ -59,8 +60,10 @@ def main():
|
||||
vllm.entrypoints.cli.benchmark.main,
|
||||
vllm.entrypoints.cli.collect_env,
|
||||
vllm.entrypoints.cli.run_batch,
|
||||
vllm.entrypoints.cli.snapshot,
|
||||
]
|
||||
|
||||
if sys.argv[1:2] != ["snapshot"]:
|
||||
cli_env_setup()
|
||||
|
||||
vllm.entrypoints.cli.benchmark.main.maybe_exec_rust_bench()
|
||||
@@ -94,6 +97,11 @@ def main():
|
||||
subparsers = parser.add_subparsers(required=False, dest="subparser")
|
||||
cmds = {}
|
||||
for cmd_module in CMD_MODULES:
|
||||
if cmd_module is vllm.entrypoints.cli.snapshot:
|
||||
new_cmds = cmd_module.cmd_init(
|
||||
create_requested=sys.argv[1:3] == ["snapshot", "create"]
|
||||
)
|
||||
else:
|
||||
new_cmds = cmd_module.cmd_init()
|
||||
for cmd in new_cmds:
|
||||
cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import typing
|
||||
from pathlib import Path
|
||||
|
||||
from vllm.entrypoints.cli.types import CLISubcommand
|
||||
|
||||
if typing.TYPE_CHECKING:
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
else:
|
||||
FlexibleArgumentParser = argparse.ArgumentParser
|
||||
|
||||
|
||||
def validate_create_args(args: argparse.Namespace) -> None:
|
||||
model = str(getattr(args, "model_tag", None) or args.model)
|
||||
revision = str(getattr(args, "revision", None) or "")
|
||||
if (
|
||||
Path(model).exists()
|
||||
or len(revision) != 40
|
||||
or any(character not in "0123456789abcdefABCDEF" for character in revision)
|
||||
):
|
||||
raise ValueError(
|
||||
"snapshot create requires an immutable remote model and "
|
||||
"40-character --revision"
|
||||
)
|
||||
parallel_sizes = (
|
||||
("tensor parallel", args.tensor_parallel_size),
|
||||
("pipeline parallel", args.pipeline_parallel_size),
|
||||
("data parallel", args.data_parallel_size),
|
||||
("prefill context parallel", args.prefill_context_parallel_size),
|
||||
)
|
||||
for name, size in parallel_sizes:
|
||||
if size != 1:
|
||||
raise ValueError(f"snapshot create currently requires {name} size 1")
|
||||
unsupported_frontend = (
|
||||
args.api_server_count not in (None, 1)
|
||||
or args.grpc
|
||||
or args.headless
|
||||
or args.api_key
|
||||
or os.environ.get("VLLM_API_KEY")
|
||||
or args.uds
|
||||
or args.middleware
|
||||
)
|
||||
if unsupported_frontend or any(
|
||||
(
|
||||
args.ssl_keyfile,
|
||||
args.ssl_certfile,
|
||||
args.ssl_ca_certs,
|
||||
args.enable_ssl_refresh,
|
||||
args.ssl_ciphers,
|
||||
)
|
||||
):
|
||||
raise ValueError(
|
||||
"snapshot create requires one unauthenticated plaintext HTTP server"
|
||||
)
|
||||
if args.logprobs_mode in {"raw_logits", "processed_logits"}:
|
||||
raise ValueError("snapshot canary requires a log-probability mode")
|
||||
if any(
|
||||
getattr(args, name, None) is not None
|
||||
for name in ("speculative_config", "spec_method", "spec_model", "spec_tokens")
|
||||
):
|
||||
raise ValueError("snapshot create does not support speculative decoding")
|
||||
if platform.system() != "Linux" or platform.machine() != "x86_64":
|
||||
raise ValueError("snapshot create requires Linux x86_64")
|
||||
|
||||
|
||||
def run_create(args: argparse.Namespace) -> None:
|
||||
validate_create_args(args)
|
||||
from vllm.snapshot.controller import create_snapshot
|
||||
|
||||
create_snapshot(args)
|
||||
|
||||
|
||||
def run_inspect(args: argparse.Namespace) -> None:
|
||||
from vllm.snapshot.manifest import inspect_snapshot
|
||||
|
||||
print(json.dumps(inspect_snapshot(Path(args.snapshot_dir)), indent=2))
|
||||
|
||||
|
||||
def run_restore(args: argparse.Namespace) -> None:
|
||||
from vllm.snapshot.controller import restore_snapshot
|
||||
|
||||
restore_snapshot(args)
|
||||
|
||||
|
||||
class SnapshotSubcommand(CLISubcommand):
|
||||
"""The `snapshot` subcommand for the vLLM CLI."""
|
||||
|
||||
name = "snapshot"
|
||||
|
||||
def __init__(self, *, create_requested: bool = False) -> None:
|
||||
self.create_requested = create_requested
|
||||
|
||||
@staticmethod
|
||||
def cmd(args: argparse.Namespace) -> None:
|
||||
args.snapshot_dispatch(args)
|
||||
|
||||
def subparser_init(
|
||||
self, subparsers: argparse._SubParsersAction
|
||||
) -> FlexibleArgumentParser:
|
||||
parser = subparsers.add_parser(
|
||||
self.name,
|
||||
help="Create, inspect, or restore an initialized vLLM snapshot.",
|
||||
usage="vllm snapshot <create|inspect|restore> [options]",
|
||||
)
|
||||
actions = parser.add_subparsers(required=True, dest="snapshot_action")
|
||||
|
||||
create_parser = actions.add_parser(
|
||||
"create", help="Create a snapshot from an initialized TP1 engine."
|
||||
)
|
||||
if self.create_requested:
|
||||
from vllm.entrypoints.launchers.cli_args import make_arg_parser
|
||||
|
||||
create_parser = make_arg_parser(create_parser)
|
||||
else:
|
||||
create_parser.add_argument("model_tag", nargs="?")
|
||||
create_parser.add_argument("--snapshot-dir", required=True)
|
||||
create_parser.set_defaults(snapshot_dispatch=run_create)
|
||||
|
||||
inspect_parser = actions.add_parser(
|
||||
"inspect", help="Inspect a snapshot without restoring it."
|
||||
)
|
||||
inspect_parser.add_argument("snapshot_dir")
|
||||
inspect_parser.set_defaults(snapshot_dispatch=run_inspect)
|
||||
|
||||
restore_parser = actions.add_parser(
|
||||
"restore", help="Restore a same-host TP1 snapshot."
|
||||
)
|
||||
restore_parser.add_argument("snapshot_dir")
|
||||
restore_parser.add_argument("--host", default=None)
|
||||
restore_parser.add_argument("--port", type=int, default=8000)
|
||||
restore_parser.set_defaults(snapshot_dispatch=run_restore)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def cmd_init(*, create_requested: bool = False) -> list[CLISubcommand]:
|
||||
return [SnapshotSubcommand(create_requested=create_requested)]
|
||||
@@ -0,0 +1,4 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Create and restore explicitly compatible vLLM engine snapshots."""
|
||||
@@ -0,0 +1,132 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Create and restore transaction orchestration for local snapshots."""
|
||||
|
||||
import argparse
|
||||
import shutil
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from vllm.snapshot.manifest import (
|
||||
read_manifest,
|
||||
validate_artifact_root,
|
||||
validate_identity,
|
||||
)
|
||||
from vllm.snapshot.runtime import (
|
||||
LocalSnapshotTools,
|
||||
SnapshotCreateError,
|
||||
SnapshotRestoreError,
|
||||
_error_detail,
|
||||
)
|
||||
from vllm.snapshot.types import Oracle, oracles_match
|
||||
|
||||
|
||||
def _child_engine_argv(
|
||||
args: argparse.Namespace, engine_argv: tuple[str, ...] | None
|
||||
) -> tuple[str, ...]:
|
||||
"""Build the child argv from the first argv source that is present.
|
||||
|
||||
In order: an explicit ``engine_argv`` from the caller, the process argv
|
||||
tail when invoked as ``vllm snapshot create ...``, then the bare model from
|
||||
the parsed args. Whichever source wins, snapshot options are stripped and
|
||||
sleep mode is appended.
|
||||
"""
|
||||
if not engine_argv:
|
||||
process_argv = tuple(sys.argv[1:])
|
||||
if process_argv[:2] == ("snapshot", "create"):
|
||||
engine_argv = process_argv[2:]
|
||||
else:
|
||||
model = getattr(args, "model_tag", None) or getattr(args, "model", None)
|
||||
if not model:
|
||||
raise SnapshotCreateError("snapshot create requires a model argument")
|
||||
engine_argv = (str(model),)
|
||||
child_argv: list[str] = []
|
||||
iterator = iter(engine_argv)
|
||||
for item in iterator:
|
||||
option, separator, _ = item.partition("=")
|
||||
option = option.replace("_", "-")
|
||||
if len(option) > 2 and "--snapshot-dir".startswith(option):
|
||||
if not separator:
|
||||
next(iterator, None)
|
||||
else:
|
||||
child_argv.append(item)
|
||||
if "--enable-sleep-mode" not in child_argv:
|
||||
child_argv.append("--enable-sleep-mode")
|
||||
return tuple(child_argv)
|
||||
|
||||
|
||||
def create_snapshot(
|
||||
args: argparse.Namespace,
|
||||
*,
|
||||
engine_argv: tuple[str, ...] | None = None,
|
||||
tools: LocalSnapshotTools | None = None,
|
||||
) -> None:
|
||||
target = Path(args.snapshot_dir).absolute()
|
||||
toolset = tools or LocalSnapshotTools()
|
||||
toolset.preflight("create", target)
|
||||
if target.exists() or target.is_symlink():
|
||||
raise SnapshotCreateError(f"snapshot target already exists: {target}")
|
||||
validate_artifact_root(target, creating=True)
|
||||
target.parent.mkdir(mode=0o700, parents=False, exist_ok=True)
|
||||
validate_artifact_root(target.parent, creating=False)
|
||||
target.mkdir(mode=0o700)
|
||||
validate_artifact_root(target, creating=False)
|
||||
published = False
|
||||
root_pid: int | None = None
|
||||
try:
|
||||
child_argv = _child_engine_argv(args, engine_argv)
|
||||
root_pid = toolset.launch_child(target, child_argv)
|
||||
oracle = toolset.wait_ready(target, root_pid)
|
||||
inventory = toolset.inventory(root_pid, target)
|
||||
toolset.dump(target, inventory)
|
||||
toolset.verify_dead(inventory)
|
||||
manifest = toolset.make_manifest(args, child_argv, inventory, oracle, target)
|
||||
toolset.publish(target, manifest)
|
||||
published = True
|
||||
except BaseException:
|
||||
if root_pid is not None:
|
||||
toolset.abort_create(root_pid)
|
||||
raise
|
||||
finally:
|
||||
# Runs even when abort_create raised, so no partial artifact is left.
|
||||
if not published and target.exists():
|
||||
shutil.rmtree(target)
|
||||
|
||||
|
||||
def restore_snapshot(
|
||||
args: argparse.Namespace, *, tools: LocalSnapshotTools | None = None
|
||||
) -> None:
|
||||
artifact = Path(args.snapshot_dir).absolute()
|
||||
toolset = tools or LocalSnapshotTools()
|
||||
toolset.preflight("restore", artifact)
|
||||
manifest = read_manifest(artifact)
|
||||
validate_identity(manifest, toolset.current_identity(manifest.gpu_uuid))
|
||||
|
||||
root_pid: int | None = None
|
||||
try:
|
||||
root_pid = toolset.restore(artifact, manifest)
|
||||
toolset.release(artifact, args.host, args.port)
|
||||
toolset.wait_listener(root_pid, args.host, args.port)
|
||||
actual = toolset.request_oracle(args.host, args.port, manifest)
|
||||
expected = Oracle(
|
||||
token_ids=manifest.oracle_token_ids,
|
||||
text=manifest.oracle_text,
|
||||
sampled_token_logprob=manifest.oracle_sampled_token_logprob,
|
||||
)
|
||||
if not oracles_match(expected, actual):
|
||||
raise SnapshotRestoreError(
|
||||
f"snapshot oracle mismatch: expected={expected!r}, actual={actual!r}"
|
||||
)
|
||||
toolset.complete_restore(root_pid)
|
||||
except BaseException as error:
|
||||
if root_pid is not None:
|
||||
try:
|
||||
toolset.cleanup(root_pid)
|
||||
except BaseException as cleanup_error:
|
||||
raise SnapshotRestoreError(
|
||||
f"snapshot restore failed: {_error_detail(error)}; "
|
||||
f"cleanup failed: {_error_detail(cleanup_error)}"
|
||||
) from error
|
||||
if isinstance(error, SnapshotRestoreError):
|
||||
raise
|
||||
raise SnapshotRestoreError(str(error)) from error
|
||||
@@ -0,0 +1,277 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import stat
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Any, Literal
|
||||
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
ConfigDict,
|
||||
Field,
|
||||
StrictFloat,
|
||||
StrictInt,
|
||||
StrictStr,
|
||||
ValidationError,
|
||||
field_validator,
|
||||
model_validator,
|
||||
)
|
||||
from typing_extensions import Self
|
||||
|
||||
|
||||
class SnapshotCompatibilityError(RuntimeError):
|
||||
"""The snapshot identity does not match the requested runtime."""
|
||||
|
||||
|
||||
class SnapshotSecurityError(RuntimeError):
|
||||
"""The snapshot artifact does not meet the private-path contract."""
|
||||
|
||||
|
||||
_MAX_PID = 2**31 - 1
|
||||
_PositivePid = Annotated[int, Field(gt=0, le=_MAX_PID)]
|
||||
_OracleTokenId = Annotated[int, Field(ge=0)]
|
||||
|
||||
|
||||
class SnapshotRuntimeIdentity(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
|
||||
|
||||
vllm_version: str
|
||||
python_version: str
|
||||
torch_version: str
|
||||
cuda_runtime: str
|
||||
driver_version: str
|
||||
criu_version: str
|
||||
cuda_checkpoint_sha256: str
|
||||
kernel_release: str
|
||||
host_id: str
|
||||
gpu_name: str
|
||||
gpu_uuid: str
|
||||
environment: tuple[tuple[str, str], ...]
|
||||
|
||||
|
||||
class SnapshotManifest(SnapshotRuntimeIdentity):
|
||||
schema_version: Literal[1]
|
||||
boundary: Literal["post-engine-init-reloadable-state-released"]
|
||||
created_at: str
|
||||
artifact_bytes: Annotated[int, Field(ge=0)]
|
||||
model: str
|
||||
served_model_name: str
|
||||
model_revision: str
|
||||
tokenizer_revision: str
|
||||
engine_argv: tuple[str, ...]
|
||||
process_tree: Annotated[tuple[_PositivePid, ...], Field(min_length=1)]
|
||||
cuda_holders: Annotated[tuple[_PositivePid, ...], Field(min_length=1)]
|
||||
external_cache_files: tuple[tuple[str, str], ...] = ()
|
||||
oracle_token_ids: Annotated[
|
||||
tuple[_OracleTokenId, ...], Field(min_length=1, max_length=1)
|
||||
]
|
||||
oracle_text: str
|
||||
oracle_sampled_token_logprob: float
|
||||
|
||||
@field_validator("schema_version", mode="before")
|
||||
@classmethod
|
||||
def _require_schema_version_one(cls, value: Any) -> Any:
|
||||
if type(value) is not int or value != 1:
|
||||
raise ValueError("schema version must be integer 1")
|
||||
return value
|
||||
|
||||
@field_validator("oracle_sampled_token_logprob", mode="before")
|
||||
@classmethod
|
||||
def _require_finite_float(cls, value: Any) -> Any:
|
||||
if type(value) is not float or not math.isfinite(value):
|
||||
raise ValueError("sampled token log probability must be a finite float")
|
||||
return value
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_process_relationships(self) -> Self:
|
||||
if len(set(self.process_tree)) != len(self.process_tree):
|
||||
raise ValueError("process_tree contains duplicate PIDs")
|
||||
if len(set(self.cuda_holders)) != len(self.cuda_holders):
|
||||
raise ValueError("cuda_holders contains duplicate PIDs")
|
||||
if not set(self.cuda_holders).issubset(self.process_tree):
|
||||
raise ValueError("cuda_holders must be a subset of process_tree")
|
||||
return self
|
||||
|
||||
|
||||
class ReleaseMarker(BaseModel):
|
||||
"""Release marker the controller writes; unknown keys are ignored."""
|
||||
|
||||
release: Literal[True]
|
||||
host: StrictStr | None = None
|
||||
port: Annotated[StrictInt, Field(ge=1, le=65535)]
|
||||
|
||||
@field_validator("release", mode="before")
|
||||
@classmethod
|
||||
def _require_json_true(cls, value: Any) -> Any:
|
||||
# A literal accepts 1 as True; the barrier requires JSON true itself.
|
||||
if value is not True:
|
||||
raise ValueError("release must be JSON true")
|
||||
return value
|
||||
|
||||
|
||||
class ReadyMarker(BaseModel):
|
||||
"""Ready marker the child writes; unknown keys are ignored."""
|
||||
|
||||
token_ids: tuple[StrictInt, ...]
|
||||
text: StrictStr
|
||||
sampled_token_logprob: StrictFloat | StrictInt
|
||||
|
||||
|
||||
def validate_identity(
|
||||
expected: SnapshotRuntimeIdentity, actual: SnapshotRuntimeIdentity
|
||||
) -> None:
|
||||
"""Require an exact match for every field that can affect compatibility."""
|
||||
# ``expected`` may be a manifest, so compare the identity fields only.
|
||||
fields = set(SnapshotRuntimeIdentity.model_fields)
|
||||
expected_values = expected.model_dump(include=fields)
|
||||
actual_values = actual.model_dump(include=fields)
|
||||
differing = [
|
||||
name
|
||||
for name in SnapshotRuntimeIdentity.model_fields
|
||||
if expected_values[name] != actual_values[name]
|
||||
]
|
||||
if differing:
|
||||
raise SnapshotCompatibilityError(f"snapshot mismatch: {', '.join(differing)}")
|
||||
|
||||
|
||||
def _existing_path_chain(path: Path):
|
||||
current = path
|
||||
while True:
|
||||
if current.exists() or current.is_symlink():
|
||||
yield current
|
||||
if current.parent == current:
|
||||
break
|
||||
current = current.parent
|
||||
|
||||
|
||||
def _is_trusted_sticky_directory(path_stat: os.stat_result) -> bool:
|
||||
return (
|
||||
path_stat.st_uid == 0
|
||||
and stat.S_ISDIR(path_stat.st_mode)
|
||||
and bool(path_stat.st_mode & stat.S_ISVTX)
|
||||
)
|
||||
|
||||
|
||||
def validate_artifact_root(path: Path, *, creating: bool) -> None:
|
||||
"""Validate that an artifact path cannot be replaced by another user."""
|
||||
path = Path(os.path.abspath(path))
|
||||
if path.is_symlink():
|
||||
raise SnapshotSecurityError(f"snapshot path is a symlink: {path}")
|
||||
if not creating and not path.exists():
|
||||
raise SnapshotSecurityError(f"snapshot path does not exist: {path}")
|
||||
|
||||
effective_uid = os.geteuid()
|
||||
for component in _existing_path_chain(path):
|
||||
component_stat = component.lstat()
|
||||
if stat.S_ISLNK(component_stat.st_mode):
|
||||
raise SnapshotSecurityError(
|
||||
f"snapshot path contains a symlink: {component}"
|
||||
)
|
||||
if component_stat.st_uid not in (0, effective_uid):
|
||||
raise SnapshotSecurityError(
|
||||
f"snapshot path has an untrusted owner: {component}"
|
||||
)
|
||||
if component_stat.st_mode & (
|
||||
stat.S_IWGRP | stat.S_IWOTH
|
||||
) and not _is_trusted_sticky_directory(component_stat):
|
||||
raise SnapshotSecurityError(
|
||||
f"snapshot path has a group- or world-writable ancestor: {component}"
|
||||
)
|
||||
|
||||
if path.exists():
|
||||
path_stat = path.lstat()
|
||||
if not stat.S_ISDIR(path_stat.st_mode):
|
||||
raise SnapshotSecurityError(f"snapshot path is not a directory: {path}")
|
||||
if path_stat.st_uid != effective_uid:
|
||||
raise SnapshotSecurityError(
|
||||
f"snapshot directory is not owned by the current user: {path}"
|
||||
)
|
||||
if stat.S_IMODE(path_stat.st_mode) & 0o077:
|
||||
raise SnapshotSecurityError(
|
||||
f"snapshot directory must have mode 0700: {path}"
|
||||
)
|
||||
|
||||
|
||||
def _validation_path(error: ValidationError) -> str:
|
||||
detail = error.errors(include_url=False)[0]
|
||||
location = detail["loc"]
|
||||
if not location:
|
||||
return "JSON" if detail["type"] == "json_invalid" else "root"
|
||||
return ".".join(map(str, location))
|
||||
|
||||
|
||||
def _fsync_directory(path: Path) -> None:
|
||||
directory_fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
|
||||
try:
|
||||
os.fsync(directory_fd)
|
||||
finally:
|
||||
os.close(directory_fd)
|
||||
|
||||
|
||||
def _write_json_atomic(path: Path, payload: object, *, overwrite: bool = False) -> None:
|
||||
temporary = path.with_name(f"{path.name}.tmp")
|
||||
if overwrite:
|
||||
temporary.unlink(missing_ok=True)
|
||||
elif path.exists() or path.is_symlink():
|
||||
raise FileExistsError(path)
|
||||
flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
|
||||
descriptor = os.open(temporary, flags, 0o600)
|
||||
try:
|
||||
with os.fdopen(descriptor, "wb") as output:
|
||||
output.write(
|
||||
json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
|
||||
)
|
||||
output.flush()
|
||||
os.fsync(output.fileno())
|
||||
os.replace(temporary, path)
|
||||
_fsync_directory(path.parent)
|
||||
except BaseException:
|
||||
temporary.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
|
||||
def write_manifest_atomic(path: Path, manifest: SnapshotManifest) -> None:
|
||||
"""Write a new private manifest and make its directory entry durable."""
|
||||
path = Path(path)
|
||||
validate_artifact_root(path, creating=False)
|
||||
destination = path / "manifest.json"
|
||||
try:
|
||||
_write_json_atomic(destination, manifest.model_dump(mode="json"))
|
||||
except FileExistsError as error:
|
||||
raise SnapshotSecurityError(
|
||||
f"manifest already exists: {destination}"
|
||||
) from error
|
||||
|
||||
|
||||
def read_manifest(path: Path) -> SnapshotManifest:
|
||||
"""Read a snapshot published by atomically installing its manifest."""
|
||||
path = Path(path)
|
||||
validate_artifact_root(path, creating=False)
|
||||
manifest_path = path / "manifest.json"
|
||||
if manifest_path.is_symlink():
|
||||
raise SnapshotSecurityError(f"manifest is a symlink: {manifest_path}")
|
||||
flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
|
||||
try:
|
||||
file_descriptor = os.open(manifest_path, flags)
|
||||
except FileNotFoundError as error:
|
||||
raise SnapshotSecurityError(
|
||||
f"manifest does not exist: {manifest_path}"
|
||||
) from error
|
||||
try:
|
||||
with os.fdopen(file_descriptor, encoding="utf-8") as manifest_file:
|
||||
payload = manifest_file.read()
|
||||
except UnicodeDecodeError as error:
|
||||
raise SnapshotCompatibilityError("invalid snapshot manifest: JSON") from error
|
||||
try:
|
||||
return SnapshotManifest.model_validate_json(payload, strict=True)
|
||||
except ValidationError as error:
|
||||
raise SnapshotCompatibilityError(
|
||||
f"invalid snapshot manifest: {_validation_path(error)}"
|
||||
) from error
|
||||
|
||||
|
||||
def inspect_snapshot(path: Path) -> dict[str, object]:
|
||||
return read_manifest(path).model_dump(mode="json")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,252 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import stat
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
from vllm.snapshot.manifest import ReleaseMarker, _validation_path, _write_json_atomic
|
||||
from vllm.snapshot.types import Oracle, oracles_match
|
||||
|
||||
_CANARY_PROMPT = "The capital of France is"
|
||||
|
||||
|
||||
class SnapshotCanaryError(RuntimeError):
|
||||
"""The initialized engine did not produce a valid snapshot oracle."""
|
||||
|
||||
|
||||
class SnapshotBarrierError(RuntimeError):
|
||||
"""The controller supplied an invalid snapshot release marker."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ListenerConfig:
|
||||
host: str | None
|
||||
port: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ControlArgs:
|
||||
ready_file: Path
|
||||
release_file: Path
|
||||
release_timeout_s: float
|
||||
|
||||
|
||||
def write_ready_atomic(path: Path, oracle: Oracle) -> None:
|
||||
_write_json_atomic(
|
||||
path,
|
||||
{
|
||||
"sampled_token_logprob": oracle.sampled_token_logprob,
|
||||
"token_ids": oracle.token_ids,
|
||||
"text": oracle.text,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def read_release_marker(path: Path) -> ListenerConfig:
|
||||
path = Path(path)
|
||||
path_stat = path.lstat()
|
||||
if stat.S_ISLNK(path_stat.st_mode):
|
||||
raise SnapshotBarrierError("release marker must not be a symlink")
|
||||
try:
|
||||
payload = json.loads(path.read_text())
|
||||
except (OSError, json.JSONDecodeError) as error:
|
||||
raise SnapshotBarrierError("release marker is not valid JSON") from error
|
||||
try:
|
||||
marker = ReleaseMarker.model_validate(payload)
|
||||
except ValidationError as error:
|
||||
raise SnapshotBarrierError(
|
||||
f"release marker is invalid: {_validation_path(error)}"
|
||||
) from error
|
||||
return ListenerConfig(host=marker.host, port=marker.port)
|
||||
|
||||
|
||||
def parse_control_args(argv: list[str]) -> tuple[ControlArgs, list[str]]:
|
||||
parser = argparse.ArgumentParser(add_help=False)
|
||||
parser.add_argument("--ready-file", type=Path, required=True)
|
||||
parser.add_argument("--release-file", type=Path, required=True)
|
||||
parser.add_argument("--release-timeout-s", type=float, default=900.0)
|
||||
control_args, remaining = parser.parse_known_args(argv)
|
||||
if remaining and remaining[0] == "--":
|
||||
remaining = remaining[1:]
|
||||
if control_args.release_timeout_s <= 0:
|
||||
raise ValueError("release timeout must be positive")
|
||||
return (
|
||||
ControlArgs(
|
||||
ready_file=control_args.ready_file,
|
||||
release_file=control_args.release_file,
|
||||
release_timeout_s=control_args.release_timeout_s,
|
||||
),
|
||||
remaining,
|
||||
)
|
||||
|
||||
|
||||
async def _release_reloadable_state(engine: Any) -> None:
|
||||
"""Discard model and KV state before the process image is captured."""
|
||||
await engine.sleep(level=2)
|
||||
|
||||
|
||||
async def _restore_reloadable_state(engine: Any) -> None:
|
||||
"""Rebuild state discarded by ``_release_reloadable_state``."""
|
||||
await engine.wake_up(tags=["weights"])
|
||||
await engine.collective_rpc("reload_weights")
|
||||
await engine.wake_up(tags=["kv_cache"])
|
||||
|
||||
|
||||
def oracle_from_request_output(request_output: Any) -> Oracle:
|
||||
try:
|
||||
candidate = request_output.outputs[0]
|
||||
token_ids = tuple(candidate.token_ids)
|
||||
(sampled_token_id,) = token_ids
|
||||
return Oracle(
|
||||
token_ids=token_ids,
|
||||
text=candidate.text,
|
||||
sampled_token_logprob=candidate.logprobs[0][sampled_token_id].logprob,
|
||||
)
|
||||
except ValueError as error:
|
||||
raise SnapshotCanaryError(
|
||||
"snapshot canary must produce exactly one finite token logprob"
|
||||
) from error
|
||||
except (AttributeError, IndexError, KeyError, TypeError) as error:
|
||||
raise SnapshotCanaryError(
|
||||
"snapshot canary did not return sampled token logprob"
|
||||
) from error
|
||||
|
||||
|
||||
async def run_engine_canary(engine: Any) -> Oracle:
|
||||
from vllm import SamplingParams
|
||||
|
||||
final_output = None
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0,
|
||||
min_tokens=1,
|
||||
max_tokens=1,
|
||||
seed=0,
|
||||
logprobs=0,
|
||||
)
|
||||
async for output in engine.generate(
|
||||
_CANARY_PROMPT,
|
||||
sampling_params,
|
||||
request_id="vllm-snapshot-canary",
|
||||
):
|
||||
final_output = output
|
||||
if final_output is None:
|
||||
raise SnapshotCanaryError("canary generation returned no output")
|
||||
return oracle_from_request_output(final_output)
|
||||
|
||||
|
||||
def detach_snapshot_streams() -> None:
|
||||
"""Remove launch-log descriptors from the reusable process image."""
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
sink = os.open(os.devnull, os.O_WRONLY)
|
||||
try:
|
||||
os.dup2(sink, sys.stdout.fileno())
|
||||
os.dup2(sink, sys.stderr.fileno())
|
||||
finally:
|
||||
os.close(sink)
|
||||
|
||||
|
||||
async def wait_for_release_marker(
|
||||
path: Path,
|
||||
*,
|
||||
timeout_s: float,
|
||||
poll_interval_s: float = 0.05,
|
||||
) -> ListenerConfig:
|
||||
if timeout_s <= 0:
|
||||
raise ValueError("release timeout must be positive")
|
||||
if poll_interval_s <= 0:
|
||||
raise ValueError("release poll interval must be positive")
|
||||
remaining_s = timeout_s
|
||||
while not path.exists():
|
||||
if remaining_s <= 0:
|
||||
raise TimeoutError(f"release marker not found before timeout: {path}")
|
||||
delay_s = min(poll_interval_s, remaining_s)
|
||||
await asyncio.sleep(delay_s)
|
||||
remaining_s -= delay_s
|
||||
return read_release_marker(path)
|
||||
|
||||
|
||||
def parse_vllm_args(argv: list[str]) -> Any:
|
||||
from vllm.entrypoints.launchers.cli_args import (
|
||||
make_arg_parser,
|
||||
validate_parsed_serve_args,
|
||||
)
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
parser = FlexibleArgumentParser(prog="python -m vllm.snapshot.server")
|
||||
args = make_arg_parser(parser).parse_args(argv)
|
||||
if args.model_tag is not None:
|
||||
args.model = args.model_tag
|
||||
if args.grpc or args.headless:
|
||||
raise ValueError("snapshot server requires the HTTP frontend")
|
||||
if args.api_server_count not in (None, 1):
|
||||
raise ValueError("snapshot server supports one API server")
|
||||
args.api_server_count = None
|
||||
validate_parsed_serve_args(args)
|
||||
return args
|
||||
|
||||
|
||||
async def run_vllm_snapshot_child(control: ControlArgs, args: Any) -> None:
|
||||
from vllm.entrypoints.launchers.api_server.entry import (
|
||||
build_and_serve,
|
||||
build_async_engine_client,
|
||||
)
|
||||
from vllm.entrypoints.launchers.launcher import setup_server
|
||||
|
||||
args.enable_sleep_mode = True
|
||||
|
||||
async with build_async_engine_client(args) as engine:
|
||||
oracle = await run_engine_canary(engine)
|
||||
try:
|
||||
await _release_reloadable_state(engine)
|
||||
await _restore_reloadable_state(engine)
|
||||
rehearsal_oracle = await run_engine_canary(engine)
|
||||
if not oracles_match(oracle, rehearsal_oracle):
|
||||
raise SnapshotCanaryError("snapshot rehearsal changed canary output")
|
||||
await _release_reloadable_state(engine)
|
||||
except SnapshotCanaryError:
|
||||
raise
|
||||
except Exception as error:
|
||||
raise SnapshotCanaryError("snapshot rehearsal failed") from error
|
||||
|
||||
detach_snapshot_streams()
|
||||
write_ready_atomic(control.ready_file, oracle)
|
||||
listener = await wait_for_release_marker(
|
||||
control.release_file,
|
||||
timeout_s=control.release_timeout_s,
|
||||
)
|
||||
args.host = listener.host
|
||||
args.port = listener.port
|
||||
await _restore_reloadable_state(engine)
|
||||
listen_address, sock = setup_server(args, reuse_port=False)
|
||||
try:
|
||||
shutdown_task = await build_and_serve(
|
||||
engine,
|
||||
listen_address,
|
||||
sock,
|
||||
args,
|
||||
)
|
||||
await shutdown_task
|
||||
finally:
|
||||
sock.close()
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> None:
|
||||
control, remaining = parse_control_args(sys.argv[1:] if argv is None else argv)
|
||||
args = parse_vllm_args(remaining)
|
||||
import uvloop
|
||||
|
||||
uvloop.run(run_vllm_snapshot_child(control, args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,37 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Oracle:
|
||||
token_ids: tuple[int, ...]
|
||||
text: str
|
||||
sampled_token_logprob: float
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if len(self.token_ids) != 1:
|
||||
raise ValueError("snapshot oracle must contain exactly one token")
|
||||
value = self.sampled_token_logprob
|
||||
if (
|
||||
isinstance(value, bool)
|
||||
or not isinstance(value, (int, float))
|
||||
or not math.isfinite(value)
|
||||
):
|
||||
raise ValueError("sampled token logprob must be finite")
|
||||
object.__setattr__(self, "sampled_token_logprob", float(value))
|
||||
|
||||
|
||||
def oracles_match(expected: Oracle, actual: Oracle) -> bool:
|
||||
return (
|
||||
expected.token_ids == actual.token_ids
|
||||
and expected.text == actual.text
|
||||
and math.isclose(
|
||||
expected.sampled_token_logprob,
|
||||
actual.sampled_token_logprob,
|
||||
rel_tol=0.0,
|
||||
abs_tol=1e-3,
|
||||
)
|
||||
)
|
||||
@@ -519,7 +519,7 @@ class FlexibleArgumentParser(ArgumentParser):
|
||||
config_args = self.load_config_file(file_path)
|
||||
|
||||
# 0th index might be the sub command {serve,chat,complete,...}
|
||||
# optionally followed by model_tag (only for serve)
|
||||
# optionally followed by model_tag (serve or snapshot create)
|
||||
# followed by config args
|
||||
# followed by rest of cli args.
|
||||
# maintaining this order will enforce the precedence
|
||||
@@ -527,8 +527,11 @@ class FlexibleArgumentParser(ArgumentParser):
|
||||
if args[0].startswith("-"):
|
||||
# No sub command (e.g., api_server entry point)
|
||||
args = config_args + args[0:index] + args[index + 2 :]
|
||||
elif args[0] == "serve":
|
||||
model_in_cli = len(args) > 1 and not args[1].startswith("-")
|
||||
elif args[0] == "serve" or args[:2] == ["snapshot", "create"]:
|
||||
model_index = 1 if args[0] == "serve" else 2
|
||||
model_in_cli = len(args) > model_index and not args[model_index].startswith(
|
||||
"-"
|
||||
)
|
||||
model_in_config = any(arg == "--model" for arg in config_args)
|
||||
|
||||
if not model_in_cli and not model_in_config:
|
||||
@@ -540,15 +543,19 @@ class FlexibleArgumentParser(ArgumentParser):
|
||||
if model_in_cli:
|
||||
# Model specified as positional arg, keep CLI version
|
||||
args = (
|
||||
[args[0]]
|
||||
+ [args[1]]
|
||||
args[: model_index + 1]
|
||||
+ config_args
|
||||
+ args[2:index]
|
||||
+ args[model_index + 1 : index]
|
||||
+ args[index + 2 :]
|
||||
)
|
||||
else:
|
||||
# No model in CLI, use config if available
|
||||
args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
|
||||
args = (
|
||||
args[:model_index]
|
||||
+ config_args
|
||||
+ args[model_index:index]
|
||||
+ args[index + 2 :]
|
||||
)
|
||||
else:
|
||||
args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user