From 194148b21ad850747d64abd1b3f2a378d2dffaf0 Mon Sep 17 00:00:00 2001 From: Zhehao Hu Date: Fri, 16 Oct 2026 21:36:12 +0000 Subject: [PATCH 1/4] launcher: harden the vLLM smoke test and newest-checkpoint selection - Add latest_checkpoint() to service_utils.sh: picks /final if present, else the highest numeric step, keyed on the directory's own name. It replaces three ad hoc sorts; ar_eval_mtbench.sh and vllm_smoke_test.sh keyed on a dash-separated field of the whole path, so a dash in a parent directory (e.g. /data/my-drafts/...) picked the wrong checkpoint. vllm_smoke_test.sh now prefers exported-checkpoint-final, as quantize_drafter.sh already did. - vllm_smoke_test.sh: compute the mean acceptance length over the whole run from vLLM's /metrics counters (1 + accepted / drafts). The log's "Mean acceptance length" line covers only the last logging interval. - vllm_smoke_test.sh: with MIN_ACCEPTANCE_LENGTH set, fail when no speculative drafts were recorded instead of warning and passing. - vllm_smoke_test.sh: time prompts with shell arithmetic and python3 instead of bc, which the vLLM image lacks (every tok/s figure printed "?"). - vllm_smoke_test.sh: optional SMOKE_JSON writes the verdict, acceptance length and throughput as JSON for downstream gates, creating its parent directory if needed. Sourcing service_utils.sh is now required; it runs after the EXIT trap, so a failure there still writes this run's verdict instead of leaving an earlier run's SMOKE_JSON in place. Signed-off-by: Zhehao Hu --- tools/launcher/common/service_utils.sh | 21 +++ .../common/specdec/ar_eval_mtbench.sh | 2 +- .../common/specdec/quantize_drafter.sh | 7 +- .../common/specdec/vllm_smoke_test.sh | 167 +++++++++++++----- tools/launcher/tests/test_service_utils.py | 65 +++++++ 5 files changed, 215 insertions(+), 47 deletions(-) create mode 100644 tools/launcher/tests/test_service_utils.py diff --git a/tools/launcher/common/service_utils.sh b/tools/launcher/common/service_utils.sh index 50bd530ed3e..8c45c34b1b2 100755 --- a/tools/launcher/common/service_utils.sh +++ b/tools/launcher/common/service_utils.sh @@ -47,6 +47,27 @@ function report_result { fi } +function latest_checkpoint { + local prefix=$2 best="" best_step=-1 candidate step + for candidate in "$1"/"${prefix}"*; do + [ -d "$candidate" ] || continue + step=${candidate##*/} + step=${step#"$prefix"} + if [ "$step" = final ]; then + echo "$candidate" + return + fi + [[ "$step" =~ ^[0-9]+$ ]] || continue + if [ "$step" -gt "$best_step" ]; then + best=$candidate + best_step=$step + fi + done + if [ -n "$best" ]; then + echo "$best" + fi +} + function util_install_extra_dep { local _marker=/tmp/.nmm_extra_dep_installed if [[ -f "$_marker" ]]; then diff --git a/tools/launcher/common/specdec/ar_eval_mtbench.sh b/tools/launcher/common/specdec/ar_eval_mtbench.sh index 4062e6b2fed..8b5c734a6a2 100644 --- a/tools/launcher/common/specdec/ar_eval_mtbench.sh +++ b/tools/launcher/common/specdec/ar_eval_mtbench.sh @@ -42,7 +42,7 @@ done MODEL_PATH="" if [ -n "$CKPT_DIR" ]; then # Find latest checkpoint subdir - LAST_CKPT=$(ls -d ${CKPT_DIR}/checkpoint-* 2>/dev/null | sort -t- -k2 -n | tail -1) + LAST_CKPT=$(latest_checkpoint "${CKPT_DIR}" checkpoint-) if [ -f "${CKPT_DIR}/model.safetensors" ]; then MODEL_PATH="${CKPT_DIR}" elif [ -n "$LAST_CKPT" ]; then diff --git a/tools/launcher/common/specdec/quantize_drafter.sh b/tools/launcher/common/specdec/quantize_drafter.sh index ddaac78ed00..6eba5218a36 100644 --- a/tools/launcher/common/specdec/quantize_drafter.sh +++ b/tools/launcher/common/specdec/quantize_drafter.sh @@ -30,15 +30,14 @@ trap 'error_handler $0 $LINENO' ERR DRAFTER="${DRAFTER_CKPT}" # Training writes exported-checkpoint-/ under output_dir; take the newest. Only for a # local directory -- anything else (an HF repo id) is passed through for the script to -# resolve. -V sorts numerically, so checkpoint-1000 beats checkpoint-900. +# resolve. if [ -d "${DRAFTER}" ] && [ ! -f "${DRAFTER}/config.json" ]; then - latest=$(find "${DRAFTER}" -maxdepth 1 -mindepth 1 -type d \ - -name 'exported-checkpoint-*' -printf '%f\n' 2>/dev/null | sort -V | tail -1) + latest=$(latest_checkpoint "${DRAFTER}" exported-checkpoint-) if [ -z "${latest}" ]; then echo "ERROR: ${DRAFTER} is not a checkpoint and holds no exported-checkpoint-* directory." exit 1 fi - DRAFTER="${DRAFTER}/${latest}" + DRAFTER="${latest}" echo "Auto-detected drafter: ${DRAFTER}" fi diff --git a/tools/launcher/common/specdec/vllm_smoke_test.sh b/tools/launcher/common/specdec/vllm_smoke_test.sh index 8028a570a8e..97419de589e 100644 --- a/tools/launcher/common/specdec/vllm_smoke_test.sh +++ b/tools/launcher/common/specdec/vllm_smoke_test.sh @@ -28,21 +28,82 @@ # VLLM_PORT — server port (default: 8000) # REASONING_PARSER — reasoning parser (e.g., "qwen3" for Qwen3.5) # DISABLE_PREFIX_CACHING — set to "1" to disable prefix caching +# MIN_ACCEPTANCE_LENGTH — fail unless the mean acceptance length (from vLLM's /metrics +# counters) reaches this; also fails if no speculative drafts were recorded +# SMOKE_JSON — also write the result to this file as JSON, on every exit; "stage" +# records how far the run got (starting, draft_missing, server_died, +# server_timeout, serving, completed) SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" -source ${SCRIPT_DIR}/../service_utils.sh 2>/dev/null || true # Ensure pandas is available (missing in some vLLM nightly builds) pip install pandas 2>/dev/null || true -cleanup() { kill $SERVER_PID 2>/dev/null; sleep 2; kill -9 $SERVER_PID 2>/dev/null; rm -f "${VLLM_LOG:-}" 2>/dev/null; } -trap cleanup EXIT - MODEL=${HF_MODEL_CKPT} DRAFT=${DRAFT_MODEL:-} +METHOD=${SPEC_METHOD:-eagle} +NUM_SPEC=${NUM_SPEC_TOKENS:-15} +STAGE=starting +VERDICT=FAIL +REASON="" +PASS=0 +FAIL=0 +TOTAL_TOKENS=0 +AVG_TPS="?" +AVG_ACCEPT="" + +write_smoke_json() { + [ -n "${SMOKE_JSON:-}" ] || return 0 + python3 - "$SMOKE_JSON" "$STAGE" "$VERDICT" "$REASON" "$METHOD" "$MODEL" "${DRAFT:-}" "$NUM_SPEC" \ + "$PASS" "$FAIL" "$TOTAL_TOKENS" "$AVG_TPS" "${AVG_ACCEPT:-}" "${MIN_ACCEPTANCE_LENGTH:-}" <<'EOF' +import json +import os +import sys + +(path, stage, verdict, reason, method, target, draft, num_spec, n_pass, n_fail, tokens, tps, al, + min_al) = sys.argv[1:] + + +def num(value): + try: + return float(value) + except ValueError: + return None + + +result = { + "stage": stage, + "pass": verdict == "PASS", + "reason": reason or None, + "method": method, + "target": target, + "draft": draft or None, + "num_speculative_tokens": int(num_spec), + "prompts_passed": int(n_pass), + "prompts_failed": int(n_fail), + "total_tokens": int(tokens), + "avg_tok_s": num(tps), + "mean_acceptance_length": num(al), + "min_acceptance_length": num(min_al), +} +# A smoke-only experiment has nothing else that creates the directory. +os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) +with open(path, "w") as f: + json.dump(result, f, indent=2) +EOF +} + +# A result left by an earlier run must never be read as this run's. +[ -n "${SMOKE_JSON:-}" ] && rm -f "$SMOKE_JSON" +cleanup() { write_smoke_json; kill $SERVER_PID 2>/dev/null; sleep 2; kill -9 $SERVER_PID 2>/dev/null; rm -f "${VLLM_LOG:-}" 2>/dev/null; } +trap cleanup EXIT + +# After the trap, so a failure here still writes this run's SMOKE_JSON. +source "${SCRIPT_DIR}/../service_utils.sh" || { REASON="cannot source service_utils.sh"; echo "ERROR: $REASON"; exit 1; } + # Auto-detect exported checkpoint from training output dir if [ -z "$DRAFT" ] && [ -n "${DRAFT_CKPT_DIR:-}" ]; then - DRAFT=$(ls -d ${DRAFT_CKPT_DIR}/exported-checkpoint-* 2>/dev/null | sort -t- -k3 -n | tail -1) + DRAFT=$(latest_checkpoint "${DRAFT_CKPT_DIR}" exported-checkpoint-) if [ -n "$DRAFT" ]; then echo "Auto-detected draft model: ${DRAFT}" else @@ -52,11 +113,11 @@ if [ -z "$DRAFT" ] && [ -n "${DRAFT_CKPT_DIR:-}" ]; then # "num_speculative_tokens was provided but without speculative model" error. echo "ERROR: DRAFT_CKPT_DIR=${DRAFT_CKPT_DIR} contains no exported-checkpoint-* directory." echo " The upstream training task likely failed to produce a draft model." + STAGE=draft_missing + REASON="DRAFT_CKPT_DIR=${DRAFT_CKPT_DIR} contains no exported-checkpoint-* directory" exit 1 fi fi -METHOD=${SPEC_METHOD:-eagle} -NUM_SPEC=${NUM_SPEC_TOKENS:-15} PORT=${VLLM_PORT:-8000} TP=${TP_SIZE:-1} @@ -111,23 +172,25 @@ for i in $(seq 1 180); do break fi if ! kill -0 $SERVER_PID 2>/dev/null; then - echo "ERROR: Server died"; wait $SERVER_PID; exit 1 + echo "ERROR: Server died"; wait $SERVER_PID + STAGE=server_died; REASON="vLLM server exited before becoming healthy"; exit 1 fi sleep 1 done if ! curl -s http://localhost:${PORT}/health > /dev/null 2>&1; then - echo "ERROR: Server timeout"; exit 1 + echo "ERROR: Server timeout" + STAGE=server_timeout; REASON="vLLM server not healthy after 180s"; exit 1 fi +STAGE=serving +REASON="exited while sending test prompts" # Run quick test prompts using chat completions API MAX_TOKENS=${MAX_OUTPUT_TOKENS:-1024} echo "" echo "=== Test Prompts (max_tokens=${MAX_TOKENS}) ===" -PASS=0 -FAIL=0 -TOTAL_TOKENS=0 -TOTAL_TIME=0 +TOTAL_NS=0 +ratio() { python3 -c 'import sys; a, b = float(sys.argv[1]), float(sys.argv[2]); print(f"{a / b:.{sys.argv[3]}f}" if b else "?")' "$@"; } # 8 prompts mimicking MT-Bench categories: writing, roleplay, reasoning, # math, coding, extraction, stem, humanities for PROMPT in \ @@ -145,14 +208,15 @@ for PROMPT in \ -d "{\"model\": \"${MODEL}\", \"messages\": [{\"role\": \"user\", \"content\": \"${PROMPT}\"}], \"max_tokens\": ${MAX_TOKENS}, \"temperature\": 0}" \ 2>/dev/null) END=$(date +%s%N) - ELAPSED=$(echo "scale=2; ($END - $START) / 1000000000" | bc 2>/dev/null || echo "0") + ELAPSED_NS=$((END - START)) + ELAPSED=$(ratio "$ELAPSED_NS" 1e9 2) TOKENS=$(echo "$RESULT" | python3 -c "import json,sys; r=json.load(sys.stdin); print(r.get('usage',{}).get('completion_tokens',0))" 2>/dev/null) if [ -n "$TOKENS" ] && [ "$TOKENS" -gt 0 ] 2>/dev/null; then - TPS=$(echo "scale=1; $TOKENS / $ELAPSED" | bc 2>/dev/null || echo "?") + TPS=$(ratio "${TOKENS}e9" "$ELAPSED_NS" 1) echo " PASS: ${TOKENS} tokens in ${ELAPSED}s (${TPS} tok/s) — \"${PROMPT:0:50}...\"" PASS=$((PASS + 1)) TOTAL_TOKENS=$((TOTAL_TOKENS + TOKENS)) - TOTAL_TIME=$(echo "$TOTAL_TIME + $ELAPSED" | bc 2>/dev/null || echo "0") + TOTAL_NS=$((TOTAL_NS + ELAPSED_NS)) else echo " FAIL: \"${PROMPT}\"" echo " Response: $(echo "$RESULT" | head -c 200)" @@ -163,44 +227,63 @@ done echo "" echo "Results: ${PASS} passed, ${FAIL} failed" if [ "$TOTAL_TOKENS" -gt 0 ] 2>/dev/null; then - AVG_TPS=$(echo "scale=1; $TOTAL_TOKENS / $TOTAL_TIME" | bc 2>/dev/null || echo "?") - echo "Total: ${TOTAL_TOKENS} tokens in ${TOTAL_TIME}s (${AVG_TPS} tok/s avg)" + AVG_TPS=$(ratio "${TOTAL_TOKENS}e9" "$TOTAL_NS" 1) + echo "Total: ${TOTAL_TOKENS} tokens in $(ratio "$TOTAL_NS" 1e9 2)s (${AVG_TPS} tok/s avg)" fi # Fetch speculative decoding metrics if available echo "" -METRICS=$(curl -s http://localhost:${PORT}/metrics 2>/dev/null | grep -i "spec\|accept\|draft" | head -10) +ALL_METRICS=$(curl -s http://localhost:${PORT}/metrics 2>/dev/null || true) +METRICS=$(echo "$ALL_METRICS" | grep -i "spec\|accept\|draft" | head -10) if [ -n "$METRICS" ]; then echo "=== Speculative Decoding Metrics ===" echo "$METRICS" fi -if [ $FAIL -gt 0 ]; then - echo "ERROR: Some prompts failed" - exit 1 -fi +# Mean acceptance length over all prompts, from vLLM's cumulative Prometheus counters: +# 1 + accepted draft tokens / drafts (the bonus token counts, as in vLLM's own log). The +# "SpecDecoding metrics: Mean acceptance length" log line covers only the last logging +# interval (vLLM resets its stats after each log), so it is not used. +AVG_ACCEPT=$(echo "$ALL_METRICS" | python3 -c ' +import sys -# Regression check: minimum acceptance length for speculative decoding -if [ -n "${MIN_ACCEPTANCE_LENGTH:-}" ]; then - # Parse mean acceptance length from vLLM's SpecDecoding metrics log. - # vLLM logs: "SpecDecoding metrics: Mean acceptance length: X.XX, ..." - # Take the last reported value (most accurate, covers all prompts). - AVG_ACCEPT=$(grep -oP 'Mean acceptance length: \K[0-9.]+' "$VLLM_LOG" 2>/dev/null | tail -1 || true) - if [ -n "$AVG_ACCEPT" ]; then - echo "" - echo "=== Acceptance Length Regression Check ===" - echo " Mean acceptance length: ${AVG_ACCEPT}" - echo " Threshold: ${MIN_ACCEPTANCE_LENGTH}" - PASS_CHECK=$(python3 -c "print('yes' if float('${AVG_ACCEPT}') >= float('${MIN_ACCEPTANCE_LENGTH}') else 'no')") - if [ "$PASS_CHECK" = "yes" ]; then - echo " PASS: ${AVG_ACCEPT} >= ${MIN_ACCEPTANCE_LENGTH}" - else - echo " REGRESSION: ${AVG_ACCEPT} < ${MIN_ACCEPTANCE_LENGTH}" - exit 1 - fi +drafts = accepted = 0.0 +for line in sys.stdin: + name = line.split("{")[0].split(" ")[0] + if name == "vllm:spec_decode_num_drafts_total": + drafts += float(line.split()[-1]) + elif name == "vllm:spec_decode_num_accepted_tokens_total": + accepted += float(line.split()[-1]) +if drafts > 0: + print(f"{1 + accepted / drafts:.4f}") +' 2>/dev/null || true) + +STAGE=completed +VERDICT=PASS +REASON="" +if [ $FAIL -gt 0 ]; then + VERDICT=FAIL + REASON="ERROR: Some prompts failed" +elif [ -n "${MIN_ACCEPTANCE_LENGTH:-}" ]; then + echo "" + echo "=== Acceptance Length Regression Check ===" + echo " Mean acceptance length: ${AVG_ACCEPT:-not reported}" + echo " Threshold: ${MIN_ACCEPTANCE_LENGTH}" + if [ -z "$AVG_ACCEPT" ]; then + # A threshold that was never evaluated must not pass. + VERDICT=FAIL + REASON="REGRESSION: MIN_ACCEPTANCE_LENGTH=${MIN_ACCEPTANCE_LENGTH} was set but vLLM /metrics reports no speculative drafts" + elif python3 -c 'import sys; sys.exit(float(sys.argv[1]) < float(sys.argv[2]))' "$AVG_ACCEPT" "$MIN_ACCEPTANCE_LENGTH"; then + echo " PASS: ${AVG_ACCEPT} >= ${MIN_ACCEPTANCE_LENGTH}" else - echo "WARNING: Could not parse acceptance length from vLLM log, skipping regression check" + VERDICT=FAIL + REASON="REGRESSION: ${AVG_ACCEPT} < ${MIN_ACCEPTANCE_LENGTH}" fi fi +if [ "$VERDICT" = FAIL ]; then + echo " $REASON" + exit 1 +fi + echo "Done" diff --git a/tools/launcher/tests/test_service_utils.py b/tools/launcher/tests/test_service_utils.py new file mode 100644 index 00000000000..fb8a670fc25 --- /dev/null +++ b/tools/launcher/tests/test_service_utils.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""latest_checkpoint in common/service_utils.sh.""" + +import subprocess +from pathlib import Path + +import pytest + +_SERVICE_UTILS = Path(__file__).resolve().parents[1] / "common" / "service_utils.sh" + + +def _latest(root, prefix="exported-checkpoint-"): + # Sourcing appends to ./modules/Model-Optimizer/modelopt/__init__.py; from the launcher + # dir that is this checkout, so run in the temp dir. + out = subprocess.run( + [ + "bash", + "-c", + 'source "$0" >/dev/null 2>&1; latest_checkpoint "$1" "$2"', + _SERVICE_UTILS, + root, + prefix, + ], + cwd=root, + capture_output=True, + text=True, + check=True, + ) + return out.stdout.strip() + + +@pytest.mark.parametrize( + ("names", "expected"), + [ + # Numeric order, not string order; non-numeric suffixes are ignored. + (["900", "1000", "tmp"], "1000"), + # final wins over any step. + (["900", "1000", "final"], "final"), + ([], None), + ], +) +def test_latest_checkpoint(tmp_path, names, expected): + """Final wins, else the highest step; a dash in a parent dir once broke the parsing.""" + root = tmp_path / "my-drafts" / "run-1" + root.mkdir(parents=True) + for name in names: + (root / f"exported-checkpoint-{name}").mkdir() + (root / "exported-checkpoint-5000.txt").touch() # files are not checkpoints + + want = str(root / f"exported-checkpoint-{expected}") if expected else "" + assert _latest(root) == want From 06f49803d8544d185c39ca9abbefd075f444b4ed Mon Sep 17 00:00:00 2001 From: Zhehao Hu Date: Fri, 23 Oct 2026 00:04:31 +0000 Subject: [PATCH 2/4] launcher: fix smoke-test startup outside a launcher workdir and on bad numbers Sourcing service_utils.sh returns the status of its last command, a rank-0 append to modules/Model-Optimizer/modelopt/__init__.py. Outside a launcher working directory that append fails, so the smoke test stopped with "cannot source service_utils.sh". Only a missing file is fatal now. A non-integer NUM_SPEC_TOKENS crashed write_smoke_json inside the EXIT trap, so no SMOKE_JSON was written; a non-numeric MIN_ACCEPTANCE_LENGTH was reported as "REGRESSION: X < Y". Both are checked before any work and fail with the new stage invalid_config. The writer also turns a non-integer or non-finite value into null, so the file is always written and is strict JSON. The prompt-failure reason drops its "ERROR:" prefix and gives the count, and the prompt counters are renamed PROMPTS_OK/PROMPTS_FAILED so they no longer share FAIL with service_utils.sh. cleanup waits for the vLLM server to stop only when one was started, so an early exit no longer sleeps 2 seconds. Signed-off-by: Zhehao Hu --- .../common/specdec/vllm_smoke_test.sh | 51 +++++++---- tools/launcher/tests/test_vllm_smoke_test.py | 86 +++++++++++++++++++ 2 files changed, 122 insertions(+), 15 deletions(-) create mode 100644 tools/launcher/tests/test_vllm_smoke_test.py diff --git a/tools/launcher/common/specdec/vllm_smoke_test.sh b/tools/launcher/common/specdec/vllm_smoke_test.sh index 97419de589e..8cd0dd76ee9 100644 --- a/tools/launcher/common/specdec/vllm_smoke_test.sh +++ b/tools/launcher/common/specdec/vllm_smoke_test.sh @@ -31,8 +31,8 @@ # MIN_ACCEPTANCE_LENGTH — fail unless the mean acceptance length (from vLLM's /metrics # counters) reaches this; also fails if no speculative drafts were recorded # SMOKE_JSON — also write the result to this file as JSON, on every exit; "stage" -# records how far the run got (starting, draft_missing, server_died, -# server_timeout, serving, completed) +# records how far the run got (invalid_config, starting, draft_missing, +# server_died, server_timeout, serving, completed) SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" @@ -46,8 +46,8 @@ NUM_SPEC=${NUM_SPEC_TOKENS:-15} STAGE=starting VERDICT=FAIL REASON="" -PASS=0 -FAIL=0 +PROMPTS_OK=0 +PROMPTS_FAILED=0 TOTAL_TOKENS=0 AVG_TPS="?" AVG_ACCEPT="" @@ -55,8 +55,9 @@ AVG_ACCEPT="" write_smoke_json() { [ -n "${SMOKE_JSON:-}" ] || return 0 python3 - "$SMOKE_JSON" "$STAGE" "$VERDICT" "$REASON" "$METHOD" "$MODEL" "${DRAFT:-}" "$NUM_SPEC" \ - "$PASS" "$FAIL" "$TOTAL_TOKENS" "$AVG_TPS" "${AVG_ACCEPT:-}" "${MIN_ACCEPTANCE_LENGTH:-}" <<'EOF' + "$PROMPTS_OK" "$PROMPTS_FAILED" "$TOTAL_TOKENS" "$AVG_TPS" "${AVG_ACCEPT:-}" "${MIN_ACCEPTANCE_LENGTH:-}" <<'EOF' import json +import math import os import sys @@ -66,9 +67,14 @@ import sys def num(value): try: - return float(value) + value = float(value) except ValueError: return None + return value if math.isfinite(value) else None + + +def count(value): + return int(value) if value.isdigit() else None result = { @@ -78,7 +84,7 @@ result = { "method": method, "target": target, "draft": draft or None, - "num_speculative_tokens": int(num_spec), + "num_speculative_tokens": count(num_spec), "prompts_passed": int(n_pass), "prompts_failed": int(n_fail), "total_tokens": int(tokens), @@ -95,11 +101,26 @@ EOF # A result left by an earlier run must never be read as this run's. [ -n "${SMOKE_JSON:-}" ] && rm -f "$SMOKE_JSON" -cleanup() { write_smoke_json; kill $SERVER_PID 2>/dev/null; sleep 2; kill -9 $SERVER_PID 2>/dev/null; rm -f "${VLLM_LOG:-}" 2>/dev/null; } +cleanup() { + write_smoke_json + if [ -n "${SERVER_PID:-}" ]; then kill $SERVER_PID 2>/dev/null; sleep 2; kill -9 $SERVER_PID 2>/dev/null; fi + rm -f "${VLLM_LOG:-}" 2>/dev/null +} trap cleanup EXIT -# After the trap, so a failure here still writes this run's SMOKE_JSON. -source "${SCRIPT_DIR}/../service_utils.sh" || { REASON="cannot source service_utils.sh"; echo "ERROR: $REASON"; exit 1; } +# Checked after the trap, so a bad value still writes this run's SMOKE_JSON. +if [ -n "${MIN_ACCEPTANCE_LENGTH:-}" ] && ! python3 -c 'import math, sys; sys.exit(not math.isfinite(float(sys.argv[1])))' "$MIN_ACCEPTANCE_LENGTH" 2>/dev/null; then + REASON="MIN_ACCEPTANCE_LENGTH='${MIN_ACCEPTANCE_LENGTH}' is not a finite number" +elif ! [[ "$NUM_SPEC" =~ ^[0-9]+$ ]] || [ "$NUM_SPEC" -eq 0 ]; then + REASON="NUM_SPEC_TOKENS='${NUM_SPEC}' is not a positive integer" +fi +if [ -n "$REASON" ]; then + STAGE=invalid_config; echo "ERROR: $REASON"; exit 1 +fi + +# source returns the status of service_utils.sh's last command, an unrelated rank-0 append. +[ -f "${SCRIPT_DIR}/../service_utils.sh" ] || { REASON="cannot find service_utils.sh"; echo "ERROR: $REASON"; exit 1; } +source "${SCRIPT_DIR}/../service_utils.sh" # Auto-detect exported checkpoint from training output dir if [ -z "$DRAFT" ] && [ -n "${DRAFT_CKPT_DIR:-}" ]; then @@ -214,18 +235,18 @@ for PROMPT in \ if [ -n "$TOKENS" ] && [ "$TOKENS" -gt 0 ] 2>/dev/null; then TPS=$(ratio "${TOKENS}e9" "$ELAPSED_NS" 1) echo " PASS: ${TOKENS} tokens in ${ELAPSED}s (${TPS} tok/s) — \"${PROMPT:0:50}...\"" - PASS=$((PASS + 1)) + PROMPTS_OK=$((PROMPTS_OK + 1)) TOTAL_TOKENS=$((TOTAL_TOKENS + TOKENS)) TOTAL_NS=$((TOTAL_NS + ELAPSED_NS)) else echo " FAIL: \"${PROMPT}\"" echo " Response: $(echo "$RESULT" | head -c 200)" - FAIL=$((FAIL + 1)) + PROMPTS_FAILED=$((PROMPTS_FAILED + 1)) fi done echo "" -echo "Results: ${PASS} passed, ${FAIL} failed" +echo "Results: ${PROMPTS_OK} passed, ${PROMPTS_FAILED} failed" if [ "$TOTAL_TOKENS" -gt 0 ] 2>/dev/null; then AVG_TPS=$(ratio "${TOTAL_TOKENS}e9" "$TOTAL_NS" 1) echo "Total: ${TOTAL_TOKENS} tokens in $(ratio "$TOTAL_NS" 1e9 2)s (${AVG_TPS} tok/s avg)" @@ -261,9 +282,9 @@ if drafts > 0: STAGE=completed VERDICT=PASS REASON="" -if [ $FAIL -gt 0 ]; then +if [ $PROMPTS_FAILED -gt 0 ]; then VERDICT=FAIL - REASON="ERROR: Some prompts failed" + REASON="${PROMPTS_FAILED} of $((PROMPTS_OK + PROMPTS_FAILED)) prompts failed" elif [ -n "${MIN_ACCEPTANCE_LENGTH:-}" ]; then echo "" echo "=== Acceptance Length Regression Check ===" diff --git a/tools/launcher/tests/test_vllm_smoke_test.py b/tools/launcher/tests/test_vllm_smoke_test.py new file mode 100644 index 00000000000..9edcd6c22a0 --- /dev/null +++ b/tools/launcher/tests/test_vllm_smoke_test.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""vllm_smoke_test.sh paths that stop before vLLM starts, and the SMOKE_JSON they leave.""" + +import json +import os +import shutil +import subprocess +from pathlib import Path + +import pytest + +_COMMON = Path(__file__).resolve().parents[1] / "common" +# service_utils.sh appends to modelopt/__init__.py only on local rank 0, so a rank set by an +# enclosing Slurm/MPI job would hide the failure under test. +_RANK_VARS = ("PMIX_LOCAL_RANK", "OMPI_COMM_WORLD_LOCAL_RANK", "SLURM_LOCALID", "SLURM_PROCID") + + +def _run(tmp_path, **env_overrides): + """Run the script from a directory with no modules/, as rank 0. Returns (proc, JSON).""" + (tmp_path / "common" / "specdec").mkdir(parents=True) + for name in ("service_utils.sh", "specdec/vllm_smoke_test.sh"): + shutil.copy(_COMMON / name, tmp_path / "common" / name) + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + for tool, body in (("pip", "exit 0"), ("nvidia-smi", "echo 1")): + (bin_dir / tool).write_text(f"#!/bin/bash\n{body}\n") + (bin_dir / tool).chmod(0o755) + (tmp_path / "train_out").mkdir() # training left no exported-checkpoint-* + + smoke_json = tmp_path / "out" / "smoke.json" + env = {k: v for k, v in os.environ.items() if k not in _RANK_VARS} + env.update( + PATH=f"{bin_dir}:{os.environ['PATH']}", + HF_MODEL_CKPT="/hf-local/target", + DRAFT_CKPT_DIR=str(tmp_path / "train_out"), + SMOKE_JSON=str(smoke_json), + **env_overrides, + ) + proc = subprocess.run( + ["bash", "common/specdec/vllm_smoke_test.sh"], + cwd=tmp_path, + env=env, + capture_output=True, + text=True, + ) + assert smoke_json.exists(), proc.stdout + proc.stderr + # Strict JSON: NaN/Infinity would break non-Python readers of the result. + return proc, json.loads(smoke_json.read_text(), parse_constant=pytest.fail) + + +def test_sources_service_utils_outside_a_launcher_workdir(tmp_path): + """The rank-0 append failing must not stop the script at `source`.""" + proc, result = _run(tmp_path) + assert proc.returncode == 1 + assert result["stage"] == "draft_missing", result["reason"] + + +@pytest.mark.parametrize( + ("env", "reason"), + [ + ({"MIN_ACCEPTANCE_LENGTH": "1,5"}, "MIN_ACCEPTANCE_LENGTH='1,5' is not a finite number"), + ({"MIN_ACCEPTANCE_LENGTH": "nan"}, "MIN_ACCEPTANCE_LENGTH='nan' is not a finite number"), + ({"NUM_SPEC_TOKENS": "x"}, "NUM_SPEC_TOKENS='x' is not a positive integer"), + ], +) +def test_invalid_config_fails_with_a_result(tmp_path, env, reason): + """A bad number fails before any work, and the trap still writes a readable result.""" + proc, result = _run(tmp_path, **env) + assert proc.returncode == 1 + assert result["stage"] == "invalid_config" + assert result["reason"] == reason + assert result["pass"] is False From 6941971bb149c5aabec8549a6a306ca88be01923 Mon Sep 17 00:00:00 2001 From: Zhehao Hu Date: Fri, 23 Oct 2026 04:37:27 +0000 Subject: [PATCH 3/4] launcher: make the smoke test's server timeout configurable vllm_smoke_test.sh waited a fixed 180 s for /health. Qwen3.8-27B with a DFlash2 drafter needs longer on GB300 (torch.compile of target and draft head, FlashInfer CuTe-DSL JIT, CUDA graph capture, GDN/Mamba Triton warmups), so every smoke attempt stopped at server_timeout. The script's exit then tore the container down mid-warmup, and the engine's next exec failed with ENOENT on ptxas, which read like the cause. SERVER_TIMEOUT sets the wait (default 180) and appears in the timeout reason. A non-integer or zero value fails up front as invalid_config, like NUM_SPEC_TOKENS. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Zhehao Hu --- tools/launcher/common/specdec/vllm_smoke_test.sh | 9 +++++++-- tools/launcher/tests/test_vllm_smoke_test.py | 1 + 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/tools/launcher/common/specdec/vllm_smoke_test.sh b/tools/launcher/common/specdec/vllm_smoke_test.sh index 8cd0dd76ee9..8fa2a0db27f 100644 --- a/tools/launcher/common/specdec/vllm_smoke_test.sh +++ b/tools/launcher/common/specdec/vllm_smoke_test.sh @@ -28,6 +28,8 @@ # VLLM_PORT — server port (default: 8000) # REASONING_PARSER — reasoning parser (e.g., "qwen3" for Qwen3.5) # DISABLE_PREFIX_CACHING — set to "1" to disable prefix caching +# SERVER_TIMEOUT — seconds to wait for the server to become healthy (default: 180). Large or +# hybrid targets need longer for compile, kernel JIT and CUDA graph capture # MIN_ACCEPTANCE_LENGTH — fail unless the mean acceptance length (from vLLM's /metrics # counters) reaches this; also fails if no speculative drafts were recorded # SMOKE_JSON — also write the result to this file as JSON, on every exit; "stage" @@ -43,6 +45,7 @@ MODEL=${HF_MODEL_CKPT} DRAFT=${DRAFT_MODEL:-} METHOD=${SPEC_METHOD:-eagle} NUM_SPEC=${NUM_SPEC_TOKENS:-15} +SERVER_TIMEOUT=${SERVER_TIMEOUT:-180} STAGE=starting VERDICT=FAIL REASON="" @@ -113,6 +116,8 @@ if [ -n "${MIN_ACCEPTANCE_LENGTH:-}" ] && ! python3 -c 'import math, sys; sys.ex REASON="MIN_ACCEPTANCE_LENGTH='${MIN_ACCEPTANCE_LENGTH}' is not a finite number" elif ! [[ "$NUM_SPEC" =~ ^[0-9]+$ ]] || [ "$NUM_SPEC" -eq 0 ]; then REASON="NUM_SPEC_TOKENS='${NUM_SPEC}' is not a positive integer" +elif ! [[ "$SERVER_TIMEOUT" =~ ^[0-9]+$ ]] || [ "$SERVER_TIMEOUT" -eq 0 ]; then + REASON="SERVER_TIMEOUT='${SERVER_TIMEOUT}' is not a positive integer" fi if [ -n "$REASON" ]; then STAGE=invalid_config; echo "ERROR: $REASON"; exit 1 @@ -187,7 +192,7 @@ SERVER_PID=$! # Wait for server echo "Waiting for vLLM server..." -for i in $(seq 1 180); do +for i in $(seq 1 "$SERVER_TIMEOUT"); do if curl -s http://localhost:${PORT}/health > /dev/null 2>&1; then echo "Server ready after ${i}s" break @@ -201,7 +206,7 @@ done if ! curl -s http://localhost:${PORT}/health > /dev/null 2>&1; then echo "ERROR: Server timeout" - STAGE=server_timeout; REASON="vLLM server not healthy after 180s"; exit 1 + STAGE=server_timeout; REASON="vLLM server not healthy after ${SERVER_TIMEOUT}s"; exit 1 fi STAGE=serving REASON="exited while sending test prompts" diff --git a/tools/launcher/tests/test_vllm_smoke_test.py b/tools/launcher/tests/test_vllm_smoke_test.py index 9edcd6c22a0..5efde9859e8 100644 --- a/tools/launcher/tests/test_vllm_smoke_test.py +++ b/tools/launcher/tests/test_vllm_smoke_test.py @@ -75,6 +75,7 @@ def test_sources_service_utils_outside_a_launcher_workdir(tmp_path): ({"MIN_ACCEPTANCE_LENGTH": "1,5"}, "MIN_ACCEPTANCE_LENGTH='1,5' is not a finite number"), ({"MIN_ACCEPTANCE_LENGTH": "nan"}, "MIN_ACCEPTANCE_LENGTH='nan' is not a finite number"), ({"NUM_SPEC_TOKENS": "x"}, "NUM_SPEC_TOKENS='x' is not a positive integer"), + ({"SERVER_TIMEOUT": "15m"}, "SERVER_TIMEOUT='15m' is not a positive integer"), ], ) def test_invalid_config_fails_with_a_result(tmp_path, env, reason): From 381e48ae757b16a2c854b29c9c60d6a97e51776a Mon Sep 17 00:00:00 2001 From: Zhehao Hu Date: Fri, 23 Oct 2026 21:33:29 +0000 Subject: [PATCH 4/4] launcher: smoke test rejects zero-padded counts and compares AL unrounded NUM_SPEC_TOKENS=08 passed validation, then went into SPEC_CONFIG as a bare JSON number, which is invalid JSON: vLLM died at startup and the run read as server_died instead of invalid_config. NUM_SPEC_TOKENS and SERVER_TIMEOUT now have to be canonical positive integers. The mean acceptance length was rounded to 4 decimals before the MIN_ACCEPTANCE_LENGTH check, so 1.16666... passed a floor of 1.1667, and the rounded value went into SMOKE_JSON, where gate_smoke repeats the check. The check and SMOKE_JSON now use full precision; only the log is rounded. Signed-off-by: Zhehao Hu --- .../common/specdec/vllm_smoke_test.sh | 20 +++++++++++-------- tools/launcher/tests/test_vllm_smoke_test.py | 2 ++ 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/tools/launcher/common/specdec/vllm_smoke_test.sh b/tools/launcher/common/specdec/vllm_smoke_test.sh index 8fa2a0db27f..437c9066e82 100644 --- a/tools/launcher/common/specdec/vllm_smoke_test.sh +++ b/tools/launcher/common/specdec/vllm_smoke_test.sh @@ -114,9 +114,10 @@ trap cleanup EXIT # Checked after the trap, so a bad value still writes this run's SMOKE_JSON. if [ -n "${MIN_ACCEPTANCE_LENGTH:-}" ] && ! python3 -c 'import math, sys; sys.exit(not math.isfinite(float(sys.argv[1])))' "$MIN_ACCEPTANCE_LENGTH" 2>/dev/null; then REASON="MIN_ACCEPTANCE_LENGTH='${MIN_ACCEPTANCE_LENGTH}' is not a finite number" -elif ! [[ "$NUM_SPEC" =~ ^[0-9]+$ ]] || [ "$NUM_SPEC" -eq 0 ]; then +# Canonical form only: NUM_SPEC goes into SPEC_CONFIG as a bare JSON number, so "08" breaks it. +elif ! [[ "$NUM_SPEC" =~ ^[1-9][0-9]*$ ]]; then REASON="NUM_SPEC_TOKENS='${NUM_SPEC}' is not a positive integer" -elif ! [[ "$SERVER_TIMEOUT" =~ ^[0-9]+$ ]] || [ "$SERVER_TIMEOUT" -eq 0 ]; then +elif ! [[ "$SERVER_TIMEOUT" =~ ^[1-9][0-9]*$ ]]; then REASON="SERVER_TIMEOUT='${SERVER_TIMEOUT}' is not a positive integer" fi if [ -n "$REASON" ]; then @@ -269,8 +270,9 @@ fi # Mean acceptance length over all prompts, from vLLM's cumulative Prometheus counters: # 1 + accepted draft tokens / drafts (the bonus token counts, as in vLLM's own log). The # "SpecDecoding metrics: Mean acceptance length" log line covers only the last logging -# interval (vLLM resets its stats after each log), so it is not used. -AVG_ACCEPT=$(echo "$ALL_METRICS" | python3 -c ' +# interval (vLLM resets its stats after each log), so it is not used. Printed twice: full +# precision for the floor check and SMOKE_JSON, rounded for the log. +AL_OUT=$(echo "$ALL_METRICS" | python3 -c ' import sys drafts = accepted = 0.0 @@ -281,8 +283,10 @@ for line in sys.stdin: elif name == "vllm:spec_decode_num_accepted_tokens_total": accepted += float(line.split()[-1]) if drafts > 0: - print(f"{1 + accepted / drafts:.4f}") + al = 1 + accepted / drafts + print(repr(al), f"{al:.4f}") ' 2>/dev/null || true) +read -r AVG_ACCEPT AL_SHOWN <<< "$AL_OUT" STAGE=completed VERDICT=PASS @@ -293,17 +297,17 @@ if [ $PROMPTS_FAILED -gt 0 ]; then elif [ -n "${MIN_ACCEPTANCE_LENGTH:-}" ]; then echo "" echo "=== Acceptance Length Regression Check ===" - echo " Mean acceptance length: ${AVG_ACCEPT:-not reported}" + echo " Mean acceptance length: ${AL_SHOWN:-not reported}" echo " Threshold: ${MIN_ACCEPTANCE_LENGTH}" if [ -z "$AVG_ACCEPT" ]; then # A threshold that was never evaluated must not pass. VERDICT=FAIL REASON="REGRESSION: MIN_ACCEPTANCE_LENGTH=${MIN_ACCEPTANCE_LENGTH} was set but vLLM /metrics reports no speculative drafts" elif python3 -c 'import sys; sys.exit(float(sys.argv[1]) < float(sys.argv[2]))' "$AVG_ACCEPT" "$MIN_ACCEPTANCE_LENGTH"; then - echo " PASS: ${AVG_ACCEPT} >= ${MIN_ACCEPTANCE_LENGTH}" + echo " PASS: ${AL_SHOWN} >= ${MIN_ACCEPTANCE_LENGTH}" else VERDICT=FAIL - REASON="REGRESSION: ${AVG_ACCEPT} < ${MIN_ACCEPTANCE_LENGTH}" + REASON="REGRESSION: ${AL_SHOWN} < ${MIN_ACCEPTANCE_LENGTH}" fi fi diff --git a/tools/launcher/tests/test_vllm_smoke_test.py b/tools/launcher/tests/test_vllm_smoke_test.py index 5efde9859e8..5dc3d21a480 100644 --- a/tools/launcher/tests/test_vllm_smoke_test.py +++ b/tools/launcher/tests/test_vllm_smoke_test.py @@ -75,6 +75,8 @@ def test_sources_service_utils_outside_a_launcher_workdir(tmp_path): ({"MIN_ACCEPTANCE_LENGTH": "1,5"}, "MIN_ACCEPTANCE_LENGTH='1,5' is not a finite number"), ({"MIN_ACCEPTANCE_LENGTH": "nan"}, "MIN_ACCEPTANCE_LENGTH='nan' is not a finite number"), ({"NUM_SPEC_TOKENS": "x"}, "NUM_SPEC_TOKENS='x' is not a positive integer"), + # Not valid as the bare JSON number SPEC_CONFIG passes to vLLM. + ({"NUM_SPEC_TOKENS": "08"}, "NUM_SPEC_TOKENS='08' is not a positive integer"), ({"SERVER_TIMEOUT": "15m"}, "SERVER_TIMEOUT='15m' is not a positive integer"), ], )