#!/usr/bin/env bash

set -euo pipefail

OUTPUT_ROOT="${OUTPUT_ROOT:-/home/admin/cpfs/wjh/frontier-community-qwen235-smoke-20260715}"
FRONTIER_ROOT="${FRONTIER_ROOT:-${OUTPUT_ROOT}/Frontier-d9cfeb6}"
VENV_ROOT="${VENV_ROOT:-/tmp/wjh-frontier-vllm0102-smoke/.venv}"
MODEL_ROOT="${MODEL_ROOT:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-235B-A22B-FP8}"
PROFILE_ROOT="${OUTPUT_ROOT}/profiles"
LOG_DIR="${OUTPUT_ROOT}/logs"
RESULT_DIR="${OUTPUT_ROOT}/results"
SERVER_PORT="${SERVER_PORT:-18900}"
SKIP_LINEAR="${SKIP_LINEAR:-0}"
SERVING_ONLY="${SERVING_ONLY:-0}"
SERVED_MODEL="qwen3-235b-community-smoke"
SERVER_PID=""

mkdir -p "${PROFILE_ROOT}" "${LOG_DIR}" "${RESULT_DIR}"
exec > >(tee -a "${LOG_DIR}/gpu_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 name exactly four allocated GPUs" >&2
  exit 1
fi
IFS=',' read -r -a GPU_IDS <<< "${CUDA_VISIBLE_DEVICES}"
if [[ "${#GPU_IDS[@]}" -ne 4 ]]; then
  echo "ERROR: expected four allocated GPUs, got CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" >&2
  exit 1
fi

echo "LAUNCH_ECHO host=$(hostname) gpus=${CUDA_VISIBLE_DEVICES} model=${MODEL_ROOT} frontier=d9cfeb6 vllm=community-0.10.2 transformers=4.55.2 backend=FLASHINFER execution=eager kv=BF16 spec=off tasks=representative-FP8-linear/attention/MoE+TP4-allreduce+TP4-model-load+one-request skip_linear=${SKIP_LINEAR} serving_only=${SERVING_ONLY} hard_wall_cap=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/python"
test -f "${FRONTIER_ROOT}/pyproject.toml"
test -f "${MODEL_ROOT}/config.json"

export PYTHONPATH="${FRONTIER_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"
export TOKENIZERS_PARALLELISM=false
export VLLM_USE_V1=1
export VLLM_ATTENTION_BACKEND=FLASHINFER
export TORCH_CUDA_ARCH_LIST=9.0

cd "${FRONTIER_ROOT}"

if [[ "${SERVING_ONLY}" -eq 0 && "${SKIP_LINEAR}" -eq 0 ]]; then
  echo "STAGE linear_op"
  timeout 300 "${VENV_ROOT}/bin/python" -m frontier.profiling.linear_op.main \
    --disable_ray \
    --models Qwen3-235B-A22B-FP8 \
    --num_gpus 1 \
    --max_tokens 16 \
    --num_tokens_list 16 \
    --num_tensor_parallel_workers 4 \
    --profile_method cuda_event \
    --device h20 \
    --output_dir "${PROFILE_ROOT}" \
    --is_moe \
    --yes
elif [[ "${SERVING_ONLY}" -eq 0 ]]; then
  echo "STAGE linear_op SKIPPED (existing artifact retained)"
fi

if [[ "${SERVING_ONLY}" -eq 0 ]]; then
  echo "STAGE attention"
  timeout 300 "${VENV_ROOT}/bin/python" -m frontier.profiling.attention.main \
  --disable_ray \
  --models Qwen3-235B-A22B-FP8 \
  --num_gpus 1 \
  --max_model_len 40960 \
  --max_seq_len 128 \
  --min_batch_size 1 \
  --max_batch_size 1 \
  --batch_size_list 1 \
  --num_tensor_parallel_workers 4 \
  --max_pipeline_parallel_size 1 \
  --attention_backend FLASHINFER \
  --block_size 16 \
  --profile_only_prefill \
  --fixed_chunked_prefill_size 128 \
  --device h20 \
  --profile_method cuda_event \
  --output_dir "${PROFILE_ROOT}" \
    --yes

  echo "STAGE moe"
  timeout 300 "${VENV_ROOT}/bin/python" -m frontier.profiling.moe.main \
  --disable_ray \
  --models Qwen3-235B-A22B-FP8 \
  --device h20 \
  --num_gpus 1 \
  --max_tokens 16 \
  --num_tokens_list 16 \
  --num_tensor_parallel_workers 4 \
  --expert_parallel_sizes 1 \
  --load_distributions uniform \
  --num_samples_per_distribution 1 \
  --routing_runtime_path standard_fused_topk \
  --gating_runtime_context prefill_hot \
  --profile_method cuda_event \
  --output_dir "${PROFILE_ROOT}" \
    --yes

  echo "STAGE allreduce"
  timeout 180 "${VENV_ROOT}/bin/torchrun" \
    --standalone \
    --nnodes=1 \
    --nproc-per-node=4 \
    "${OUTPUT_ROOT}/scripts/allreduce_smoke.py" \
    | tee "${RESULT_DIR}/allreduce_tp4.jsonl"
else
  echo "STAGES Frontier profiles and allreduce SKIPPED (existing artifacts retained)"
fi

echo "STAGE serving"
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

curl -fsS --max-time 120 \
  -H 'Content-Type: application/json' \
  -d '{"model":"qwen3-235b-community-smoke","prompt":"Hello","max_tokens":1,"temperature":0}' \
  "http://127.0.0.1:${SERVER_PORT}/v1/completions" \
  | tee "${RESULT_DIR}/one_request.json"
echo
jq -e '.choices | length == 1' "${RESULT_DIR}/one_request.json" >/dev/null

cleanup
SERVER_PID=""

find "${PROFILE_ROOT}" -type f -maxdepth 5 -print -exec sha256sum {} \;
sha256sum \
  "${RESULT_DIR}/allreduce_tp4.jsonl" \
  "${RESULT_DIR}/models.json" \
  "${RESULT_DIR}/one_request.json" \
  > "${RESULT_DIR}/results.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 "GPU_SMOKE_COMPLETE"