#!/usr/bin/env bash
set -u

MAX_BATCHED_TOKENS="${1:?usage: run_batched_tokens_experiment.sh MAX_BATCHED_TOKENS LABEL}"
LABEL="${2:?usage: run_batched_tokens_experiment.sh MAX_BATCHED_TOKENS LABEL}"

OUT_DIR="/root/work/logs/${LABEL}"
MODEL_PATH="/root/public-storage/models/Qwen/Qwen3.6-35B-A3B"
DATASET="/root/work/logs/synthetic_cumulative_8_s3_cap40k.jsonl"
MAX_REQUESTS="${MAX_REQUESTS:-8}"
MAX_TOKENS="${MAX_TOKENS:-256}"
RESULT_JSON="${OUT_DIR}/eager_100k_b${MAX_BATCHED_TOKENS}_c1_r${MAX_REQUESTS}_t${MAX_TOKENS}_cap40k.json"
SERVER_LOG="${OUT_DIR}/server.log"
DRIVER_LOG="${OUT_DIR}/driver.log"

mkdir -p "${OUT_DIR}"
exec > >(tee -a "${DRIVER_LOG}") 2>&1

echo "[exp] label=${LABEL}"
echo "[exp] max_num_batched_tokens=${MAX_BATCHED_TOKENS}"
echo "[exp] out_dir=${OUT_DIR}"
date

echo "[exp] stopping existing api_server"
PIDS="$(pgrep -f 'vllm.entrypoints.openai.api_server' || true)"
if [ -n "${PIDS}" ]; then
  kill -TERM ${PIDS} 2>/dev/null || true
  sleep 8
  kill -KILL ${PIDS} 2>/dev/null || true
fi

export PYTHONPATH=/usr/local/corex/lib64/python3/dist-packages
export LD_LIBRARY_PATH=/usr/local/corex/lib64:/usr/local/iluvatar/lib64:/usr/local/openmpi/lib
export PATH=/usr/local/corex/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/openmpi/bin
export MOE_NATIVE=1
export CHUNK_PARALLEL=1
export PREFIX_FLASH=1
export RMSNORM_NATIVE=1
export LOOP1_NATIVE=1
export PROF_TRACE="${PROF_TRACE:-0}"
export PROF_MOE_SUMMARY="${PROF_MOE_SUMMARY:-0}"
export PROF_MOE_SYNC="${PROF_MOE_SYNC:-1}"
export PROF_MOE_INTERVAL="${PROF_MOE_INTERVAL:-400}"
export VLLM_ENGINE_ITERATION_TIMEOUT_S=3600

echo "[exp] gpu before start"
LD_LIBRARY_PATH=/usr/local/corex-3.2.3/lib64:/usr/local/corex/lib64:/usr/local/corex/lib:/usr/local/iluvatar/lib64 \
  timeout 15 /usr/local/corex/bin/ixsmi || true

echo "[exp] starting server"
cd /root || exit 1
nohup python3 -m vllm.entrypoints.openai.api_server \
  --model "${MODEL_PATH}" \
  --served-model-name llm \
  --host 0.0.0.0 \
  --port 1111 \
  --max-model-len 100000 \
  --enforce-eager \
  --gpu-memory-utilization 0.9 \
  --trust-remote-code \
  -tp 4 \
  --max-num-seqs 1 \
  --disable-log-requests \
  --disable-frontend-multiprocessing \
  --max-num-batched-tokens "${MAX_BATCHED_TOKENS}" \
  --enable-chunked-prefill \
  --max-seq-len-to-capture 32768 \
  --enable-auto-tool-choice \
  --tool-call-parser qwen3_coder \
  --reasoning-parser qwen3 \
  --enable-prefix-caching \
  --chat-template /workspace/chat_template_multi_system.jinja \
  > "${SERVER_LOG}" 2>&1 &
SERVER_PID="$!"
echo "[exp] server_pid=${SERVER_PID}"

echo "[exp] waiting for health"
for i in $(seq 1 90); do
  if python3 -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:1111/health', timeout=3).read()" >/dev/null 2>&1; then
    echo "[exp] health ok after ${i} checks"
    break
  fi
  if ! kill -0 "${SERVER_PID}" 2>/dev/null; then
    echo "[exp] server exited before health"
    tail -120 "${SERVER_LOG}" || true
    exit 2
  fi
  sleep 5
done

if ! python3 -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:1111/health', timeout=5).read()" >/dev/null 2>&1; then
  echo "[exp] health failed"
  tail -120 "${SERVER_LOG}" || true
  exit 3
fi

echo "[exp] running benchmark"
cd /root/work || exit 1
python3 /root/work/formal_perf_bench.py \
  --url http://127.0.0.1:1111 \
  --model llm \
  --dataset "${DATASET}" \
  --concurrency 1 \
  --max-requests "${MAX_REQUESTS}" \
  --max-tokens "${MAX_TOKENS}" \
  --timeout 1800 \
  --out "${RESULT_JSON}"
BENCH_RC="$?"

echo "[exp] benchmark rc=${BENCH_RC}"
cat "${RESULT_JSON}" || true
echo "[exp] gpu after benchmark"
LD_LIBRARY_PATH=/usr/local/corex-3.2.3/lib64:/usr/local/corex/lib64:/usr/local/corex/lib:/usr/local/iluvatar/lib64 \
  timeout 15 /usr/local/corex/bin/ixsmi || true

exit "${BENCH_RC}"