[2/8] launcher: harden the vLLM smoke test and newest-checkpoint selection

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

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

tools/launcher/common/specdec/quantize_drafter.sh

@@ -30,15 +30,14 @@ trap 'error_handler $0 $LINENO' ERR
 DRAFTER="${DRAFTER_CKPT}"
 # Training writes exported-checkpoint-<step>/ 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
 

tools/launcher/common/specdec/vllm_smoke_test.sh

@@ -28,21 +28,109 @@
 #   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"
+#                    records how far the run got (invalid_config, 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}
+SERVER_TIMEOUT=${SERVER_TIMEOUT:-180}
+STAGE=starting
+VERDICT=FAIL
+REASON=""
+PROMPTS_OK=0
+PROMPTS_FAILED=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" \
+        "$PROMPTS_OK" "$PROMPTS_FAILED" "$TOTAL_TOKENS" "$AVG_TPS" "${AVG_ACCEPT:-}" "${MIN_ACCEPTANCE_LENGTH:-}" <<'EOF'
+import json
+import math
+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:
+        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 = {
+    "stage": stage,
+    "pass": verdict == "PASS",
+    "reason": reason or None,
+    "method": method,
+    "target": target,
+    "draft": draft or None,
+    "num_speculative_tokens": count(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
+    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
+
+# 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"
+# 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" =~ ^[1-9][0-9]*$ ]]; then
+    REASON="SERVER_TIMEOUT='${SERVER_TIMEOUT}' 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
-    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 +140,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}
 
@@ -105,29 +193,31 @@ 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
     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 ${SERVER_TIMEOUT}s"; 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,62 +235,85 @@ 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))
+        PROMPTS_OK=$((PROMPTS_OK + 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)"
-        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=$(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. 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
+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:
+    al = 1 + accepted / drafts
+    print(repr(al), f"{al:.4f}")
+' 2>/dev/null || true)
+read -r AVG_ACCEPT AL_SHOWN <<< "$AL_OUT"
 
-# 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
+STAGE=completed
+VERDICT=PASS
+REASON=""
+if [ $PROMPTS_FAILED -gt 0 ]; then
+    VERDICT=FAIL
+    REASON="${PROMPTS_FAILED} of $((PROMPTS_OK + PROMPTS_FAILED)) prompts failed"
+elif [ -n "${MIN_ACCEPTANCE_LENGTH:-}" ]; then
+    echo ""
+    echo "=== Acceptance Length Regression Check ==="
+    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: ${AL_SHOWN} >= ${MIN_ACCEPTANCE_LENGTH}"
     else
-        echo "WARNING: Could not parse acceptance length from vLLM log, skipping regression check"
+        VERDICT=FAIL
+        REASON="REGRESSION: ${AL_SHOWN} < ${MIN_ACCEPTANCE_LENGTH}"
     fi
 fi
 
+if [ "$VERDICT" = FAIL ]; then
+    echo "  $REASON"
+    exit 1
+fi
+
 echo "Done"