#!/bin/sh
#
# llama-bench sweep, the controlled cross-check for the corpus harness.
# Runs ON THE POD. Normally invoked by bench.sh.
#
# Usage: lbsweep.sh <build_dir> <label>
# Env overrides: RESULTS, MODEL, LB_P, LB_N, LB_R
#
# Everything that must be compared lives inside ONE llama-bench invocation
# (-p takes a comma-separated list), because cross-invocation variance on this
# box is about 7 percent while the within-run error bars are about 0.3 percent.
#
# Flag traps, learned the hard way:
#   * -tb is NOT a llama-bench flag. It exists on llama-server only.
#   * -ncmoe is the llama-bench spelling of --n-cpu-moe.
#   * -mmp 0 is the llama-bench equivalent of the server's --no-mmap.
#   * A bad flag makes llama-bench print usage and exit silently, which looks
#     exactly like a sweep that produced no rows. Always check row count.
#   * A ubatch larger than -p never fills, so keep -p a multiple of -ub.
#
set -e
BUILD_DIR="$1"
LABEL="$2"
if [ -z "$BUILD_DIR" ] || [ -z "$LABEL" ]; then
  echo "usage: $0 <build_dir> <label>" >&2
  exit 2
fi

RESULTS="${RESULTS:-/root/bench/results.txt}"
MODEL="${MODEL:-/root/.cache/huggingface/hub/models--unsloth--Qwen3.6-35B-A3B-GGUF/snapshots/a483e9e6cbd595906af30beda3187c2663a1118c/Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf}"
LB_P="${LB_P:-2048,8192}"
LB_N="${LB_N:-32}"
LB_R="${LB_R:-3}"

STAMP=$(date -u +%Y-%m-%dT%H:%M:%S)
export LD_LIBRARY_PATH="$BUILD_DIR/bin"
export LABEL BUILD_DIR STAMP RESULTS

OUT=$("$BUILD_DIR/bin/llama-bench" \
  -m "$MODEL" \
  -t 6 -ngl 99 -ncmoe 40 -b 2048 -ub 2048 -fa 1 -mmp 0 \
  -p "$LB_P" -n "$LB_N" -r "$LB_R" -o json 2>/dev/null)

if [ -z "$OUT" ]; then
  echo "ERROR: llama-bench produced no output for $LABEL (bad flag?)" >&2
  exit 3
fi

printf '%s' "$OUT" | python3 -c '
import json, sys, os
label = os.environ["LABEL"]
build = os.environ["BUILD_DIR"]
stamp = os.environ["STAMP"]
results = os.environ["RESULTS"]
rows = json.load(sys.stdin)
if not rows:
    sys.stderr.write("ERROR: llama-bench returned zero rows\n")
    sys.exit(3)
out = open(results, "a")
for r in rows:
    test = "pp%d" % r["n_prompt"] if r["n_prompt"] else "tg%d" % r["n_gen"]
    line = ("LBRESULT label=%s build=%s test=%s tps=%.2f stddev=%.2f "
            "samples=%s time=%s"
            % (label, build, test, r["avg_ts"], r["stddev_ts"],
               r["samples_ts"], stamp))
    out.write(line + "\n")
    print(line)
out.close()
'