#!/usr/bin/env bash

set -euo pipefail

OUTPUT_ROOT="${OUTPUT_ROOT:-$(pwd)/artifacts/t0-smoke-20260716}"
VENV_ROOT="${VENV_ROOT:-/tmp/wjh-frontier-vllm0102-smoke/.venv}"
MODEL_ROOT="${MODEL_ROOT:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-235B-A22B-FP8}"
FRONTIER_ROOT="${FRONTIER_ROOT:-/home/admin/cpfs/wjh/frontier-community-qwen235-smoke-20260715/Frontier-d9cfeb6}"
SERVER_PORT="${SERVER_PORT:-18910}"
SERVED_MODEL="qwen3-235b-t0-smoke"
LOG_DIR="${OUTPUT_ROOT}/logs"
RESULT_DIR="${OUTPUT_ROOT}/results"
PROVENANCE_DIR="${OUTPUT_ROOT}/provenance"
SERVER_PID=""

mkdir -p "${LOG_DIR}" "${RESULT_DIR}" "${PROVENANCE_DIR}"
exec > >(tee -a "${LOG_DIR}/smoke.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
      if ! kill -0 "${SERVER_PID}" 2>/dev/null; then
        break
      fi
      sleep 1
    done
    kill -KILL -- "-${SERVER_PID}" 2>/dev/null || true
  fi
}
trap cleanup EXIT INT TERM

if [[ -z "${CUDA_VISIBLE_DEVICES:-}" ]]; then
  echo "ERROR: CUDA_VISIBLE_DEVICES must contain the four fleet-allocated GPUs" >&2
  exit 1
fi
IFS=',' read -r -a GPU_IDS <<< "${CUDA_VISIBLE_DEVICES}"
if [[ "${#GPU_IDS[@]}" -ne 4 ]]; then
  echo "ERROR: expected exactly four GPUs, got ${CUDA_VISIBLE_DEVICES}" >&2
  exit 1
fi

echo "LAUNCH_ECHO host=$(hostname) gpus=${CUDA_VISIBLE_DEVICES} model=${MODEL_ROOT} trace=fixed_ISL2048_OSL128 prefix=off qps=single_then_concurrency2 runtime=community_vllm_0.10.2 topology=TP4_DP1 execution=eager kv=BF16 spec=off cuda_graph=off output=${OUTPUT_ROOT} expected_wall=20-30m hard_wall=1800s hard_gpu_cap=2_H20h"
date -u +"START_UTC=%Y-%m-%dT%H:%M:%SZ"
nvidia-smi --query-gpu=index,name,memory.used,utilization.gpu --format=csv,noheader

test -x "${VENV_ROOT}/bin/vllm"
test -x "${VENV_ROOT}/bin/python"
test -f "${MODEL_ROOT}/config.json"
test -f "${FRONTIER_ROOT}/pyproject.toml"
test -f "$(pwd)/t0_smoke_client.py"

sha256sum run_t0_smoke.sh t0_smoke_client.py > "${PROVENANCE_DIR}/source.sha256"
"${VENV_ROOT}/bin/python" - <<'PY' > "${PROVENANCE_DIR}/environment.json"
import importlib.metadata as metadata
import json
import platform

print(json.dumps({
    "python": platform.python_version(),
    "torch": metadata.version("torch"),
    "transformers": metadata.version("transformers"),
    "vllm": metadata.version("vllm"),
}, indent=2, sort_keys=True))
PY

export TOKENIZERS_PARALLELISM=false
export VLLM_USE_V1=1
export VLLM_ATTENTION_BACKEND=FLASHINFER
export TORCH_CUDA_ARCH_LIST=9.0

echo "STAGE server_start"
setsid "${VENV_ROOT}/bin/vllm" serve "${MODEL_ROOT}" \
  --host 127.0.0.1 \
  --port "${SERVER_PORT}" \
  --served-model-name "${SERVED_MODEL}" \
  --tensor-parallel-size 4 \
  --disable-custom-all-reduce \
  --quantization fp8 \
  --gpu-memory-utilization 0.80 \
  --kv-cache-dtype auto \
  --max-model-len 40960 \
  --max-num-batched-tokens 8192 \
  --max-num-seqs 64 \
  --no-enable-prefix-caching \
  --enable-chunked-prefill \
  --enforce-eager \
  --disable-log-requests \
  > "${LOG_DIR}/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" \
    > "${RESULT_DIR}/models.json" 2>/dev/null; then
    READY=1
    break
  fi
  if ! kill -0 "${SERVER_PID}" 2>/dev/null; then
    echo "ERROR: vLLM server exited before readiness" >&2
    tail -200 "${LOG_DIR}/server.log" >&2 || true
    exit 1
  fi
  sleep 5
done
if [[ "${READY}" -ne 1 ]]; then
  echo "ERROR: vLLM server did not become ready within 900 seconds" >&2
  tail -200 "${LOG_DIR}/server.log" >&2 || true
  exit 1
fi

echo "STAGE fixed_shape_single"
"${VENV_ROOT}/bin/python" t0_smoke_client.py \
  --port "${SERVER_PORT}" \
  --served-model "${SERVED_MODEL}" \
  --model-path "${MODEL_ROOT}" \
  --input-tokens 2048 \
  --output-tokens 128 \
  --concurrency 1 \
  --requests 1 \
  --output "${RESULT_DIR}/single.json"

echo "STAGE fixed_shape_concurrency2"
"${VENV_ROOT}/bin/python" t0_smoke_client.py \
  --port "${SERVER_PORT}" \
  --served-model "${SERVED_MODEL}" \
  --model-path "${MODEL_ROOT}" \
  --input-tokens 2048 \
  --output-tokens 128 \
  --concurrency 2 \
  --requests 2 \
  --output "${RESULT_DIR}/concurrency2.json"

cleanup
SERVER_PID=""

sha256sum \
  "${PROVENANCE_DIR}/environment.json" \
  "${PROVENANCE_DIR}/source.sha256" \
  "${RESULT_DIR}/models.json" \
  "${RESULT_DIR}/single.json" \
  "${RESULT_DIR}/concurrency2.json" \
  > "${PROVENANCE_DIR}/artifacts.sha256"
nvidia-smi --query-gpu=index,name,memory.used,utilization.gpu --format=csv,noheader
date -u +"END_UTC=%Y-%m-%dT%H:%M:%SZ"
echo "T0_SMOKE_COMPLETE"