#!/bin/bash
# Stage 2: engineering rung 1 = CUDA graphs + piecewise compile (drop enforce-eager)
# Chain T (GPU2): 8B twin draft.  smoke(gate) -> eager refs -> graph ladder
# Chain F (GPU3): 4B blind draft. smoke(gate) -> graph ladder, eager fallback on gate fail
S=/testessfs10/users/zeyu.zhang/wangyu_ssd
ST=$S/stage0
PY=$S/vllm_env/bin/python
cd $ST
export PYTHONPATH=$ST PYTHONDONTWRITEBYTECODE=1
export PATH=$S/vllm_env/bin:$PATH
export HF_HOME=$S/hf_cache HF_ENDPOINT=https://hf-mirror.com
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
M8=Qwen/Qwen3-VL-8B-Thinking
D8=cyankiwi/Qwen3-VL-8B-Thinking-AWQ-4bit
D4=cyankiwi/Qwen3-VL-4B-Thinking-AWQ-8bit
COMMON="--parquet testmini.parquet --pids-from res05_qwen_thinking.jsonl --n 20 --max-tokens 1024 --gpu-mem-util 0.60 --max-model-len 3584"
SMOKEC="--parquet testmini.parquet --pids-from res05_qwen_thinking.jsonl --n 4 --max-tokens 256 --gpu-mem-util 0.60 --max-model-len 3584"
B=$S/stage1
rm -f $B/og8_*.jsonl $B/og4_*.jsonl $B/DONE_og8 $B/DONE_og4 $B/FAIL_og8 $B/FAIL_og4

gate() { # $1=log $2=threshold(percent)
  A=$(grep -o "Draft acceptance rate: [0-9.]*" "$1" | tail -1 | sed "s/.*: //")
  awk -v a="${A:-0}" -v t="$2" "BEGIN{exit (a>=t)?0:1}"
}

nohup bash -c "
CUDA_VISIBLE_DEVICES=2 $PY stage1_spec_smoke.py --mode spec --model $M8 --draft $D8 --gamma 6 $SMOKEC --out $B/og8_smoke.jsonl > $B/logs/og8_smoke.log 2>&1
$(declare -f gate)
if gate $B/logs/og8_smoke.log 50; then GOK=1; else GOK=0; touch $B/FAIL_og8; fi
CUDA_VISIBLE_DEVICES=2 $PY stage1_spec_smoke.py --mode vanilla --model $M8 $COMMON --enforce-eager --out $B/og8_evan.jsonl > $B/logs/og8_evan.log 2>&1
CUDA_VISIBLE_DEVICES=2 $PY stage1_spec_smoke.py --mode spec --model $M8 --draft $D8 --gamma 6 $COMMON --enforce-eager --out $B/og8_estrict_g6.jsonl > $B/logs/og8_estrict_g6.log 2>&1
if [ \$GOK = 1 ]; then
  CUDA_VISIBLE_DEVICES=2 $PY stage1_spec_smoke.py --mode vanilla --model $M8 $COMMON --out $B/og8_van1.jsonl > $B/logs/og8_van1.log 2>&1
  CUDA_VISIBLE_DEVICES=2 $PY stage1_spec_smoke.py --mode spec --model $M8 --draft $D8 --gamma 6 $COMMON --out $B/og8_strict_g6.jsonl > $B/logs/og8_strict_g6.log 2>&1
  VLLM_SPEC_RELAX_LOGTAU=0.693 CUDA_VISIBLE_DEVICES=2 $PY stage1_spec_smoke.py --mode spec --model $M8 --draft $D8 --gamma 6 $COMMON --out $B/og8_relax_g6.jsonl > $B/logs/og8_relax_g6.log 2>&1
  CUDA_VISIBLE_DEVICES=2 $PY stage1_spec_smoke.py --mode vanilla --model $M8 $COMMON --out $B/og8_vanm.jsonl > $B/logs/og8_vanm.log 2>&1
  CUDA_VISIBLE_DEVICES=2 $PY stage1_spec_smoke.py --mode spec --model $M8 --draft $D8 --gamma 8 $COMMON --out $B/og8_strict_g8.jsonl > $B/logs/og8_strict_g8.log 2>&1
  VLLM_SPEC_RELAX_LOGTAU=0.693 CUDA_VISIBLE_DEVICES=2 $PY stage1_spec_smoke.py --mode spec --model $M8 --draft $D8 --gamma 8 $COMMON --out $B/og8_relax_g8.jsonl > $B/logs/og8_relax_g8.log 2>&1
  CUDA_VISIBLE_DEVICES=2 $PY stage1_spec_smoke.py --mode vanilla --model $M8 $COMMON --out $B/og8_van2.jsonl > $B/logs/og8_van2.log 2>&1
fi
touch $B/DONE_og8
" > /dev/null 2>&1 &
disown

nohup bash -c "
sleep 240
VLLM_SPEC_DRAFT_BLIND_MM=1 CUDA_VISIBLE_DEVICES=3 $PY stage1_spec_smoke.py --mode spec --model $M8 --draft $D4 --gamma 6 $SMOKEC --out $B/og4_smoke.jsonl > $B/logs/og4_smoke.log 2>&1
$(declare -f gate)
if gate $B/logs/og4_smoke.log 40; then
  CUDA_VISIBLE_DEVICES=3 $PY stage1_spec_smoke.py --mode vanilla --model $M8 $COMMON --out $B/og4_van1.jsonl > $B/logs/og4_van1.log 2>&1
  VLLM_SPEC_DRAFT_BLIND_MM=1 CUDA_VISIBLE_DEVICES=3 $PY stage1_spec_smoke.py --mode spec --model $M8 --draft $D4 --gamma 6 $COMMON --out $B/og4_strict_g6.jsonl > $B/logs/og4_strict_g6.log 2>&1
  VLLM_SPEC_DRAFT_BLIND_MM=1 VLLM_SPEC_RELAX_LOGTAU=0.693 CUDA_VISIBLE_DEVICES=3 $PY stage1_spec_smoke.py --mode spec --model $M8 --draft $D4 --gamma 6 $COMMON --out $B/og4_relax_g6.jsonl > $B/logs/og4_relax_g6.log 2>&1
  CUDA_VISIBLE_DEVICES=3 $PY stage1_spec_smoke.py --mode vanilla --model $M8 $COMMON --out $B/og4_vanm.jsonl > $B/logs/og4_vanm.log 2>&1
  VLLM_SPEC_DRAFT_BLIND_MM=1 CUDA_VISIBLE_DEVICES=3 $PY stage1_spec_smoke.py --mode spec --model $M8 --draft $D4 --gamma 8 $COMMON --out $B/og4_strict_g8.jsonl > $B/logs/og4_strict_g8.log 2>&1
  VLLM_SPEC_DRAFT_BLIND_MM=1 VLLM_SPEC_RELAX_LOGTAU=0.693 CUDA_VISIBLE_DEVICES=3 $PY stage1_spec_smoke.py --mode spec --model $M8 --draft $D4 --gamma 8 $COMMON --out $B/og4_relax_g8.jsonl > $B/logs/og4_relax_g8.log 2>&1
  CUDA_VISIBLE_DEVICES=3 $PY stage1_spec_smoke.py --mode vanilla --model $M8 $COMMON --out $B/og4_van2.jsonl > $B/logs/og4_van2.log 2>&1
else
  touch $B/FAIL_og4
  CUDA_VISIBLE_DEVICES=3 $PY stage1_spec_smoke.py --mode vanilla --model $M8 $COMMON --enforce-eager --out $B/og4_van1.jsonl > $B/logs/og4_van1.log 2>&1
  VLLM_SPEC_DRAFT_BLIND_MM=1 CUDA_VISIBLE_DEVICES=3 $PY stage1_spec_smoke.py --mode spec --model $M8 --draft $D4 --gamma 6 $COMMON --enforce-eager --out $B/og4_strict_g6.jsonl > $B/logs/og4_strict_g6.log 2>&1
  VLLM_SPEC_DRAFT_BLIND_MM=1 VLLM_SPEC_RELAX_LOGTAU=0.693 CUDA_VISIBLE_DEVICES=3 $PY stage1_spec_smoke.py --mode spec --model $M8 --draft $D4 --gamma 6 $COMMON --enforce-eager --out $B/og4_relax_g6.jsonl > $B/logs/og4_relax_g6.log 2>&1
  CUDA_VISIBLE_DEVICES=3 $PY stage1_spec_smoke.py --mode vanilla --model $M8 $COMMON --enforce-eager --out $B/og4_van2.jsonl > $B/logs/og4_van2.log 2>&1
fi
touch $B/DONE_og4
" > /dev/null 2>&1 &
disown
echo LAUNCHED_STAGE2