#!/usr/bin/env bash
# mi50_power_context_sweep_v2.sh
#
# Corrected sweep. Fixes from v1:
#   - -ngl 99 (NOT 999; 999 broke offload and produced garbage sub-1 tok/s)
#   - large prefill sizes (up to 32768) and LONG generation (2048 tok) so
#     decode reaches steady state instead of measuring 128-token startup noise
#   - drops draw_w / efficiency entirely: these MI50s report
#     "get_power_avg: Not supported", so actual watts are unmeasurable.
#     We reason about the SET cap vs throughput retention only.
#   - multiple reps with min/max/mean so variance is visible
#
# Power CAPPING works and is confirmed settable. Power READING does not.
#
# Set MODEL if path differs. Run in tmux. Uses llama-bench (verified: at
# stock power, -ngl 99 gives pp512~472 t/s, tg128~64 t/s on this rig).

set -uo pipefail

# ------------------------------- config -------------------------------
MODEL="${MODEL:-$HOME/.cache/llama.cpp/openai_gpt-oss-120b-Q4_K_M/openai_gpt-oss-120b-Q4_K_M-00001-of-00002.gguf}"
NGL=99
GPUS=(0 1)

POWER_LEVELS=(225 200 180 160)

# prefill sizes (compute-bound) and a long generation (decode steady-state)
PREFILL_SIZES=(512 8192 32768)
GEN_TOKENS=1536

REPS=2
COOLDOWN=15

# hard wall-clock budget: the sweep will stop launching new cells once
# this many seconds have elapsed, so the whole run fits your time box.
TIME_BUDGET_S=5400          # 90 minutes
START_EPOCH=$(date +%s)

LOGFILE="mi50_pxc_v2_$(date +%Y%m%d_%H%M%S).csv"
# ----------------------------------------------------------------------

command -v rocm-smi   >/dev/null || { echo "rocm-smi not found"; exit 1; }
command -v llama-bench >/dev/null || { echo "llama-bench not found in PATH (sudo drops PATH; run: sudo env \"PATH=\$PATH\" MODEL=... ./thisscript)"; exit 1; }
[[ -f "$MODEL" ]] || { echo "Model not found: $MODEL"; exit 1; }

echo "power_w,ctx_prompt,gen_tokens,pp_mean,pp_min,pp_max,tg_mean,tg_min,tg_max,edge_temp_c" > "$LOGFILE"
echo "Logging to $LOGFILE"
echo "Model: $MODEL"
echo "NOTE: draw_w is not collected; these MI50s do not expose a power sensor."

restore_power() {
  echo ""
  echo "Restoring stock (225W) power on all GPUs..."
  for g in "${GPUS[@]}"; do
    rocm-smi -d "$g" --resetpoweroverdrive >/dev/null 2>&1 || true
  done
}
trap restore_power EXIT INT TERM

grab_float() { grep -oE '[0-9]+\.[0-9]+' | head -1; }

# llama-bench csv: the t/s value is the final field of each result row.
# Rows are tagged by test name (pp<N>, tg<N>). Pull by tag, positional fallback.
extract_toks() {
  local out="$1" tag="$2"
  local v
  v=$(echo "$out" | grep -iE "(^|[\",])${tag}([\",]|$)" | tail -1 \
        | awk -F',' '{gsub(/"/,"",$NF); print $NF}')
  echo "$v"
}

sample_temp() { rocm-smi -d 0 --showtemp 2>/dev/null | grep -i edge | grab_float; }

# min/max/mean over a space-separated list
stats() {
  awk '{
    n=NF; if(n==0){print "NA NA NA"; exit}
    mn=$1; mx=$1; s=0;
    for(i=1;i<=n;i++){ s+=$i; if($i<mn)mn=$i; if($i>mx)mx=$i }
    printf "%.2f %.2f %.2f", s/n, mn, mx
  }'
}

total=$(( ${#POWER_LEVELS[@]} * ${#PREFILL_SIZES[@]} ))
cell=0

for pw in "${POWER_LEVELS[@]}"; do
  # stop launching new power levels if we are out of time budget
  ELAPSED=$(( $(date +%s) - START_EPOCH ))
  if (( ELAPSED >= TIME_BUDGET_S )); then
    echo ""
    echo "### time budget (${TIME_BUDGET_S}s) reached, stopping before ${pw}W ###"
    break
  fi

  echo ""
  echo "############ POWER CAP = ${pw}W ############"
  set_ok=1
  for g in "${GPUS[@]}"; do
    if ! rocm-smi -d "$g" --setpoweroverdrive "$pw" >/dev/null 2>&1; then
      echo "  WARN: failed to set ${pw}W on GPU $g"
      set_ok=0
    fi
  done
  # verify the cap actually applied, per GPU, via showmaxpower (clean value)
  cap_ok=1
  for g in "${GPUS[@]}"; do
    CAP=$(rocm-smi -d "$g" --showmaxpower 2>/dev/null \
            | grep -i "Package Power" | grab_float)
    echo "  GPU$g cap now: ${CAP:-unknown}W"
    # compare as integers
    if [[ -z "${CAP:-}" ]] || (( ${CAP%.*} != pw )); then
      cap_ok=0
    fi
  done
  if (( cap_ok == 0 )); then
    echo "  ERROR: cap did not apply as expected on all GPUs; skipping ${pw}W"
    continue
  fi
  sleep 4

  for pp in "${PREFILL_SIZES[@]}"; do
    ELAPSED=$(( $(date +%s) - START_EPOCH ))
    if (( ELAPSED >= TIME_BUDGET_S )); then
      echo "  ### time budget reached mid-sweep, stopping ###"
      break 2
    fi
    cell=$((cell+1))
    echo ""
    echo "--- [$cell/$total] cap=${pw}W prefill=${pp} gen=${GEN_TOKENS} (x${REPS}) [${ELAPSED}s elapsed] ---"

    pp_vals=""; tg_vals=""
    for r in $(seq 1 "$REPS"); do
      OUT=$(llama-bench -m "$MODEL" -ngl "$NGL" -p "$pp" -n "$GEN_TOKENS" -o csv 2>/dev/null)
      PP=$(extract_toks "$OUT" "pp${pp}")
      TG=$(extract_toks "$OUT" "tg${GEN_TOKENS}")
      [[ -z "${PP:-}" ]] && PP=$(echo "$OUT" | awk -F',' 'NR==2{gsub(/"/,"",$NF);print $NF}')
      [[ -z "${TG:-}" ]] && TG=$(echo "$OUT" | awk -F',' 'NR==3{gsub(/"/,"",$NF);print $NF}')
      if [[ -n "${PP:-}" && -n "${TG:-}" ]]; then
        pp_vals="$pp_vals $PP"; tg_vals="$tg_vals $TG"
        echo "    rep $r: pp=${PP} tg=${TG}"
      else
        echo "    rep $r: PARSE MISS (pp=${PP:-NA} tg=${TG:-NA})"
      fi
      sleep "$COOLDOWN"
    done

    read PP_MEAN PP_MIN PP_MAX < <(echo "$pp_vals" | stats)
    read TG_MEAN TG_MIN TG_MAX < <(echo "$tg_vals" | stats)
    TEMP=$(sample_temp)

    echo "    => pp ${PP_MEAN} (${PP_MIN}-${PP_MAX}) | tg ${TG_MEAN} (${TG_MIN}-${TG_MAX}) | ${TEMP:-NA}C"
    echo "${pw},${pp},${GEN_TOKENS},${PP_MEAN},${PP_MIN},${PP_MAX},${TG_MEAN},${TG_MIN},${TG_MAX},${TEMP:-NA}" >> "$LOGFILE"
  done
done

echo ""
echo "======================================================================"
echo "Sweep complete: $LOGFILE"
echo ""
column -t -s',' "$LOGFILE"
echo ""
echo "Read tg (decode) down each prefill column: if it holds flat as the cap"
echo "drops, decode is bandwidth-bound and the low cap is free. Read pp"
echo "(prefill) the same way: expect it to fall off at low caps, worse at"
echo "large prefill. draw_w intentionally absent (no sensor on these cards)."