| 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" |