#!/usr/bin/env bash

set -euo pipefail

TP="${TP:?TP is required}"
OUTPUT_ROOT="${OUTPUT_ROOT:?OUTPUT_ROOT is required}"
SERVER_PORT="${SERVER_PORT:?SERVER_PORT is required}"
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}"
GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.92}"
SERVER_READY_ATTEMPTS="${SERVER_READY_ATTEMPTS:-900}"
WARMUP_REQUESTS="${WARMUP_REQUESTS:-2}"
SCRIPT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MATERIALIZER="${SCRIPT_ROOT}/materialize_prefill_request.py"
CLIENT="${SCRIPT_ROOT}/run_request.py"
SERVER_PID=""

IFS=',' read -r -a GPU_IDS <<< "${CUDA_VISIBLE_DEVICES:?GPU allocation is required}"
if [[ "${#GPU_IDS[@]}" -ne "${TP}" ]]; then
  echo "ERROR: TP=${TP}, but CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" >&2
  exit 1
fi
if [[ ! "${WARMUP_REQUESTS}" =~ ^[1-9][0-9]*$ ]]; then
  echo "ERROR: WARMUP_REQUESTS must be a positive integer" >&2
  exit 1
fi

mkdir -p "${OUTPUT_ROOT}"
OUTPUT_ROOT="$(cd "${OUTPUT_ROOT}" && pwd)"
mkdir -p "${OUTPUT_ROOT}/logs" "${OUTPUT_ROOT}/provenance" \
  "${OUTPUT_ROOT}/requests" "${OUTPUT_ROOT}/results" \
  "${OUTPUT_ROOT}/trace-staging" "${OUTPUT_ROOT}/traces/profile"
exec > >(tee -a "${OUTPUT_ROOT}/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

sha256sum "${BASH_SOURCE[0]}" "${MATERIALIZER}" "${CLIENT}" \
  "${MODEL_ROOT}/config.json" > "${OUTPUT_ROOT}/provenance/inputs.sha256"
"${VENV_ROOT}/bin/python" -c \
  'import torch, transformers, vllm; print(f"torch={torch.__version__}"); print(f"transformers={transformers.__version__}"); print(f"vllm={vllm.__version__}")' \
  > "${OUTPUT_ROOT}/provenance/runtime.versions"
nvidia-smi --query-gpu=index,name,uuid,driver_version,memory.total \
  --format=csv,noheader > "${OUTPUT_ROOT}/provenance/gpus.before.csv"
ps -eo user,pid,ppid,etimes,pcpu,pmem,args --sort=pid \
  > "${OUTPUT_ROOT}/provenance/processes.before.txt"
uptime > "${OUTPUT_ROOT}/provenance/uptime.before.txt"
env | sort > "${OUTPUT_ROOT}/provenance/environment.txt"

PROFILE_CONFIG="$("${VENV_ROOT}/bin/python" - "${OUTPUT_ROOT}/trace-staging" <<'PY'
import json
import sys
print(json.dumps({
    "profiler": "torch",
    "torch_profiler_dir": sys.argv[1],
    "torch_profiler_with_stack": False,
    "torch_profiler_record_shapes": True,
    "torch_profiler_use_gzip": True,
    "ignore_frontend": True,
    "wait_iterations": 0,
    "warmup_iterations": 0,
    "active_iterations": 2,
}, separators=(",", ":")))
PY
)"
printf '%s\n' "${PROFILE_CONFIG}" > "${OUTPUT_ROOT}/provenance/profiler-config.json"

export TOKENIZERS_PARALLELISM=false
export VLLM_USE_V1=1
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
export FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE:-${OUTPUT_ROOT}/flashinfer-workspace}"
mkdir -p "${FLASHINFER_WORKSPACE_BASE}"
ulimit -n 65536

"${VENV_ROOT}/bin/python" "${MATERIALIZER}" --model "${MODEL_ROOT}" \
  --input-tokens 8192 --output-tokens 2 \
  --output "${OUTPUT_ROOT}/requests/q8192-o2.json"

printf 'LAUNCH host=%s tp=%s gpus=%s output=%s warmups=%s\n' \
  "$(hostname)" "${TP}" "${CUDA_VISIBLE_DEVICES}" "${OUTPUT_ROOT}" \
  "${WARMUP_REQUESTS}"
setsid "${VENV_ROOT}/bin/vllm" serve "${MODEL_ROOT}" \
  --host 127.0.0.1 --port "${SERVER_PORT}" \
  --served-model-name qwen30-prefill-profile \
  --tensor-parallel-size "${TP}" \
  --gpu-memory-utilization "${GPU_MEMORY_UTILIZATION}" \
  --max-model-len 40960 --max-num-batched-tokens 8192 --max-num-seqs 16 \
  --no-enable-prefix-caching --enable-chunked-prefill --no-enable-log-requests \
  --enable-logging-iteration-details --profiler-config "${PROFILE_CONFIG}" \
  > "${OUTPUT_ROOT}/logs/server.log" 2>&1 &
SERVER_PID=$!

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

for index in $(seq 1 "${WARMUP_REQUESTS}"); do
  "${VENV_ROOT}/bin/python" "${CLIENT}" --port "${SERVER_PORT}" \
    --request "${OUTPUT_ROOT}/requests/q8192-o2.json" \
    --output "${OUTPUT_ROOT}/results/warmup-${index}.json"
done

curl -fsS -X POST "http://127.0.0.1:${SERVER_PORT}/start_profile" \
  > "${OUTPUT_ROOT}/logs/start-profile.txt"
"${VENV_ROOT}/bin/python" "${CLIENT}" --port "${SERVER_PORT}" \
  --request "${OUTPUT_ROOT}/requests/q8192-o2.json" \
  --output "${OUTPUT_ROOT}/results/profile.json"
curl -fsS -X POST "http://127.0.0.1:${SERVER_PORT}/stop_profile" \
  > "${OUTPUT_ROOT}/logs/stop-profile.txt"

deadline=$((SECONDS + 120))
while (( SECONDS < deadline )); do
  trace_count="$(find "${OUTPUT_ROOT}/trace-staging" -maxdepth 1 -type f \
    -name '*.pt.trace.json*' | wc -l)"
  if (( trace_count >= TP )); then
    break
  fi
  sleep 2
done
trace_count="$(find "${OUTPUT_ROOT}/trace-staging" -maxdepth 1 -type f \
  -name '*.pt.trace.json*' | wc -l)"
if (( trace_count < TP )); then
  echo "ERROR: expected ${TP} rank traces, found ${trace_count}" >&2
  exit 1
fi
find "${OUTPUT_ROOT}/trace-staging" -maxdepth 1 -type f \
  -name '*.pt.trace.json*' -exec mv -t "${OUTPUT_ROOT}/traces/profile" {} +

cleanup
nvidia-smi --query-gpu=index,name,uuid,driver_version,memory.total \
  --format=csv,noheader > "${OUTPUT_ROOT}/provenance/gpus.after.csv"
find "${OUTPUT_ROOT}" -type f \
  ! -path '*/provenance/artifacts.sha256' -print0 \
  | sort -z | xargs -0 sha256sum \
  > "${OUTPUT_ROOT}/provenance/artifacts.sha256"
echo PREFILL_PROFILE_COMPLETE