#!/usr/bin/env bash
# bench-ubatch-rocmi4.sh — sweep ubatch size at 8k context on the
# frank-default-large ROCmFPX standalone-MTP stack (dankclaw, gfx1151).
#
# Boots one llama-server instance per ubatch on a scratch port with a fixed
# -b and -c 8192, runs a standardized ~75%-of-ctx prose workload, and reads
# authoritative prefill/generation tok/s from the server's own print_timing
# lines. Steps from UBATCH_MIN to UBATCH_MAX inclusive by UBATCH_STEP.
#
# Usage:
#   ./scripts/bench-ubatch-rocmi4.sh
#   UBATCH_MIN=512 UBATCH_MAX=4096 UBATCH_STEP=512 CTX=8192 RUNS=3 \
#     ./scripts/bench-ubatch-rocmi4.sh
#
# Benchmarking policy: stop the production service first so the GPU is not
# starved. With STOPSVC=1 the script restarts frank-default-large on exit.
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ENV_FILE="${FRANK_LARGE_ENV_FILE:-$SCRIPT_DIR/frank-default-large.env}"
if [[ -f "$ENV_FILE" ]]; then
    set -a
    # shellcheck disable=SC1090
    source "$ENV_FILE"
    set +a
fi

# --- sweep parameters -------------------------------------------------------
: "${UBATCH_MIN:=256}"
: "${UBATCH_MAX:=6144}"
: "${UBATCH_STEP:=256}"
: "${UBATCH_VALUES:=256 512 768 1024 1280 1536 2048 3072 4096 5120 6144}"
: "${CTX:=8192}"                 # 8k contexts
: "${BENCH_BATCH:=$CTX}"         # fixed batch = ctx, so the full 8k prompt fits
: "${OUT_TOKENS:=256}"
: "${RUNS:=3}"                   # decode runs per ubatch
: "${PORT_BASE:=7894}"           # scratch ports 7894..
: "${STOPSVC:=1}"                # stop/restart the live unit for a clean GPU

# --- resolved paths (fall back to serve-script defaults) --------------------
: "${LLAMA_BIN:=/home/dank/ciru-qwen38/ROCmFPX/build-promptforge/bin/llama-server}"
: "${MODEL_DIR:=/home/dank/models/kingjones-Qwen3.8-27B-STRIX-MTP}"
: "${MODEL_FILE:=Qwen3.8-27B-Uncensored-Q6_0_ROCMFPX_AGENT.gguf}"
: "${MTP_DRAFT_FILE:=$MODEL_DIR/mtp-Qwen3.8-27B-Q4_0.gguf}"
: "${MMPROJ_FILE:=mmproj-Qwen3.8-27B-Q8_0.gguf}"
: "${CHAT_TEMPLATE_FILE:=$SCRIPT_DIR/chat-template.jinja}"
: "${ALIAS:=frank-default-large}"
: "${HSA_OVERRIDE_GFX_VERSION:=11.5.1}"
: "${GGML_HIP_ENABLE_UNIFIED_MEMORY:=1}"
: "${MTP_N_MAX:=16}"
: "${MTP_P_MIN:=0.60}"

MODEL="$MODEL_DIR/$MODEL_FILE"
MMPROJ="$MODEL_DIR/$MMPROJ_FILE"

for f in "$MODEL" "$CHAT_TEMPLATE_FILE"; do
    if [[ ! -f "$f" ]]; then
        echo "missing required file: $f" >&2
        exit 1
    fi
done
if [[ ! -x "$LLAMA_BIN" ]]; then
    echo "missing llama-server binary: $LLAMA_BIN" >&2
    exit 1
fi

MMPROJ_ARGS=()
if [[ -f "$MMPROJ" ]]; then
    MMPROJ_ARGS=(--mmproj "$MMPROJ")
fi

STopped=0
SCRATCH_PID=""
cleanup() {
    if [[ -n "$SCRATCH_PID" ]] && kill -0 "$SCRATCH_PID" 2>/dev/null; then
        echo "=== stopping scratch llama-server (pid $SCRATCH_PID) ==="
        kill "$SCRATCH_PID" 2>/dev/null || true
        wait "$SCRATCH_PID" 2>/dev/null || true
        SCRATCH_PID=""
    fi
    if [ "$STopped" = 1 ]; then
        echo "=== restarting frank-default-large ==="
        sudo systemctl restart frank-default-large || true
        STopped=0
    fi
}
trap cleanup EXIT

stop_live_service() {
    if [ "$STOPSVC" = 1 ]; then
        echo "=== stopping frank-default-large for a clean GPU ==="
        sudo systemctl stop frank-default-large || true
        STopped=1
    fi
}

# wait_health <pid> <port> — poll /health until ready or the server dies.
wait_health() {
    local pid="$1" port="$2" i ready=0 crashed=0
    for i in $(seq 1 900); do
        if ! kill -0 "$pid" 2>/dev/null; then crashed=1; break; fi
        if curl -sf "http://127.0.0.1:${port}/health" >/dev/null 2>&1; then
            ready=1; break
        fi
        [ $((i % 10)) = 0 ] && echo -n .
        sleep 2
    done
    [ "$crashed" = 1 ] && { echo " BACKEND_CRASHED (before ready)"; return 2; }
    [ "$ready" = 1 ] || { echo " FAIL (timeout)"; return 1; }
    echo " [ready]"
    return 0
}

# run_one <ubatch> <port> <out_tokens> <runs> — boot one server instance, run
# the workload, parse print_timing, print a RESULT line.
run_one() {
    local ubatch="$1" port="$2" out_tokens="$3" runs="$4"
    local log="/tmp/ubatch-${ubatch}-p${port}.log"
    : > "$log"

    echo "=== ubatch=${ubatch} ctx=${CTX} port=${port} runs=${runs} ==="
    "$LLAMA_BIN" \
        --model "$MODEL" \
        "${MMPROJ_ARGS[@]}" \
        --chat-template-file "$CHAT_TEMPLATE_FILE" \
        --chat-template-kwargs '{"enable_thinking":true,"preserve_thinking":true,"preserve_reasoning":true}' \
        --alias bench-ubatch \
        --host 127.0.0.1 --port "$port" \
        -dev ROCm0 -ngl 999 -np 1 \
        -t 16 -tb 32 \
        -c "$CTX" -b "$BENCH_BATCH" -ub "$ubatch" \
        -fa on -ctk f16 -ctv f16 -ctxcp 16 \
        --cache-ram 8192 --cache-prompt --metrics --no-mmap \
        --spec-type draft-mtp --model-draft "$MTP_DRAFT_FILE" \
        --spec-draft-device ROCm0 --spec-draft-ngl 999 \
        --spec-draft-type-k f16 --spec-draft-type-v f16 \
        --spec-draft-n-max "$MTP_N_MAX" --spec-draft-n-min 0 \
        --spec-draft-p-min "$MTP_P_MIN" \
        --spec-draft-backend-sampling \
        --temp 0.7 --top-p 0.95 \
        --reasoning "$REASONING" --reasoning-format "$REASONING_FORMAT" --reasoning-budget -1 \
        -lv 4 >>"$log" 2>&1 &
    local pid=$!
    SCRATCH_PID="$pid"
    wait_health "$pid" "$port" || { tail -30 "$log"; return 2; }

    # one python pass per run; reuse the same boot across runs
    local r
    for r in $(seq 1 "$runs"); do
        echo "--- run ${r}/${runs} ---"
        python3 - "$ubatch" "$port" "$out_tokens" "$CTX" "$log" "$r" <<'PY'
import sys, time, json, urllib.request
ubatch, port, out_tokens, ctx, log, run = sys.argv[1], sys.argv[2], int(sys.argv[3]), int(sys.argv[4]), sys.argv[5], sys.argv[6]

# Flowing prose at ~75% of ctx so MTP acceptance is realistic.
VAR = (
    "narrative nonfiction blends observation with reflection, ",
    "and a careful writer lets each clause earn its place ",
    "in the sentence's rising line, so meaning accumulates slowly. ",
    "The reader is not hurried; the reader is led. Loose and periodic, ",
    "the style turns on concrete nouns and verbs that carry weight, ",
    "and it resists the abstract noun that hides a vague claim. ",
    "A paragraph is not a list. It is a single turn of thought that ",
    "opens, develops, and lands. This is the craft of prose. Choose, ",
    "arrange, and above all omit. Omission is the discipline that ",
    "keeps the whole clear, and clarity is the courtesy we owe the ",
    "willing reader."
)
n = len(VAR)
parts = []
t = 0
while t * 76 < (ctx * 3) // 4:
    for k in range(5):
        parts.append(VAR[(t + k) % n])
    parts.append(f" This is cued by note {t}. ")
    t += 1
prompt = "".join(parts)
prompt += ("\n\nNow, in a detailed and flowing paragraph, continue this essay about "
           "the philosophy of language and its influence on modern computation. ")

body = json.dumps({
    "prompt": prompt, "max_tokens": out_tokens,
    "temperature": 0.0, "stream": False,
}).encode()
req = urllib.request.Request(f"http://127.0.0.1:{port}/v1/completions",
                             data=body, headers={"Content-Type": "application/json"})
t0 = time.time()
try:
    with urllib.request.urlopen(req, timeout=7200) as resp:
        out = json.load(resp)
except Exception:
    print("request failed; preserving benchmark cleanup", file=sys.stderr)
    raise
dt = time.time() - t0
pt = out.get("usage", {}).get("prompt_tokens", 0)
ct = out.get("usage", {}).get("completion_tokens", 0)

lines = open(log, encoding="utf-8", errors="ignore").read().splitlines()
prefill = [l for l in lines if "prompt eval time" in l]
evals = [l for l in lines if "eval time" in l and "prompt eval" not in l]

def grab(arr):
    return arr[-1].strip() if arr else "n/a"

g_tps = ""
if evals and ct:
    # eval timing line: "eval time = .. ms / N tokens (... tokens per second)"
    import re
    match = re.search(r"eval time =\s*([0-9.]+) ms", evals[-1])
    if match:
        ms = float(match.group(1))
        g_tps = f"{ct / (ms / 1000):.2f}"
print(f"== RESULT ubatch={ubatch} run={run} "
      f"PREFILL={grab(prefill)} GEN_TPS={g_tps} "
      f"pt={pt} ct={ct} wall={dt:.1f}s")
PY
    done
    kill "$pid" 2>/dev/null || true
    wait "$pid" 2>/dev/null || true
    SCRATCH_PID=""
}

main() {
    stop_live_service
    local idx=0 ubatch port
    printf "%-10s %-14s %-10s\n" "UBATCH" "PREFILL" "GEN_TPS"
    for ubatch in $UBATCH_VALUES; do
        port=$((PORT_BASE + idx))
        run_one "$ubatch" "$port" "$OUT_TOKENS" "$RUNS"
        idx=$((idx + 1))
    done
    echo "=== sweep complete ==="
}

main "$@"