#!/usr/bin/env bash
# =============================================================================
# run_benchmarks.sh
# llama-bench harness for 2 x RTX PRO 6000 Blackwell + Ryzen 9 9950X.
#
# What it does:
#   1. Sets both GPUs to persistence mode and a configurable power limit.
#   2. Collects system / build info.
#   3. For each model: runs llama-bench, samples nvidia-smi (temp / power /
#      clocks / utilisation / VRAM / throttle) for ONLY the GPUs that run
#      uses, captures the markdown results table, and records per-GPU
#      thermal maxima.
#   4. Writes per-model sections to benchmark_results.md, a manifest.tsv for
#      the report generator, and a final summary. All files land under
#      bench_logs/<power_limit>/.
#   5. Automatically renders the HTML report with make_report_html.py.
#
# Usage:
#   ./run_benchmarks.sh
#   GPU_POWER_LIMIT_W=450 ./run_benchmarks.sh
#   ONLY="qwen2.5-coder-32b hermes-4-70b" ./run_benchmarks.sh
#   COOLDOWN_S=60 ./run_benchmarks.sh        # extra cool-down between models
#   NO_HTML=1 ./run_benchmarks.sh            # skip the HTML report step
#
# HF_TOKEN must already be in the environment (export from ~/.zshrc).
# Requirements: llama-bench in PATH, nvidia-smi, awk, python3 (for the HTML
# report). sudo is used only to set the power limit; if unavailable the
# harness continues without it.
# =============================================================================

set -uo pipefail

# ----------------------------- configuration ---------------------------------
GPU_POWER_LIMIT_W="${GPU_POWER_LIMIT_W:-400}"   # max for this card is 600 W
THREADS="${THREADS:-16}"                          # 9950X physical cores
KV_TYPE="${KV_TYPE:-q8_0}"                        # KV cache precision
REPS="${REPS:-5}"                                 # llama-bench repetitions
ONLY="${ONLY:-}"                                  # space-separated aliases to restrict to
SMI_INTERVAL_MS="${SMI_INTERVAL_MS:-500}"         # nvidia-smi sampling interval
COOLDOWN_S="${COOLDOWN_S:-20}"                    # idle seconds between models
NO_HTML="${NO_HTML:-}"                            # non-empty: skip HTML report

OUTDIR="$(pwd)"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LOGDIR="${OUTDIR}/bench_logs/${GPU_POWER_LIMIT_W}"
RESULTS_MD="${LOGDIR}/benchmark_results.md"
MANIFEST="${LOGDIR}/manifest.tsv"
REPORT_HTML="${LOGDIR}/report_${GPU_POWER_LIMIT_W}W.html"
mkdir -p "$LOGDIR"

# HF_TOKEN must be exported from ~/.zshrc or set inline.
if [[ -z "${HF_TOKEN:-}" ]]; then
  echo "ERROR: HF_TOKEN is not set." >&2
  echo "Export it in ~/.zshrc, or run: HF_TOKEN=hf_xxx ./run_benchmarks.sh" >&2
  exit 1
fi

# ----------------------------- model table -----------------------------------
# Fields: alias | hf_repo[:quant_tag] | placement | note
#   placement: gpu0      -> single card (CUDA_VISIBLE_DEVICES=0)
#              dual      -> both cards, -sm layer
#              cpumoe:N  -> both cards, --n-cpu-moe N
#              cpumoe:ALL-> both cards, --cpu-moe
MODELS=(
  "gpt-oss-120b|ggml-org/gpt-oss-120b-GGUF|gpu0|MXFP4 native, ~60 GB, single card"
  "qwen3-coder-next-80b|unsloth/Qwen3-Coder-Next-GGUF:Q8_0|gpu0|80B-A3B, Q8 ~85 GB, single card"
  "qwen2.5-coder-32b|unsloth/Qwen2.5-Coder-32B-Instruct-GGUF:Q8_0|gpu0|32B dense, Q8 ~35 GB"
  "devstral-small-2-24b|unsloth/Devstral-Small-2507-GGUF:Q8_0|gpu0|24B dense, Q8 ~26 GB (verify repo)"
  "qwen3-235b-a22b|unsloth/Qwen3-235B-A22B-Instruct-2507-GGUF:Q4_K_M|dual|235B-A22B, Q4 ~133 GB, dual"
  "minimax-m2.7|unsloth/MiniMax-M2.7-GGUF:UD-Q4_K_XL|dual|~230B-A10B, Q4, dual (verify tag)"
  "glm-4.6-355b|unsloth/GLM-4.6-GGUF:UD-Q2_K_XL|dual|355B-A32B, Q2_K_XL ~135 GB, dual"
  "qwen3-coder-480b|unsloth/Qwen3-Coder-480B-A35B-Instruct-GGUF:UD-Q2_K_XL|dual|480B-A35B, Q2 ~180 GB, dual (tight)"
  "qwen3.6-27b|unsloth/Qwen3.6-27B-GGUF:Q8_0|gpu0|27B dense, Q8 ~30 GB, single card"
  "qwen3.6-35b-a3b|unsloth/Qwen3.6-35B-A3B-GGUF:Q8_0|gpu0|35B-A3B MoE, Q8 ~38 GB, single card"
  "hermes-4.3-36b|NousResearch/Hermes-4.3-36B-GGUF:Q8_0|gpu0|36B dense (Seed-OSS-36B), Q8 ~38 GB, single card"
  "hermes-4-70b|unsloth/Hermes-4-70B-GGUF:Q8_0|gpu0|70B dense (Llama 3.1), Q8 ~75 GB, single card"
)

# ----------------------------- helpers ----------------------------------------
md() { printf '%s\n' "$*" >> "$RESULTS_MD"; }
hr() { printf '\n' >> "$RESULTS_MD"; }

# Make sure a stray smi logger never outlives the harness (Ctrl-C etc.).
SMI_PID=""
cleanup() {
  if [[ -n "$SMI_PID" ]]; then
    kill "$SMI_PID" >/dev/null 2>&1
    wait "$SMI_PID" 2>/dev/null
    SMI_PID=""
  fi
}
trap cleanup EXIT INT TERM

set_power_limit() {
  local ok=1
  for i in 0 1; do
    sudo nvidia-smi -i "$i" -pm 1                    >/dev/null 2>&1 || ok=0
    sudo nvidia-smi -i "$i" -pl "$GPU_POWER_LIMIT_W" >/dev/null 2>&1 || ok=0
  done
  if [[ "$ok" -eq 1 ]]; then
    echo "Power limit set to ${GPU_POWER_LIMIT_W} W on GPU 0 and GPU 1; persistence on."
    POWER_NOTE="${GPU_POWER_LIMIT_W} W (set via nvidia-smi -pl; persistence mode on)"
  else
    echo "WARNING: could not set power limit (sudo/permission). Continuing." >&2
    POWER_NOTE="NOT SET (sudo unavailable) — GPUs at current/default limit"
  fi
}

# $1 = csv path, $2 = comma-separated GPU indices to watch.
# Scoping the logger to the run's own GPUs (-i) keeps an idle or still-cooling
# neighbour card from a previous run out of this run's data.
start_smi_logger() {
  nvidia-smi -i "$2" \
    --query-gpu=timestamp,index,temperature.gpu,temperature.memory,power.draw,clocks.sm,clocks.mem,utilization.gpu,utilization.memory,clocks_throttle_reasons.active,memory.used \
    --format=csv,noheader,nounits -lms "$SMI_INTERVAL_MS" > "$1" 2>/dev/null &
  echo $!
}

# Emit per-GPU max temp / mem-temp / power / VRAM + observed throttle reasons.
summarize_thermals() {  # $1 = csv path
  awk -F', ' '
    { idx=$2; tg=$3; tm=$4; pw=$5; thr=$10; vr=$11;
      if (tg ~ /^[0-9.]+$/ && tg+0 > mt[idx]) mt[idx]=tg+0;
      if (tm ~ /^[0-9.]+$/ && tm+0 > mm[idx]) mm[idx]=tm+0;
      if (pw ~ /^[0-9.]+$/ && pw+0 > mp[idx]) mp[idx]=pw+0;
      if (vr ~ /^[0-9.]+$/ && vr+0 > mv[idx]) mv[idx]=vr+0;
      if (thr != "" && thr != "0x0000000000000000" && thr !~ /N\/A/) seen[idx]=seen[idx] thr " ";
      idxs[idx]=1 }
    END { for (i in idxs) {
            t=(i in mm)?sprintf("%.0f",mm[i]):"N/A";
            v=(i in mv)?sprintf("%.1f",mv[i]/1024):"N/A";
            s=(i in seen)?seen[i]:"none";
            printf "%s,%.0f,%s,%.1f,%s,%s\n", i, mt[i], t, mp[i], v, s } }
  ' "$1" | sort -t',' -k1,1n
}

# ----------------------------- model prefetch ---------------------------------
# Downloads all model GGUFs into the HF cache before benchmarking starts, so
# progress bars are visible in the terminal instead of being swallowed by log
# redirection.  Works with the standalone 'hf' CLI (https://hf.co/cli/install.sh)
# or the Python 'huggingface-cli' (pip install huggingface-hub); auto-detects both.
#
# llama.cpp's -hf flag and the hf/huggingface-cli tools share the same cache
# ($HF_HOME/hub, default ~/.cache/huggingface/hub), so anything fetched here
# is found immediately when each benchmark run begins.
prefetch_models() {
  # Accept both names: 'hf' (standalone installer) and 'huggingface-cli' (pip).
  local HF_CMD=""
  if   command -v hf                &>/dev/null; then HF_CMD="hf"
  elif command -v huggingface-cli   &>/dev/null; then HF_CMD="huggingface-cli"
  else
    echo "NOTE: neither 'hf' nor 'huggingface-cli' found — skipping pre-download."
    echo "      Install: curl -LsSf https://hf.co/cli/install.sh | bash"
    echo "      Without it, models download silently as each benchmark starts."
    echo
    return 0
  fi

  echo ">>> Pre-downloading model weights  (using: ${HF_CMD})"
  echo "    Cached files resolve immediately; new files show progress below."
  echo

  for entry in "${MODELS[@]}"; do
    IFS='|' read -r alias hf placement note <<< "$entry"

    if [[ -n "$ONLY" ]] && [[ " $ONLY " != *" $alias "* ]]; then
      continue
    fi

    # Split "repo/name:QUANT_TAG" -> HF repo path and a glob for the GGUF(s).
    # If no colon is present, download all GGUFs in the repo.
    local repo="${hf%%:*}"
    local qtag="${hf#*:}"
    [[ "$qtag" == "$hf" ]] && qtag=""          # no colon — no quant tag

    local include="*.gguf"
    [[ -n "$qtag" ]] && include="*${qtag}*.gguf"

    echo "--- ${alias}  (${repo},  ${include})"
    "$HF_CMD" download "$repo" --include "$include" 2>&1
    echo
  done

  echo ">>> Pre-download complete.  Starting benchmarks."
  echo
}

# ----------------------------- system info ------------------------------------
set_power_limit

CPU_MODEL="$(lscpu 2>/dev/null | sed -n 's/^Model name:[[:space:]]*//p' | head -1)"
DRIVER="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -1)"
GPU_NAME="$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)"
GPU_COUNT="$(nvidia-smi -L | wc -l)"
LLAMA_VER="$(llama-cli --version 2>&1 | head -3 | tr '\n' ' ' || echo 'unknown')"
NOW="$(date '+%Y-%m-%d %H:%M:%S %Z')"

# RAM: read MemTotal from /proc/meminfo and round UP to the nearest power-of-2
# GiB. That matches the marketed DIMM capacity (e.g. 128 GiB); the OS reports
# slightly less because some memory is reserved by firmware/hardware.
_RAM_KIB="$(awk '/MemTotal/{print $2}' /proc/meminfo 2>/dev/null)"
if [[ -n "$_RAM_KIB" ]]; then
  _RAM_GIB=$(( _RAM_KIB / 1048576 ))
  _RAM_POW2=1
  while (( _RAM_POW2 < _RAM_GIB )); do (( _RAM_POW2 *= 2 )); done
  RAM_TOTAL="${_RAM_POW2} GiB (OS-visible: ${_RAM_GIB} GiB)"
else
  RAM_TOTAL="unknown"
fi

# ----------------------------- report header ----------------------------------
: > "$RESULTS_MD"
printf '# alias\thf_spec\tplacement\tnote\tdevices\n' > "$MANIFEST"
md "# llama.cpp Benchmark Results"
hr
md "Generated: ${NOW}"
hr
md "## System"
hr
md "| Component | Value |"
md "|---|---|"
md "| GPUs | ${GPU_COUNT} x ${GPU_NAME} |"
md "| GPU driver | ${DRIVER} |"
md "| GPU power limit | ${POWER_NOTE} |"
md "| CPU | ${CPU_MODEL:-unknown} |"
md "| System RAM | ${RAM_TOTAL} |"
md "| llama.cpp build | ${LLAMA_VER} |"
hr
md "## Global benchmark configuration"
hr
md "These settings are identical across every model unless noted in a model section."
hr
md "| Parameter | Value |"
md "|---|---|"
md "| GPU power limit (both GPUs) | ${GPU_POWER_LIMIT_W} W |"
md "| Flash attention | on (\`-fa 1\`) |"
md "| KV cache type (K and V) | ${KV_TYPE} (\`-ctk ${KV_TYPE} -ctv ${KV_TYPE}\`) |"
md "| CPU threads | ${THREADS} (\`-t ${THREADS}\`) |"
md "| GPU layers | all (\`-ngl 999\`) |"
md "| Batch / micro-batch | llama-bench defaults (\`-b 2048 -ub 512\`) |"
md "| Multi-GPU split (dual) | layer (\`-sm layer\`) |"
md "| Repetitions | ${REPS} (\`-r ${REPS}\`) |"
md "| Cool-down between models | ${COOLDOWN_S} s |"
md "| nvidia-smi sampling | every ${SMI_INTERVAL_MS} ms, scoped to the run's GPUs |"
hr
md "Sweep matrices by placement:"
md "- single card / dual card: \`-p 512,4096 -n 128 -d 0,16384\`"
hr
md "Column key: \`pp\` = prefill (tokens/sec), \`tg\` = decode (tokens/sec);"
md "\`@dN\` = measurement taken with N tokens already in context."
hr
md "Note on power readings: \`power.draw\` samples are instantaneous; NVIDIA's"
md "power controller regulates a time-averaged budget, so individual samples a"
md "few percent above the configured limit are normal transients, not faults."
hr
md "---"
hr

# ----------------------------- bench loop -------------------------------------
prefetch_models

declare -a SUMMARY_TABLES=()

run_one() {
  local alias="$1" hf="$2" placement="$3" note="$4"
  local dev sm ncmoe pp nn dd reps
  local log smicsv tmpmd

  log="${LOGDIR}/${alias}.log"
  smicsv="${LOGDIR}/${alias}.nvidia-smi.csv"
  tmpmd="${LOGDIR}/${alias}.table.md"

  case "$placement" in
    gpu0)       dev="0";   sm="";          ncmoe="";                                 pp="512,4096"; nn="128"; dd="0,16384"; reps="$REPS" ;;
    dual)       dev="0,1"; sm="-sm layer"; ncmoe="";                                 pp="512,4096"; nn="128"; dd="0,16384"; reps="$REPS" ;;
    cpumoe:ALL) dev="0,1"; sm="-sm layer"; ncmoe="--cpu-moe";                        pp="512";      nn="64";  dd="0";       reps="3" ;;
    cpumoe:*)   dev="0,1"; sm="-sm layer"; ncmoe="--n-cpu-moe ${placement#cpumoe:}"; pp="512";      nn="64";  dd="0";       reps="3" ;;
    *) echo "  unknown placement '$placement' for $alias, skipping" >&2; return 1 ;;
  esac

  echo ">>> Benchmarking ${alias} (${hf}) [${placement}]"

  printf '%s\t%s\t%s\t%s\t%s\n' "$alias" "$hf" "$placement" "$note" "$dev" >> "$MANIFEST"

  local -a cmd=( llama-bench -hf "$hf" -ngl 999 )
  [[ -n "$sm"    ]] && cmd+=( $sm )
  [[ -n "$ncmoe" ]] && cmd+=( $ncmoe )
  cmd+=( -fa 1 -ctk "$KV_TYPE" -ctv "$KV_TYPE" -t "$THREADS"
         -p "$pp" -n "$nn" -d "$dd" -r "$reps" -o md )
  local printable="CUDA_VISIBLE_DEVICES=${dev} ${cmd[*]}"

  md "## ${alias}"
  hr
  md "- Model spec: \`-hf ${hf}\`"
  md "- Placement: ${placement}"
  md "- Note: ${note}"
  hr
  md "Parameters for this run:"
  hr
  md "| Parameter | Value |"
  md "|---|---|"
  md "| GPU power limit | ${GPU_POWER_LIMIT_W} W (both GPUs) |"
  md "| Visible GPUs | ${dev} |"
  md "| Split mode | ${sm:-n/a (single GPU)} |"
  md "| CPU-MoE offload | ${ncmoe:-none} |"
  md "| Flash attention | on |"
  md "| KV cache (K/V) | ${KV_TYPE} |"
  md "| Threads | ${THREADS} |"
  md "| Prompt sizes | ${pp} |"
  md "| Gen tokens | ${nn} |"
  md "| Depths | ${dd} |"
  md "| Repetitions | ${reps} |"
  hr
  md "Exact command:"
  hr
  md '```bash'
  md "${printable}"
  md '```'
  hr

  local rc
  SMI_PID="$(start_smi_logger "$smicsv" "$dev")"
  CUDA_VISIBLE_DEVICES="$dev" "${cmd[@]}" >"$tmpmd" 2>"$log"
  rc=$?
  cleanup

  if [[ $rc -ne 0 || ! -s "$tmpmd" ]]; then
    md "**Benchmark FAILED** (exit ${rc}). Common causes: repo/quant tag changed,"
    md "out-of-memory, or unsupported flag in this build. Last log lines:"
    hr
    md '```'
    tail -n 15 "$log" >> "$RESULTS_MD"
    md '```'
    hr
    md "---"
    hr
    return 1
  fi

  md "### llama-bench results"
  hr
  cat "$tmpmd" >> "$RESULTS_MD"
  hr
  SUMMARY_TABLES+=( "${alias}::${tmpmd}" )

  md "### Thermals and power (sampled during the run)"
  hr
  md "| GPU | Max core temp (C) | Max mem temp (C) | Max power draw (W) | Max VRAM used (GiB) | Throttle reasons observed |"
  md "|---|---|---|---|---|---|"
  while IFS=',' read -r idx maxt maxm maxp maxv thr; do
    [[ -z "$idx" ]] && continue
    md "| ${idx} | ${maxt} | ${maxm} | ${maxp} | ${maxv} | ${thr} |"
  done < <(summarize_thermals "$smicsv")
  hr
  md "Logs: \`bench_logs/${GPU_POWER_LIMIT_W}/${alias}.log\` | thermals: \`bench_logs/${GPU_POWER_LIMIT_W}/${alias}.nvidia-smi.csv\`"
  hr
  md "---"
  hr
  echo "    done: ${alias}"
}

FIRST_RUN=1
for entry in "${MODELS[@]}"; do
  IFS='|' read -r alias hf placement note <<< "$entry"

  if [[ -n "$ONLY" ]] && [[ " $ONLY " != *" $alias "* ]]; then
    continue
  fi
  # Cool-down between models: lets temperatures settle so each run starts
  # from a comparable thermal baseline instead of inheriting the previous
  # run's heat soak.
  if [[ "$FIRST_RUN" -eq 0 && "$COOLDOWN_S" -gt 0 ]]; then
    echo "    cooling down ${COOLDOWN_S}s..."
    sleep "$COOLDOWN_S"
  fi
  FIRST_RUN=0
  run_one "$alias" "$hf" "$placement" "$note"
done

# ----------------------------- summary table ----------------------------------
md "## Summary (tokens/sec)"
hr
md "Extracted from each model's llama-bench table. \`pp\` = prefill, \`tg\` = decode."
hr
md "| Model | Test | tokens/sec |"
md "|---|---|---|"
for item in "${SUMMARY_TABLES[@]}"; do
  alias="${item%%::*}"; tbl="${item##*::}"
  awk -F'|' -v a="$alias" '
    {
      test=""; ts=""; lastnum="";
      for (i=1;i<=NF;i++){
        f=$i; gsub(/^[ \t]+|[ \t]+$/,"",f);
        if (f ~ /^(pp|tg)[0-9]/) test=f;
        if (f ~ /±/) ts=f;
        else if (f ~ /^[0-9]+\.[0-9]+$/) lastnum=f;
      }
      if (ts=="") ts=lastnum;
      if (test != "" && ts != "") printf "| %s | %s | %s |\n", a, test, ts;
    }' "$tbl" >> "$RESULTS_MD"
done
hr
md "Notes on interpretation:"
md "- Decode (\`tg\`) is bounded by active parameters and memory bandwidth; MoE models"
md "  decode far faster than dense models of similar total size."
md "- The second GPU adds capacity and context; do not expect dual-card \`tg\` to be ~2x"
md "  a single card on this PCIe (no-NVLink) system."
md "- Single-card (placement \`gpu0\`) and dual-card runs use different hardware budgets;"
md "  their absolute throughput numbers are not directly comparable."
md "- If \`Max mem temp\` reads \`N/A\`, the driver/board did not expose the GDDR7 sensor."
md "- A throttle reason other than \`none\` means the card hit a limit during the run."
md "  \`SW power cap\` is expected whenever the power limit is set below the card's max."
hr

# ----------------------------- HTML report ------------------------------------
REPORT_GEN="${SCRIPT_DIR}/make_report_html.py"
if [[ -n "$NO_HTML" ]]; then
  echo "NO_HTML set — skipping HTML report. Generate later with:"
  echo "  python3 ${REPORT_GEN} --logdir ${LOGDIR} --out ${REPORT_HTML}"
elif [[ ! -f "$REPORT_GEN" ]]; then
  echo "WARNING: ${REPORT_GEN} not found — skipping HTML report." >&2
elif ! command -v python3 &>/dev/null; then
  echo "WARNING: python3 not found — skipping HTML report." >&2
else
  echo
  echo ">>> Generating HTML report"
  if python3 "$REPORT_GEN" --logdir "$LOGDIR" --out "$REPORT_HTML"; then
    echo "  HTML report: ${REPORT_HTML}"
  else
    echo "WARNING: HTML report generation failed; markdown results are intact." >&2
  fi
fi

# ----------------------------- finish -----------------------------------------
DEFAULT_PL="$(nvidia-smi --query-gpu=power.default_limit --format=csv,noheader,nounits | head -1 | tr -d ' ')"
echo
echo "Done."
echo "  Report:  ${RESULTS_MD}"
echo "  HTML:    ${REPORT_HTML}"
echo "  Logs:    ${LOGDIR}/"
echo
echo "GPUs are left at the ${GPU_POWER_LIMIT_W} W limit."
echo "To restore the factory default (${DEFAULT_PL:-600} W):"
echo "  sudo nvidia-smi -i 0 -pl ${DEFAULT_PL:-600}"
echo "  sudo nvidia-smi -i 1 -pl ${DEFAULT_PL:-600}"