#!/bin/sh
#
# RX 580 prompt-processing benchmark orchestrator. Runs ON THE POD.
#
# Stops llama-swap, waits for the GPU to go idle, runs the real-corpus
# prompt-processing harness and the llama-bench sweep for every arm, appends
# parseable results, and ALWAYS restarts llama-swap on exit (success, error
# or Ctrl-C).
#
# Usage:
#   bench.sh --build DIR --label NAME
#            [--build-b DIR --label-b NAME]   second arm, interleaved
#            [--arm DIR:LABEL ...]            extra arms, repeatable (N-way)
#            [--rounds N]                     default 3 multi-arm, 1 otherwise
#            [--reps N]                       corpus repeats per size, default 3
#            [--sizes 4096,16384,32768]
#            [--corpus-only | --lb-only]
#            [--results PATH]                 default /root/bench/results.txt
#            [--keep-swap-down]               do not restart llama-swap at exit
#
set -e

BENCH_DIR=/root/bench
SUPCONF=/root/supervisord.conf
RESULTS="$BENCH_DIR/results.txt"
MODEL=/root/.cache/huggingface/hub/models--unsloth--Qwen3.6-35B-A3B-GGUF/snapshots/a483e9e6cbd595906af30beda3187c2663a1118c/Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf

ARMS=""            # space separated DIR|LABEL entries
BUILD_A=""; LABEL_A=""
BUILD_B=""; LABEL_B=""
ROUNDS=""
REPS=3
SIZES=4096,16384,32768
DO_CORPUS=1
DO_LB=1
KEEP_SWAP_DOWN=0

while [ $# -gt 0 ]; do
  case "$1" in
    --build)     BUILD_A="$2"; shift 2 ;;
    --label)     LABEL_A="$2"; shift 2 ;;
    --build-b)   BUILD_B="$2"; shift 2 ;;
    --label-b)   LABEL_B="$2"; shift 2 ;;
    --arm)       ARMS="$ARMS ${2%%:*}|${2#*:}"; shift 2 ;;
    --rounds)    ROUNDS="$2"; shift 2 ;;
    --reps)      REPS="$2"; shift 2 ;;
    --sizes)     SIZES="$2"; shift 2 ;;
    --results)   RESULTS="$2"; shift 2 ;;
    --corpus-only) DO_LB=0; shift ;;
    --lb-only)     DO_CORPUS=0; shift ;;
    --keep-swap-down) KEEP_SWAP_DOWN=1; shift ;;
    -h|--help)   sed -n '2,20p' "$0"; exit 0 ;;
    *) echo "unknown arg: $1" >&2; exit 2 ;;
  esac
done

if [ -n "$BUILD_A" ] && [ -z "$LABEL_A" ]; then
  echo "error: --build needs --label" >&2; exit 2
fi
if [ -n "$BUILD_B" ] && [ -z "$LABEL_B" ]; then
  echo "error: --build-b needs --label-b" >&2; exit 2
fi
# --build/--label and --build-b/--label-b are sugar for the first two arms.
if [ -n "$BUILD_A" ]; then ARMS="$BUILD_A|$LABEL_A $ARMS"; fi
if [ -n "$BUILD_B" ]; then ARMS="$ARMS $BUILD_B|$LABEL_B"; fi
ARMS=$(echo $ARMS)
if [ -z "$ARMS" ]; then
  echo "error: need at least one arm (--build/--label or --arm DIR:LABEL)" >&2
  exit 2
fi
NARMS=$(echo "$ARMS" | wc -w)
if [ -z "$ROUNDS" ]; then
  if [ "$NARMS" -gt 1 ]; then ROUNDS=3; else ROUNDS=1; fi
fi
for e in $ARMS; do
  d=${e%%|*}
  if [ ! -x "$d/bin/llama-server" ] || [ ! -x "$d/bin/llama-bench" ]; then
    echo "error: $d lacks bin/llama-server or bin/llama-bench" >&2
    exit 2
  fi
done

cat <<'WARN'
-------------------------------------------------------------------------
RX 580 benchmark. Read this before trusting any number.

  * Cross-invocation variance on this box is about 7 percent. llama-bench
    within-run error bars understate it by roughly 14x. Any cross-build
    delta under about 5 percent is UNPROVEN noise.
  * Comparisons you care about must live inside ONE llama-bench invocation
    (comma-separated sweeps), or be interleaved across at least 3 rounds.
    That is what multiple arms plus --rounds does.
  * Nothing else may touch the GPU while this runs. llama-swap is stopped
    for the duration and restarted on exit, including on error or Ctrl-C.
  * Traps: -tb is NOT a llama-bench flag. A ubatch larger than -p never
    fills, so prompt lengths must be multiples of the 2048 ubatch. A bad
    flag makes llama-bench print usage and exit silently mid-sweep.
-------------------------------------------------------------------------
WARN

echo "[bench] $NARMS arm(s), $ROUNDS round(s): $ARMS"

restore_swap() {
  rc=$?
  if [ "$KEEP_SWAP_DOWN" -eq 0 ]; then
    echo ""
    echo "[bench] restarting llama-swap"
    supervisorctl -c "$SUPCONF" start llama-swap || true
    supervisorctl -c "$SUPCONF" status llama-swap || true
  else
    echo "[bench] --keep-swap-down: llama-swap left STOPPED"
  fi
  exit $rc
}
trap restore_swap EXIT INT TERM

echo "[bench] stopping llama-swap"
supervisorctl -c "$SUPCONF" stop llama-swap || true

wait_gpu_idle() {
  i=0
  while [ $i -lt 60 ]; do
    if ! pgrep -f 'bin/llama-server|bin/llama-bench|bin/llama-cli' >/dev/null 2>&1; then
      return 0
    fi
    sleep 2
    i=$((i+1))
  done
  echo "error: something is still on the GPU:" >&2
  pgrep -af 'bin/llama-server|bin/llama-bench|bin/llama-cli' >&2
  exit 3
}

echo "[bench] waiting for the GPU to go idle"
wait_gpu_idle
echo "[bench] GPU idle"

run_corpus() {
  [ "$DO_CORPUS" -eq 1 ] || return 0
  echo "===== corpus $2 ====="
  wait_gpu_idle
  python3 "$BENCH_DIR/ppbench.py" "$1" "$2" \
    --sizes "$SIZES" --reps "$REPS" --results "$RESULTS"
}

run_lb() {
  [ "$DO_LB" -eq 1 ] || return 0
  echo "===== llama-bench $2 ====="
  wait_gpu_idle
  RESULTS="$RESULTS" MODEL="$MODEL" sh "$BENCH_DIR/lbsweep.sh" "$1" "$2"
}

# Interleave: every arm is measured once per round, so slow drift in the
# machine hits all arms roughly equally instead of biasing whichever ran first.
r=1
while [ "$r" -le "$ROUNDS" ]; do
  echo ""
  echo "########## ROUND $r / $ROUNDS ##########"
  for e in $ARMS; do run_corpus "${e%%|*}" "${e#*|}"; done
  for e in $ARMS; do run_lb     "${e%%|*}" "${e#*|}"; done
  r=$((r+1))
done

echo ""
echo "ALL_RUNS_DONE"
echo "[bench] raw results appended to $RESULTS"
python3 "$BENCH_DIR/summarize.py" "$RESULTS" || true