#!/usr/bin/env bash

# Frozen Qwen30 Fixed-PD/Fixed-PO Frontier-versus-real campaign. The pressure
# calibration artifacts are deliberately not reused as evaluation results.

set -euo pipefail

CAMPAIGN_ROOT="${CAMPAIGN_ROOT:?CAMPAIGN_ROOT is required}"
RUNNER_DIR="${RUNNER_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)}"
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}"
FRONTIER_SOURCE="${FRONTIER_SOURCE:-/home/admin/cpfs/wjh/aituner/frontier-t1-dash0-deadc4a}"
REPLAYSERVE_ROOT="${REPLAYSERVE_ROOT:-/home/admin/cpfs/wjh/replayserve}"
PYTHON_DEPS="${PYTHON_DEPS:-${VENV_ROOT}/lib/python3.12/site-packages}"
PROFILE_ROOT="${PROFILE_ROOT:-/home/admin/cpfs/wjh/aituner/aituner-graph-piecewise-bdc357d/runs/frontier-fidelity-envelope-v1/profiles/profile-v4-trace-final}"
KERNEL_PROFILE_ROOT="${KERNEL_PROFILE_ROOT:-/home/admin/cpfs/wjh/aituner/graph-piecewise-qwen30-20260717/full/frozen-kernel-only}"
ALLREDUCE_CSV="${ALLREDUCE_CSV:-/home/admin/cpfs/wjh/aituner/aituner-graph-piecewise-bdc357d/runs/frontier-fidelity-envelope-v1/profiles/measured-allreduce.csv}"
FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE:?FLASHINFER_WORKSPACE_BASE is required}"
REQUESTS="${REQUESTS:-257}"
PER_GPU_RATE="${PER_GPU_RATE:-1.125}"

[[ "${REQUESTS}" == "257" ]] || { echo 'ERROR: frozen campaign requires REQUESTS=257' >&2; exit 1; }
[[ "${PER_GPU_RATE}" == "1.125" ]] || { echo 'ERROR: frozen campaign requires PER_GPU_RATE=1.125' >&2; exit 1; }

MATERIALIZER="${RUNNER_DIR}/prepare_qwen30_latency_case.py"
REAL_RUNNER="${RUNNER_DIR}/run_qwen30_latency_case_real_surface.sh"
SIM_RUNNER="${RUNNER_DIR}/run_frontier_qwen30_exact_trace_surface.py"
REAL_AUDITOR="${RUNNER_DIR}/audit_qwen30_latency_case.py"
COMPARATOR="${RUNNER_DIR}/analyze_qwen30_latency_case.py"

mkdir -p "${CAMPAIGN_ROOT}/provenance" "${CAMPAIGN_ROOT}/traces" \
  "${CAMPAIGN_ROOT}/real" "${CAMPAIGN_ROOT}/sim" "${CAMPAIGN_ROOT}/analysis"
exec > >(tee -a "${CAMPAIGN_ROOT}/controller.log") 2>&1

printf '%s\n' "Q30_FIXED_PRESSURE_CAMPAIGN_LAUNCH_ECHO host=dash0 model=Qwen3-30B-A3B engine=vLLM-0.20.0+cu129 dtype=BF16 cases={fixed-pd:4096_to_256,fixed-po:4096_to_1} prefix=false requests=${REQUESTS} per_gpu_rate=${PER_GPU_RATE} global_rates={TP1:1.125,TP2:2.25,TP4:4.5} surface_per_case=TP{1,2,4}xMNS{8,16,32,64} real_trials=3 fresh_server=true simulator=Frontier-deadc4a_piecewise_graph-kernel-only metrics=mean,p90(TTFT,TPOT-if-PD,E2E) SLO=not_scored flashinfer_workspace=${FLASHINFER_WORKSPACE_BASE} expected_wall=2-5h expected_cost=12-24_H20-GPUh output=${CAMPAIGN_ROOT}"
date -u +START_UTC=%Y-%m-%dT%H:%M:%SZ

sha256sum "${BASH_SOURCE[0]}" "${MATERIALIZER}" "${REAL_RUNNER}" \
  "${SIM_RUNNER}" "${REAL_AUDITOR}" "${COMPARATOR}" \
  "${MODEL_ROOT}/config.json" "${PROFILE_ROOT}/manifest.json" \
  "${KERNEL_PROFILE_ROOT}/manifest.json" "${ALLREDUCE_CSV}" \
  > "${CAMPAIGN_ROOT}/provenance/input.sha256"
git -C "${FRONTIER_SOURCE}" rev-parse HEAD > "${CAMPAIGN_ROOT}/provenance/frontier.commit"
git -C "${RUNNER_DIR}" rev-parse HEAD > "${CAMPAIGN_ROOT}/provenance/aituner.commit"
"${VENV_ROOT}/bin/vllm" --version > "${CAMPAIGN_ROOT}/provenance/vllm.version"
nvidia-smi --query-gpu=index,name,uuid,driver_version,memory.total --format=csv,noheader \
  > "${CAMPAIGN_ROOT}/provenance/gpus.before.csv"

prepare_case() {
  local case_name="$1" output_tokens="$2"
  local tp
  for tp in 1 2 4; do
    "${VENV_ROOT}/bin/python" "${MATERIALIZER}" fixed \
      --model "${MODEL_ROOT}" --input-tokens 4096 --output-tokens "${output_tokens}" \
      --requests "${REQUESTS}" --per-gpu-rate "${PER_GPU_RATE}" --tp "${tp}" \
      --output-root "${CAMPAIGN_ROOT}/traces/${case_name}/tp${tp}"
  done
}

run_real_case() {
  local case_name="$1" base_port="$2"
  CASE_NAME="${case_name}" PREFIX_CACHING=false \
    TRACE_ROOT="${CAMPAIGN_ROOT}/traces/${case_name}" \
    OUTPUT_ROOT="${CAMPAIGN_ROOT}/real/${case_name}" \
    RUNNER_DIR="${RUNNER_DIR}" VENV_ROOT="${VENV_ROOT}" MODEL_ROOT="${MODEL_ROOT}" \
    FLASHINFER_SHARED_WORKSPACE="${FLASHINFER_WORKSPACE_BASE}" \
    BASE_PORT="${base_port}" RESUME_VALID_CELLS=true \
    bash "${REAL_RUNNER}"
}

run_sim_tp() {
  local case_name="$1" tp="$2"
  local sim_root="${CAMPAIGN_ROOT}/sim/${case_name}"
  local -a configs=()
  local mns
  for mns in 8 16 32 64; do
    configs+=(--config "tp${tp}_mns${mns}")
  done
  /usr/bin/python3 "${SIM_RUNNER}" \
    --frontier-source "${FRONTIER_SOURCE}" --replayserve-root "${REPLAYSERVE_ROOT}" \
    --profile-root "${PROFILE_ROOT}" --kernel-profile-root "${KERNEL_PROFILE_ROOT}" \
    --python-deps "${PYTHON_DEPS}" --output-root "${sim_root}" \
    --trace "tp${tp}=${CAMPAIGN_ROOT}/traces/${case_name}/tp${tp}/public/frontier.csv" \
    "${configs[@]}" --rate-contract uniform-spacing --no-prefix-caching \
    --cc-backend vidur --allreduce-csv "${ALLREDUCE_CSV}" \
    --timeout-seconds 3600 --predictor-training-job-threads 4 \
    --decode-cuda-graph-mode piecewise --align-real-graph-runtime \
    --fresh-predictor-cache --resume --continue-on-failure
}

run_sim_case() {
  local case_name="$1" failed=0 pid
  local -a pids=()
  mkdir -p "${CAMPAIGN_ROOT}/sim/${case_name}"
  for tp in 1 2 4; do
    run_sim_tp "${case_name}" "${tp}" \
      > "${CAMPAIGN_ROOT}/sim/${case_name}/launcher-tp${tp}.log" 2>&1 &
    pids+=("$!")
  done
  for pid in "${pids[@]}"; do
    wait "${pid}" || failed=1
  done
  [[ "${failed}" -eq 0 ]] || return 1
}

analyze_case() {
  local case_name="$1"
  local analysis_root="${CAMPAIGN_ROOT}/analysis/${case_name}"
  mkdir -p "${analysis_root}"
  "${VENV_ROOT}/bin/python" "${REAL_AUDITOR}" \
    --case-root "${CAMPAIGN_ROOT}/real/${case_name}" \
    --traces-root "${CAMPAIGN_ROOT}/traces/${case_name}" \
    --json-output "${analysis_root}/real-audit.json" \
    --markdown-output "${analysis_root}/real-audit.md"
  "${VENV_ROOT}/bin/python" "${COMPARATOR}" \
    --sim-root "${CAMPAIGN_ROOT}/sim/${case_name}" \
    --real-audit "${analysis_root}/real-audit.json" \
    --json-output "${analysis_root}/comparison.json" \
    --markdown-output "${analysis_root}/comparison.md"
}

prepare_case fixed-pd 256
prepare_case fixed-po 1

run_real_case fixed-pd 9000
run_real_case fixed-po 9100

run_sim_case fixed-pd
analyze_case fixed-pd
run_sim_case fixed-po
analyze_case fixed-po

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