#!/usr/bin/env bash

# Compatibility gate only: this is not a simulator-fidelity measurement.

set -euo pipefail

OUT="${OUTPUT_ROOT:?OUTPUT_ROOT is required}"
TP="${TP:?TP is required}"
MNS="${MNS:?MNS is required}"
MBT="${MBT:?MBT is required}"
SERVER_PORT="${SERVER_PORT:?SERVER_PORT is required}"
ENABLE_EXPERT_PARALLEL="${ENABLE_EXPERT_PARALLEL:?ENABLE_EXPERT_PARALLEL 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-235B-A22B-FP8}"
CLIENT="${CLIENT:?CLIENT is required}"
SERVED_MODEL="qwen3-235b-v020-smoke"
SERVER_PID=""

mkdir -p "${OUT}/logs" "${OUT}/provenance" "${OUT}/results"
exec > >(tee -a "${OUT}/logs/controller.log") 2>&1

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
  SERVER_PID=""
}
trap cleanup EXIT INT TERM

IFS=',' read -r -a GPU_IDS <<< "${CUDA_VISIBLE_DEVICES:?GPU allocation is required}"
if [[ "${#GPU_IDS[@]}" -ne "${TP}" ]]; then
  echo "ERROR: expected TP=${TP} GPUs, got ${CUDA_VISIBLE_DEVICES}" >&2
  exit 1
fi
case "${ENABLE_EXPERT_PARALLEL}" in
  true) EP_FLAG=(--enable-expert-parallel) ;;
  false) EP_FLAG=() ;;
  *) echo "ERROR: ENABLE_EXPERT_PARALLEL must be true or false" >&2; exit 1 ;;
esac

echo "QWEN235_V020_SMOKE_LAUNCH_ECHO host=$(hostname) gpus=${CUDA_VISIBLE_DEVICES} model=${MODEL_ROOT} runtime=vLLM-0.20.0+cu129 checkpoint=FP8_compute_kv=auto config=TP${TP}_EP${ENABLE_EXPERT_PARALLEL}_MNS${MNS}_MBT${MBT} prefix=false requests=1 shape=2048x1 role=compatibility_gate_not_latency_measurement output=${OUT} expected_wall=8-20m expected_gpu_cap=$((${TP} * 1))_H20h"
date -u +START_UTC=%Y-%m-%dT%H:%M:%SZ
sha256sum "${BASH_SOURCE[0]}" "${CLIENT}" "${MODEL_ROOT}/config.json" > "${OUT}/provenance/input.sha256"
"${VENV_ROOT}/bin/vllm" --version > "${OUT}/provenance/vllm.version"
"${VENV_ROOT}/bin/python" - <<'PY' > "${OUT}/provenance/runtime.json"
import importlib.metadata as metadata
import json
import vllm
print(json.dumps({"vllm": metadata.version("vllm"), "vllm_file": vllm.__file__}, sort_keys=True))
PY

export TOKENIZERS_PARALLELISM=false
export VLLM_USE_V1=1
export TORCH_CUDA_ARCH_LIST=9.0
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
export HOME=/tmp/wjh
export XDG_CACHE_HOME=/tmp/wjh/.cache
export VLLM_CACHE_ROOT=/tmp/wjh/.cache/vllm
setsid "${VENV_ROOT}/bin/vllm" serve "${MODEL_ROOT}" \
  --host 127.0.0.1 --port "${SERVER_PORT}" --served-model-name "${SERVED_MODEL}" \
  --tensor-parallel-size "${TP}" --disable-custom-all-reduce --quantization fp8 \
  --gpu-memory-utilization 0.80 --kv-cache-dtype auto --max-model-len 40960 \
  --max-num-batched-tokens "${MBT}" --max-num-seqs "${MNS}" \
  --no-enable-prefix-caching --enable-chunked-prefill --no-enable-log-requests \
  "${EP_FLAG[@]}" > "${OUT}/logs/server.log" 2>&1 &
SERVER_PID=$!
READY=0
for _ in $(seq 1 180); do
  if curl -fsS --max-time 2 "http://127.0.0.1:${SERVER_PORT}/v1/models" > "${OUT}/results/models.json" 2>/dev/null; then
    READY=1
    break
  fi
  if ! kill -0 "${SERVER_PID}" 2>/dev/null; then
    tail -200 "${OUT}/logs/server.log"
    exit 1
  fi
  sleep 3
done
if [[ "${READY}" -ne 1 ]]; then
  tail -200 "${OUT}/logs/server.log"
  exit 1
fi

"${VENV_ROOT}/bin/python" "${CLIENT}" --port "${SERVER_PORT}" --served-model "${SERVED_MODEL}" \
  --model-path "${MODEL_ROOT}" --rate 1 --requests 1 --input-tokens 2048 --output-tokens 1 \
  --output "${OUT}/results/one_request.json"
grep -E -i 'quant|fp8|gpu blocks|num_gpu_blocks|kv cache|expert.parallel|expert parallel' "${OUT}/logs/server.log" > "${OUT}/provenance/resolved-runtime.log" || true
cleanup
find "${OUT}" -type f ! -path '*/provenance/artifacts.sha256' -print0 | sort -z | xargs -0 sha256sum > "${OUT}/provenance/artifacts.sha256"
date -u +END_UTC=%Y-%m-%dT%H:%M:%SZ
echo 'QWEN235_V020_SMOKE_COMPLETE'