#!/usr/bin/env bash

# Three fresh-server real vLLM trials for the per-GPU-normalized Qwen30
# Trace-PD surface. This script and its JIT-tolerant harness are copied into
# the remote output root, not the shared/dirty AITuner checkout.

set -euo pipefail

OUT="${1:?output root is required}"
RUNNER_DIR="/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-v1/runs/frontier-fidelity-envelope-v1"
RUNNER="${OUT}/bin/run_qwen30_exact_trace_real_anchor_jit_tolerant.sh"
CLIENT="${OUT}/bin/qwen30_exact_trace_client.py"
PREFILL_CLIENT="/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-v1/runs/frontier-phase-factorial-v0/qwen30_prefill_client.py"
VENV_ROOT="/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1"
MODEL_ROOT="/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B"
TRACE_ROOT="${OUT}/traces"
# FlashInfer's JIT uses per-operation file locks. Reuse the smoke-populated
# cache across fresh servers so the 36 measurement runs do not each rebuild
# the same H20 MoE/all-reduce kernels. Server processes and metrics remain
# isolated per run under ${OUT}/real.
FLASHINFER_SHARED_WORKSPACE="${OUT}/flashinfer-shared-workspace"
PORT=8200

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 experiment 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="tp${tp}-normalized-u0p01" \
      SERVER_PORT="${port}" \
      OUTPUT_ROOT="${run_out}" \
      REQUESTS_FILE="${requests}" \
      VENV_ROOT="${VENV_ROOT}" \
      MODEL_ROOT="${MODEL_ROOT}" \
      FLASHINFER_WORKSPACE_BASE="${FLASHINFER_SHARED_WORKSPACE}" \
      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 trial=%s tp=%s mns=%s\n' "${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'
      ;;
    *)
      printf 'ERROR unsupported TP=%s\n' "${tp}" >&2
      return 1
      ;;
  esac
  wait_for_wave
  preflight_gpus
  assert_no_our_server
  printf 'WAVE_COMPLETE trial=%s tp=%s\n' "${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; trace=trace-exact-v1/u0p01; requests=129; transform=t_prime=t/TP; per_gpu_offered_rate=0.215 req/s; surface=TP{1,2,4}xMNS{8,16,32,64}; trials=3; fresh_server=true; metrics=mean,p90(TTFT,TPOT,E2E); SLO=not_scored; flashinfer_cache=${FLASHINFER_SHARED_WORKSPACE}; 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"
  mkdir -p "${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"

  # Rotate topology and MNS order across trials to avoid a fixed temporal order.
  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
