#!/usr/bin/env bash

set -euo pipefail

PORT="${1:-8123}"
CONTROL_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${CONTROL_ROOT}/../.." && pwd)"
VENV_ROOT="${VENV_ROOT:-/home/admin/cpfs/wjh/venvs/vllm-0.20.0-cu129-workload-regime-v2}"
MODEL_ROOT="${MODEL_ROOT:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B}"
OUTPUT_ROOT="${OUTPUT_ROOT:-${CONTROL_ROOT}/outputs/max-model-len-smoke}"
SERVED_MODEL="qwen3-30b-code-maxlen-smoke"
MAX_MODEL_LEN=147456
SERVER_PID=""

[[ "${ALLOW_LONG_CONTEXT_SERVER:-false}" == "true" ]] || {
  echo "ERROR: set ALLOW_LONG_CONTEXT_SERVER=true to acknowledge the model's native 40960 limit" >&2
  exit 1
}
[[ "${CUDA_VISIBLE_DEVICES:-}" == "0,1,2,3" ]] || {
  echo "ERROR: this TP4 smoke requires CUDA_VISIBLE_DEVICES=0,1,2,3" >&2
  exit 1
}
for path in \
  "${VENV_ROOT}/bin/python" \
  "${VENV_ROOT}/bin/vllm" \
  "${MODEL_ROOT}/config.json" \
  "${REPO_ROOT}/runs/frontier-s3-real-v0/qwen30_prefill_client.py"; do
  [[ -e "${path}" ]] || { echo "ERROR: missing ${path}" >&2; exit 1; }
done
[[ ! -e "${OUTPUT_ROOT}" ]] || {
  echo "ERROR: refusing to overwrite ${OUTPUT_ROOT}" >&2
  exit 1
}
mkdir -p "${OUTPUT_ROOT}/logs" "${OUTPUT_ROOT}/results" "${OUTPUT_ROOT}/provenance"

cleanup() {
  if [[ -n "${SERVER_PID}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
    kill -TERM -- "-${SERVER_PID}" 2>/dev/null || true
    for _ in $(seq 1 30); do
      kill -0 "${SERVER_PID}" 2>/dev/null || break
      sleep 1
    done
    kill -KILL -- "-${SERVER_PID}" 2>/dev/null || true
  fi
}
trap cleanup EXIT INT TERM

nvidia-smi \
  --query-gpu=index,name,memory.used,utilization.gpu,ecc.errors.uncorrected.aggregate.total \
  --format=csv,noheader,nounits > "${OUTPUT_ROOT}/provenance/gpus.before.csv"
[[ "$(wc -l < "${OUTPUT_ROOT}/provenance/gpus.before.csv")" -eq 8 ]] || {
  echo "ERROR: expected eight GPUs on host" >&2
  exit 1
}
awk -F, '
  $2 !~ /NVIDIA H20/ || $3 + 0 != 0 || $4 + 0 != 0 || $5 + 0 != 0 { bad = 1 }
  END { exit bad }
' "${OUTPUT_ROOT}/provenance/gpus.before.csv" || {
  echo "ERROR: all eight H20 GPUs must be idle and healthy before the smoke" >&2
  exit 1
}
if nvidia-smi --query-compute-apps=pid --format=csv,noheader,nounits \
  | grep -Eq '^[[:space:]]*[0-9]+'; then
  echo "ERROR: compute process appeared after the idle probe" >&2
  exit 1
fi

export TOKENIZERS_PARALLELISM=false
export VLLM_USE_V1=1
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
export FLASHINFER_WORKSPACE_BASE="/tmp/wjh/frontier-code-maxlen-smoke"
mkdir -p "${FLASHINFER_WORKSPACE_BASE}"
ulimit -n 65536

setsid "${VENV_ROOT}/bin/vllm" serve "${MODEL_ROOT}" \
  --host 127.0.0.1 \
  --port "${PORT}" \
  --served-model-name "${SERVED_MODEL}" \
  --tensor-parallel-size 4 \
  --gpu-memory-utilization 0.92 \
  --max-model-len "${MAX_MODEL_LEN}" \
  --max-num-batched-tokens 8192 \
  --max-num-seqs 16 \
  --enable-prefix-caching \
  --block-size 16 \
  --enable-chunked-prefill \
  --no-enable-log-requests \
  > "${OUTPUT_ROOT}/logs/server.log" 2>&1 &
SERVER_PID=$!

READY=0
for _ in $(seq 1 900); do
  if curl -fsS --max-time 2 "http://127.0.0.1:${PORT}/v1/models" \
    > "${OUTPUT_ROOT}/results/models.json" 2>/dev/null; then
    READY=1
    break
  fi
  kill -0 "${SERVER_PID}" 2>/dev/null || {
    tail -200 "${OUTPUT_ROOT}/logs/server.log"
    exit 1
  }
  sleep 3
done
[[ "${READY}" -eq 1 ]] || {
  echo "ERROR: server readiness timeout" >&2
  exit 1
}

run_shape() {
  local label="$1"
  local input_tokens="$2"
  local output_tokens="$3"
  "${VENV_ROOT}/bin/python" \
    "${REPO_ROOT}/runs/frontier-s3-real-v0/qwen30_prefill_client.py" \
    --port "${PORT}" \
    --served-model "${SERVED_MODEL}" \
    --model-path "${MODEL_ROOT}" \
    --rate 1 \
    --requests 1 \
    --input-tokens "${input_tokens}" \
    --output-tokens "${output_tokens}" \
    --timeout-seconds 1800 \
    --output "${OUTPUT_ROOT}/results/${label}.json"
}

run_shape warmup 512 1
run_shape code-p50 20051 78
run_shape code-candidate-p99 119702 68
run_shape code-candidate-max 136774 242

cleanup
SERVER_PID=""
nvidia-smi --query-gpu=index,name,memory.used,utilization.gpu \
  --format=csv,noheader,nounits > "${OUTPUT_ROOT}/provenance/gpus.after.csv"
sha256sum \
  "${BASH_SOURCE[0]}" \
  "${REPO_ROOT}/runs/frontier-s3-real-v0/qwen30_prefill_client.py" \
  "${MODEL_ROOT}/config.json" \
  > "${OUTPUT_ROOT}/provenance/sources.sha256"
echo "MAX_MODEL_LEN_SMOKE_COMPLETE output=${OUTPUT_ROOT}"