#!/usr/bin/env bash

# One three-trial Qwen30 latency-selection surface.  The caller supplies a
# frozen trace root and runs different cases sequentially to keep GPU use
# exclusive.

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"
PREFILL_CLIENT="${RUNNER_DIR}/../frontier-phase-factorial-v0/qwen30_prefill_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-30B-A3B}"
FLASHINFER_SHARED_WORKSPACE="${FLASHINFER_SHARED_WORKSPACE:-${OUT}/flashinfer-shared-workspace}"
SERVER_READY_ATTEMPTS="${SERVER_READY_ATTEMPTS:-300}"
PORT="${BASE_PORT:-8300}"

case "${PREFIX_CACHING}" in
  true|false) ;;
  *) echo "ERROR: PREFIX_CACHING must be true or false" >&2; exit 1 ;;
esac

case "${CASE_NAME}" in
  trace-po|fixed-pd|fixed-po) ;;
  *) echo "ERROR: unsupported CASE_NAME=${CASE_NAME}" >&2; exit 1 ;;
esac

for path in "${RUNNER}" "${CLIENT}" "${PREFILL_CLIENT}" "${MODEL_ROOT}/config.json"; do
  [[ -f "${path}" ]] || { echo "ERROR: required path is missing: ${path}" >&2; exit 1; }
done
mkdir -p "${OUT}"

declare -a WAVE_PIDS=()

wait_for_wave() {
  local failed=0 pid
  for pid in "${WAVE_PIDS[@]}"; do
    if ! wait "${pid}"; then
      failed=1
    fi
  done
  WAVE_PIDS=()
  if [[ "${failed}" -ne 0 ]]; then
    printf '%s\n' 'ERROR one or more real-serving runs failed; aborting surface.' \
      | tee -a "${OUT}/controller.log" >&2
    return 1
  fi
}

preflight_gpus() {
  nvidia-smi --query-gpu=index,name,memory.used,utilization.gpu --format=csv,noheader \
    | tee -a "${OUT}/controller.log"
  if nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits \
    | awk '$1 != 0 {exit 1}'; then
    return 0
  fi
  printf '%s\n' 'ERROR non-idle GPU before a real-serving wave; refusing to start.' \
    | tee -a "${OUT}/controller.log" >&2
  return 1
}

assert_no_our_server() {
  if pgrep -fa 'qwen3-30b-exact-trace' > "${OUT}/provenance/unexpected_server_processes.txt"; then
    printf '%s\n' 'ERROR a prior Qwen30 server remains; refusing the next wave.' \
      | tee -a "${OUT}/controller.log" >&2
    return 1
  fi
}

launch_config() {
  local trial="$1" tp="$2" mns="$3" gpus="$4"
  local config="tp${tp}_mns${mns}"
  local run_out="${OUT}/real/${config}/trial${trial}"
  local requests="${TRACE_ROOT}/tp${tp}/private/real_requests.jsonl"
  local port="${PORT}"
  PORT=$((PORT + 1))
  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}" \
      SERVER_PORT="${port}" \
      OUTPUT_ROOT="${run_out}" \
      REQUESTS_FILE="${requests}" \
      VENV_ROOT="${VENV_ROOT}" \
      MODEL_ROOT="${MODEL_ROOT}" \
      FLASHINFER_WORKSPACE_BASE="${FLASHINFER_SHARED_WORKSPACE}" \
      SERVER_READY_ATTEMPTS="${SERVER_READY_ATTEMPTS}" \
      EXACT_TRACE_CLIENT="${CLIENT}" \
      timeout --signal=TERM --kill-after=60s 4200 bash "${RUNNER}" \
      > "${run_out}/launcher.stdout.log" 2> "${run_out}/launcher.stderr.log"
    local_status="$?"
    printf '%s\n' "${local_status}" > "${run_out}/launcher.exit_code"
    exit "${local_status}"
  ) &
  WAVE_PIDS+=("$!")
}

run_tp_wave() {
  local trial="$1" tp="$2" order="$3"
  IFS=',' read -r -a mnss <<< "${order}"
  WAVE_PIDS=()
  preflight_gpus
  assert_no_our_server
  printf 'WAVE_START case=%s trial=%s tp=%s mns=%s\n' "${CASE_NAME}" "${trial}" "${tp}" "${order}" \
    | tee -a "${OUT}/controller.log"
  case "${tp}" in
    4)
      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_our_server
      launch_config "${trial}" 4 "${mnss[2]}" '0,1,2,3'
      launch_config "${trial}" 4 "${mnss[3]}" '4,5,6,7'
      ;;
    2)
      launch_config "${trial}" 2 "${mnss[0]}" '0,1'
      launch_config "${trial}" 2 "${mnss[1]}" '2,3'
      launch_config "${trial}" 2 "${mnss[2]}" '4,5'
      launch_config "${trial}" 2 "${mnss[3]}" '6,7'
      ;;
    1)
      launch_config "${trial}" 1 "${mnss[0]}" '0'
      launch_config "${trial}" 1 "${mnss[1]}" '1'
      launch_config "${trial}" 1 "${mnss[2]}" '2'
      launch_config "${trial}" 1 "${mnss[3]}" '3'
      ;;
  esac
  wait_for_wave
  preflight_gpus
  assert_no_our_server
  printf 'WAVE_COMPLETE case=%s trial=%s tp=%s\n' "${CASE_NAME}" "${trial}" "${tp}" \
    | tee -a "${OUT}/controller.log"
}

{
  printf '%s\n' "REAL_LAUNCH_ECHO host=dash0 model=Qwen3-30B-A3B engine=vLLM-0.20.0+cu129 dtype=BF16 case=${CASE_NAME} traces=${TRACE_ROOT}/tp{1,2,4}/private/real_requests.jsonl prefix=${PREFIX_CACHING} requests=129 transform=t_prime=t/TP surface=TP{1,2,4}xMNS{8,16,32,64} trials=3 fresh_server=true metrics=mean,p90(TTFT,TPOT-if-OSL-gt-1,E2E) expected_cost=13_H20-GPUh_nominal__41_H20-GPUh_max output=${OUT}/real"
  date -u +START_UTC=%Y-%m-%dT%H:%M:%SZ
  mkdir -p "${OUT}/provenance" "${FLASHINFER_SHARED_WORKSPACE}"
  sha256sum "${BASH_SOURCE[0]}" "${RUNNER}" "${CLIENT}" "${PREFILL_CLIENT}" \
    "${MODEL_ROOT}/config.json" > "${OUT}/provenance/real-input.sha256"
  for tp in 1 2 4; do
    sha256sum "${TRACE_ROOT}/tp${tp}/public/manifest.json" \
      "${TRACE_ROOT}/tp${tp}/private/real_requests.jsonl" \
      >> "${OUT}/provenance/real-input.sha256"
  done
  "${VENV_ROOT}/bin/vllm" --version > "${OUT}/provenance/vllm.version"
  "${VENV_ROOT}/bin/python" -c 'import torch, transformers, vllm; print(f"torch={torch.__version__}"); print(f"transformers={transformers.__version__}"); print(f"vllm={vllm.__version__}")' \
    > "${OUT}/provenance/runtime.versions"
  nvidia-smi --query-gpu=index,name,uuid,driver_version,memory.total --format=csv,noheader \
    > "${OUT}/provenance/gpus.before.csv"

  run_tp_wave 1 4 '8,16,32,64'
  run_tp_wave 1 2 '8,16,32,64'
  run_tp_wave 1 1 '8,16,32,64'
  run_tp_wave 2 1 '64,32,16,8'
  run_tp_wave 2 2 '64,32,16,8'
  run_tp_wave 2 4 '64,32,16,8'
  run_tp_wave 3 2 '16,32,64,8'
  run_tp_wave 3 4 '16,32,64,8'
  run_tp_wave 3 1 '16,32,64,8'

  nvidia-smi --query-gpu=index,name,uuid,driver_version,memory.total --format=csv,noheader \
    > "${OUT}/provenance/gpus.after.csv"
  find "${OUT}/real" -type f ! -path '*/provenance/artifacts.sha256' -print0 \
    | sort -z | xargs -0 sha256sum > "${OUT}/provenance/real-artifacts.sha256"
  date -u +END_UTC=%Y-%m-%dT%H:%M:%SZ
  printf '%s\n' 'REAL_SURFACE_COMPLETE'
} >> "${OUT}/controller.log" 2>&1