From 3ee29432a5529434e41de81ac844bf30d33e55fa Mon Sep 17 00:00:00 2001
From: Gahow Wang <gahow.wang@gmail.com>
Date: Sat, 18 Jul 2026 01:11:08 +0800
Subject: [PATCH] Allow fixed runtime preflight resume

---
 .../run_qwen30_fixed_runtime_preflight.sh       | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/runs/frontier-fidelity-envelope-v1/run_qwen30_fixed_runtime_preflight.sh b/runs/frontier-fidelity-envelope-v1/run_qwen30_fixed_runtime_preflight.sh
index a64da12..30e9a59 100644
--- a/runs/frontier-fidelity-envelope-v1/run_qwen30_fixed_runtime_preflight.sh
+++ b/runs/frontier-fidelity-envelope-v1/run_qwen30_fixed_runtime_preflight.sh
@@ -9,6 +9,7 @@ OUT="${1:?output root is required}"
 VENV_ROOT="${VENV_ROOT:-/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1}"
 MODEL_ROOT="${MODEL_ROOT:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B}"
 FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE:-${OUT}/flashinfer-shared-workspace}"
+TPS="${TPS:-4 2 1}"
 PORT=8600
 
 mkdir -p "${OUT}/runs" "${OUT}/provenance" "${FLASHINFER_WORKSPACE_BASE}"
@@ -19,7 +20,17 @@ date -u +START_UTC=%Y-%m-%dT%H:%M:%SZ
 sha256sum "${BASH_SOURCE[0]}" "${MODEL_ROOT}/config.json" > "${OUT}/provenance/input.sha256"
 "${VENV_ROOT}/bin/vllm" --version > "${OUT}/provenance/vllm.version"
 
-if ! nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | awk '$1 != 0 {exit 1}'; then
+wait_for_idle() {
+  for _ in $(seq 1 30); do
+    if nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | awk '$1 != 0 {exit 1}'; then
+      return 0
+    fi
+    sleep 2
+  done
+  return 1
+}
+
+if ! wait_for_idle; then
   echo 'ERROR: a GPU is not idle before runtime preflight' >&2
   exit 1
 fi
@@ -101,13 +112,13 @@ wait_wave() {
     echo 'ERROR: one or more runtime preflight servers failed' >&2
     exit 1
   fi
-  if ! nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | awk '$1 != 0 {exit 1}'; then
+  if ! wait_for_idle; then
     echo 'ERROR: a GPU remains allocated after runtime preflight wave' >&2
     exit 1
   fi
 }
 
-for tp in 4 2 1; do
+for tp in ${TPS}; do
   echo "PREFLIGHT_WAVE_START tp=${tp}"
   case "${tp}" in
     4)