#!/usr/bin/env bash

# Three fresh-server trials for the four-cell Qwen235 vLLM 0.20 surface.

set -euo pipefail

OUT="${OUTPUT_ROOT:?OUTPUT_ROOT is required}"
TRACE_ROOT="${TRACE_ROOT:?TRACE_ROOT is required}"
CASE_NAME="${CASE_NAME:?CASE_NAME is required}"
PREFIX_CACHING="${PREFIX_CACHING:?PREFIX_CACHING is required}"
RUNNER_DIR="${RUNNER_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)}"
RUNNER="${RUNNER_DIR}/run_qwen30_exact_trace_real_anchor.sh"
CLIENT="${RUNNER_DIR}/qwen30_exact_trace_client.py"
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}"
SERVER_READY_ATTEMPTS="${SERVER_READY_ATTEMPTS:-600}"
IDLE_GPU_MEMORY_TOLERANCE_MIB="${IDLE_GPU_MEMORY_TOLERANCE_MIB:-16}"
IDLE_GPU_SETTLE_ATTEMPTS="${IDLE_GPU_SETTLE_ATTEMPTS:-120}"
RESUME_VALID_CELLS="${RESUME_VALID_CELLS:-true}"
PORT="${BASE_PORT:-9300}"
REQUEST_COUNT="$(wc -l < "${TRACE_ROOT}/tp4/private/real_requests.jsonl")"

export PATH="/usr/local/cuda/bin:${PATH}"
command -v nvcc >/dev/null

[[ "$(wc -l < "${TRACE_ROOT}/tp8/private/real_requests.jsonl")" == "${REQUEST_COUNT}" ]] || {
  echo 'ERROR: TP-specific request counts differ' >&2
  exit 1
}
case "${CASE_NAME}" in fixed-pd|fixed-po|trace-pd|trace-po) ;; *) exit 2 ;; esac
case "${PREFIX_CACHING}" in true|false) ;; *) exit 2 ;; esac

mkdir -p "${OUT}/provenance"
declare -a WAVE_PIDS=()

has_valid_result() {
  local result="$1"
  [[ -s "${result}" ]] || return 1
  "${VENV_ROOT}/bin/python" - "${result}" "${REQUEST_COUNT}" <<'PY'
import json, sys
s = json.load(open(sys.argv[1])).get("summary", {})
raise SystemExit(0 if s.get("completed") == int(sys.argv[2]) and s.get("failed") == 0 else 1)
PY
}

preflight_gpus() {
  local attempt
  nvidia-smi --query-gpu=index,name,memory.used,utilization.gpu --format=csv,noheader \
    | tee -a "${OUT}/controller.log"
  for ((attempt = 1; attempt <= IDLE_GPU_SETTLE_ATTEMPTS; attempt++)); do
    if nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits \
      | awk -v tolerance="${IDLE_GPU_MEMORY_TOLERANCE_MIB}" '$1 > tolerance {exit 1}'; then
      return 0
    fi
    sleep 1
  done
  echo "ERROR: GPU memory did not settle below ${IDLE_GPU_MEMORY_TOLERANCE_MIB} MiB after ${IDLE_GPU_SETTLE_ATTEMPTS}s" >&2
  return 1
}

assert_no_server() {
  ! pgrep -fa 'qwen3-235b-v020-eval' > "${OUT}/provenance/unexpected_server_processes.txt"
}

wait_for_wave() {
  local failed=0 pid
  for pid in "${WAVE_PIDS[@]}"; do wait "${pid}" || failed=1; done
  WAVE_PIDS=()
  [[ "${failed}" -eq 0 ]]
}

launch_config() {
  local trial="$1" tp="$2" mns="$3" gpus="$4" ep=false ep_size=1
  if [[ "${tp}" == 8 ]]; then ep=true; ep_size=8; fi
  local config="tp${tp}_ep${ep_size}_mns${mns}"
  local run_out="${OUT}/real/${config}/trial${trial}"
  local requests="${TRACE_ROOT}/tp${tp}/private/real_requests.jsonl" port="${PORT}"
  PORT=$((PORT + 1))
  if [[ "${RESUME_VALID_CELLS}" == true ]] && has_valid_result "${run_out}/results/result.json"; then
    printf 'CONFIG_REUSED_VALID case=%s trial=%s config=%s\n' "${CASE_NAME}" "${trial}" "${config}"
    return 0
  fi
  mkdir -p "${run_out}"
  (
    cd "${RUNNER_DIR}"
    set +e
    env HOME=/tmp/wjh XDG_CACHE_HOME=/tmp/wjh/.cache \
      VLLM_CACHE_ROOT=/tmp/wjh/.cache/vllm CUDA_VISIBLE_DEVICES="${gpus}" \
      TP="${tp}" MNS="${mns}" TRACE_LABEL="${CASE_NAME}/tp${tp}-normalized" \
      PREFIX_CACHING="${PREFIX_CACHING}" ENABLE_EXPERT_PARALLEL="${ep}" \
      MODEL_QUANTIZATION=fp8 DISABLE_CUSTOM_ALL_REDUCE=true GPU_MEMORY_UTILIZATION=0.80 \
      SERVED_MODEL=qwen3-235b-v020-eval SERVER_PORT="${port}" OUTPUT_ROOT="${run_out}" \
      REQUESTS_FILE="${requests}" VENV_ROOT="${VENV_ROOT}" MODEL_ROOT="${MODEL_ROOT}" \
      SERVER_READY_ATTEMPTS="${SERVER_READY_ATTEMPTS}" CLIENT_TIMEOUT_SECONDS=6000 \
      EXACT_TRACE_CLIENT="${CLIENT}" \
      timeout --signal=TERM --kill-after=60s 7800 bash "${RUNNER}" \
      > "${run_out}/launcher.stdout.log" 2> "${run_out}/launcher.stderr.log"
    status="$?"
    printf '%s\n' "${status}" > "${run_out}/launcher.exit_code"
    exit "${status}"
  ) &
  WAVE_PIDS+=("$!")
}

run_trial() {
  local trial="$1" order="$2"
  IFS=',' read -r -a mnss <<< "${order}"
  preflight_gpus && assert_no_server
  printf 'WAVE_START case=%s trial=%s tp=4 mns=%s,%s\n' \
    "${CASE_NAME}" "${trial}" "${mnss[0]}" "${mnss[1]}" | tee -a "${OUT}/controller.log"
  launch_config "${trial}" 4 "${mnss[0]}" '0,1,2,3'
  launch_config "${trial}" 4 "${mnss[1]}" '4,5,6,7'
  wait_for_wave
  preflight_gpus
  assert_no_server
  for mns in "${mnss[@]}"; do
    printf 'WAVE_START case=%s trial=%s tp=8 mns=%s\n' \
      "${CASE_NAME}" "${trial}" "${mns}" | tee -a "${OUT}/controller.log"
    launch_config "${trial}" 8 "${mns}" '0,1,2,3,4,5,6,7'
    wait_for_wave
    preflight_gpus
    assert_no_server
  done
}

{
  printf '%s\n' "Q235_REAL_LAUNCH_ECHO host=dash0 model=Qwen3-235B-A22B-FP8 engine=vLLM-0.20.0+cu129 checkpoint=FP8 cases=${CASE_NAME} trace=${TRACE_ROOT}/tp{4,8} requests=${REQUEST_COUNT} transform=t_prime=t/TP prefix=${PREFIX_CACHING} surface={TP4/EP1,TP8/EP8}xMNS{64,128} MBT=8192 trials=3 fresh_server=true metrics=mean,p90(TTFT,TPOT-if-PD,E2E) SLO=not_scored output=${OUT}/real"
  date -u +START_UTC=%Y-%m-%dT%H:%M:%SZ
  sha256sum "${BASH_SOURCE[0]}" "${RUNNER}" "${CLIENT}" "${MODEL_ROOT}/config.json" \
    > "${OUT}/provenance/real-input.sha256"
  run_trial 1 '64,128'
  run_trial 2 '128,64'
  run_trial 3 '64,128'
  date -u +END_UTC=%Y-%m-%dT%H:%M:%SZ
  printf '%s\n' 'Q235_REAL_SURFACE_COMPLETE'
} 2>&1 | tee -a "${OUT}/controller.log"