#!/usr/bin/env bash

# Run llama-benchy against the local vLLM server.
#
#   benchmarks/bench.sh [PP] [TG] [CONCURRENCY] [DEPTH] [extra llama-benchy args]
#
#   PP            prompt tokens, default 32768 (use 1 for decode-only)
#   TG            generated tokens, default 2048 (use 1 for prefill-only)
#   CONCURRENCY   one level, or a space-separated list, default 1
#   DEPTH         prior-conversation tokens, one value or a list, default 0
#
# PP, TG, CONCURRENCY, and DEPTH also read the environment when their argument is absent.
#
# Examples:
#   benchmarks/bench.sh                          # pp=32768 tg=2048 conc=1 depth=0
#   benchmarks/bench.sh 32768 1                  # prefill-only
#   benchmarks/bench.sh 1 4096                   # decode-only
#   benchmarks/bench.sh 32768 2048 "1 2 4"       # concurrency sweep
#   benchmarks/bench.sh 128 256 1 "0 8192 65536" # depth sweep, small turn
#   benchmarks/bench.sh 1 512 1 --runs 5         # extra args go to llama-benchy
#   benchmarks/bench.sh 128 256 1 8192 --enable-prefix-caching
#                                                # two-phase run, the tool measures
#                                                # the cold prefix load and the warm
#                                                # turn as separate rows

set -euo pipefail

BASE_URL="${BASE_URL:-http://127.0.0.1:8080/v1}"
MODEL="${MODEL:-Hinata}"
TOKENIZER="${TOKENIZER:-/data/models/halt95/Swift1.5-Qwen3.8-Flash-Next-W4A16-Merlin}"

# The first four arguments are PP, TG, CONCURRENCY, and DEPTH. Parsing stops at the first
# argument that starts with `-`, so that argument and the rest go to llama-benchy.
names=(PP TG CON DEPTH)

for i in 0 1 2 3; do
  [ $# -gt 0 ] || break
  case "$1" in -*) break ;; esac
  printf -v "${names[$i]}" '%s' "$1"
  shift
done

PP="${PP:-32768}"; TG="${TG:-2048}"; CON="${CON:-1}"; DEPTH="${DEPTH:-0}"

for v in PP TG; do
  case "${!v}" in ''|*[!0-9]*) echo "$v must be a whole number, got '${!v}'" >&2; exit 2 ;; esac
done

for v in CON DEPTH; do
  for x in ${!v}; do
    case "$x" in ''|*[!0-9]*) echo "$v must be whole numbers, got '${!v}'" >&2; exit 2 ;; esac
  done
done

command -v llama-benchy >/dev/null || {
  echo "llama-benchy not found. Install it with: uv tool install git+https://github.com/eugr/llama-benchy" >&2
  exit 1
}

# shellcheck disable=SC2086  # PP, TG, CON, DEPTH split on spaces on purpose: benchy takes lists
exec llama-benchy \
  --base-url "$BASE_URL" \
  --model "$MODEL" \
  --latency-mode api \
  --tokenizer "$TOKENIZER" \
  --pp $PP --tg $TG \
  --no-cache \
  --exact-tg \
  --concurrency $CON \
  --depth $DEPTH \
  "$@"