#!/usr/bin/env bash
#
# bench-batch-ubatch.sh — how large may the logical batch and the micro-batch be?
#
# Sweeps llama.cpp's --batch-size against --ubatch-size with llama-bench and records, per
# cell, the prefill rate, the decode rate and the PEAK VRAM of each card. The VRAM half is
# not decoration: batch/ubatch is the knob that sizes the activation buffers, so a cell
# that looks fast because it silently fell back to a smaller allocation would otherwise be
# indistinguishable from a cell that is genuinely fast.
#
# ONE llama-bench INVOCATION PER CELL, on purpose. llama-bench calls exit(1) the moment a
# test cannot run, so a single infeasible ubatch would take every later cell with it —
# and "infeasible" is one of the things this sweep is looking for.
#
# RUN IT FROM A TRANSIENT UNIT, NOT FROM A SHELL. The deployment's unit sets
# LimitMEMLOCK=infinity because the ROCm backend pins staging buffers; an SSH shell on
# these hosts gets 8 MiB, and a load that fails on that proves nothing about the model:
#
#   sudo systemd-run --unit=fm-bench --collect --no-block --property=Type=oneshot \
#        --property=LimitMEMLOCK=infinity --property=User=dank \
#        /home/dank/serverAppConfigs/scripts/bench-batch-ubatch.sh
#
# It stops $SERVICE for the duration and starts it again on exit (trap), because a bench
# beside a live service starves the same GPU. Results land in $OUT: one directory per cell
# (bench.json, bench.log, vram.tsv) plus cells.tsv, appended as the sweep proceeds, so
# progress can be read without waiting for the end. A re-run skips every cell that already
# has a bench.json, so an interrupted sweep resumes.
#
set -uo pipefail

BIN=${BIN:-/home/dank/GitPackages/llama.cpp-dflash2/build-rdna4-rocm714/bin}
MODEL=${MODEL:-/home/dank/models/Swift-Qwen3.8-27B-Uncensored-Dynamic-MTP-GGUF/Swift-Qwen3.8-27B-Uncensored-Dynamic-MTP-UD-Q8_K_XL.gguf}
OUT=${OUT:-/home/dank/models/logs/bench-bub-$(date +%Y%m%d-%H%M%S)}

# what a cell measures: a prefill and a decode, both with a 64k KV already resident, so the
# numbers describe the agent-loop shape (prompt landing in a long conversation) rather than
# a cold start
N_PROMPT=${N_PROMPT:-65536}
N_GEN=${N_GEN:-128}
N_DEPTH=${N_DEPTH:-65536}
REPS=${REPS:-3}
BATCHES=${BATCHES:-"2048 4096 6144 8192"}
UBATCHES=${UBATCHES:-"512 768 1024 1280 1536 1792 2048"}
SEED_CELLS=${SEED_CELLS:-"6144/1536"}   # the deployed cell first: it is the reference

# how the model is loaded — must match the deployment, or the numbers describe a service
# that does not exist
NGL=${NGL:-999}
DEVICES=${DEVICES:-ROCm0/ROCm1}        # llama-bench separates with '/'; ROCm2 is the iGPU
SPLIT_MODE=${SPLIT_MODE:-layer}
TENSOR_SPLIT=${TENSOR_SPLIT:-36/29}    # slash here — the llama-server unit takes commas
FLASH_ATTN=${FLASH_ATTN:-on}
KV_TYPE=${KV_TYPE:-f16}
THREADS=${THREADS:-8}

SERVICE=${SERVICE:-frank-medium}
BOOTED=

service_stop() {
    if systemctl is-active --quiet "$SERVICE"; then
        echo "stopping $SERVICE for the sweep" >&2
        sudo systemctl stop "$SERVICE"
        BOOTED=1
    fi
}

service_start() {
    if [ -n "$BOOTED" ]; then
        echo "starting $SERVICE again" >&2
        sudo systemctl start "$SERVICE"
    fi
}

trap service_start EXIT

run_cell() {
    # two statements: `local a=1 b=$a` expands b before the builtin assigns a, and
    # set -u turns that into "b: unbound variable"
    local b=$1 ub=$2
    local dir="$OUT/b${b}_ub${ub}" sampler start rc wall
    mkdir -p "$dir"
    if [ -s "$dir/bench.json" ]; then
        echo "b=$b ub=$ub: bench.json present, skipping" >&2
        return 0
    fi
    echo "b=$b ub=$ub: running" >&2

    # 2 s sampling for the cell's whole life, load included
    (
        while :; do
            rocm-smi --showmeminfo vram --csv 2>/dev/null |
                awk -F, -v t="$(date +%s)" '/^card[01],/ { print t, $1, $2, $3 }'
            sleep 2
        done
    ) > "$dir/vram.tsv" 2>/dev/null &
    sampler=$!

    start=$(date +%s)
    "$BIN/llama-bench" -m "$MODEL" \
        -p "$N_PROMPT" -n "$N_GEN" -d "$N_DEPTH" \
        -b "$b" -ub "$ub" \
        -ngl "$NGL" -dev "$DEVICES" -sm "$SPLIT_MODE" -ts "$TENSOR_SPLIT" \
        -fa "$FLASH_ATTN" -ctk "$KV_TYPE" -ctv "$KV_TYPE" -t "$THREADS" \
        -r "$REPS" --no-warmup --progress -o json \
        > "$dir/bench.json" 2> "$dir/bench.log"
    rc=$?

    kill "$sampler" 2>/dev/null
    wait "$sampler" 2>/dev/null
    wall=$(( $(date +%s) - start ))

    if [ "$rc" != 0 ]; then
        grep -iE "error|failed|abort" "$dir/bench.log" | tail -3 >> "$OUT/failures.txt"
        echo "b=$b ub=$ub: exit $rc" >&2
    fi

    python3 - "$dir/bench.json" "$dir/vram.tsv" "$b" "$ub" "$rc" "$wall" <<'PY'
import json, sys

json_path, vram_path, b, ub, rc, wall = sys.argv[1:7]

def rates(tests, key):
    hit = [t for t in tests if t.get(key)]
    if not hit:
        return "-\t-"
    return "%.2f\t%.2f" % (hit[0]["avg_ts"], hit[0]["stddev_ts"])

try:
    tests = json.load(open(json_path))
except Exception:
    tests = []

peak = {}
for line in open(vram_path):
    parts = line.split()
    if len(parts) == 4:
        _, card, total, used = parts
        if card not in peak or int(used) > peak[card][1]:
            peak[card] = (int(total), int(used))

def free_gib(card):
    if card not in peak:
        return "-"
    total, used = peak[card]
    return "%.2f" % ((total - used) / 2**30)

row = [b, ub, rc, wall, rates(tests, "n_prompt"), rates(tests, "n_gen"),
       free_gib("card0"), free_gib("card1"), len(tests)]
print("\t".join(str(f) for f in row))
PY
}

mkdir -p "$OUT"
echo "output directory: $OUT" >&2
if [ ! -s "$OUT/cells.tsv" ]; then
    printf 'batch\tubatch\trc\twall_s\tpp_tok_s\tpp_sd\ttg_tok_s\ttg_sd\tgpu0_free_GiB\tgpu1_free_GiB\ttests\n' \
        > "$OUT/cells.tsv"
fi

service_stop

for cell in $SEED_CELLS; do
    run_cell "${cell%%/*}" "${cell##*/}" >> "$OUT/cells.tsv"
done

for b in $BATCHES; do
    for ub in $UBATCHES; do
        run_cell "$b" "$ub" >> "$OUT/cells.tsv"
    done
done

echo "sweep complete: $OUT/cells.tsv" >&2
cat "$OUT/cells.tsv" >&2