#!/bin/bash
# Bonsai 2 基准测试：llama-bench 吞吐 + llama-server 显存峰值 + 质量对比
# 用法: bench.sh <gguf> <标签> [端口] [上下文]
set -u
export LD_LIBRARY_PATH=/usr/local/cuda-13.1/lib64:${LD_LIBRARY_PATH:-}
BIN=/home/zyw/bonsai2/bin
OUT=/home/zyw/bonsai2/results
mkdir -p "$OUT"

GGUF="$1"; TAG="$2"; PORT="${3:-8081}"; CTX="${4:-262144}"
MODELS=/home/zyw/bonsai2/models

echo "================ $TAG ================"
echo "--- llama-bench: tg128 / pp512 / pp4096 ---"
$BIN/llama-bench -m "$GGUF" -ngl 99 -fa on -p 512,4096 -n 128 -r 3 2>&1 | tee -a "$OUT/llama-bench-$TAG.txt"

echo
echo "--- 启动 llama-server: ctx=$CTX, KV=q4_0, fa=on, port=$PORT ---"
pkill -x llama-server 2>/dev/null; sleep 2
nohup $BIN/llama-server -m "$GGUF" \
  --mmproj "$MODELS/mmproj-Q8_0.gguf" \
  -ngl 99 -fa on -c "$CTX" -ctk q4_0 -ctv q4_0 -np 1 \
  --jinja --host 0.0.0.0 --port "$PORT" --alias bonsai2-27b \
  > "$OUT/server-$TAG.log" 2>&1 &
echo "server pid=$!"

for i in $(seq 1 90); do
  if curl -s --noproxy '*' -m 2 "http://127.0.0.1:$PORT/health" | grep -q ok; then
    echo "服务就绪（等了 ${i}s）"; break
  fi
  sleep 1
done

echo "--- 加载后显存 ---"
nvidia-smi --query-gpu=memory.used,memory.total --format=csv,noheader

echo "--- 发送一个长需求以扩充 KV，然后测显存峰值 ---"
python3 - "$PORT" <<'PYEOF'
import json,sys,urllib.request
port=sys.argv[1]
long_text=("请阅读下面这段材料并只用一句话总结它。材料：" + "三值量化把权重映射到负一零正一三个值，"
 "每128个权重共享一个FP16缩放因子，权重存储在Hadamard旋转基下，"
 "推理时对激活施加匹配的沃尔什哈达玛变换。"*400)
payload={"model":"bonsai2-27b","messages":[{"role":"user","content":long_text}],
         "max_tokens":64,"temperature":0.7,"chat_template_kwargs":{"enable_thinking":False}}
req=urllib.request.Request("http://127.0.0.1:%s/v1/chat/completions"%port,
    data=json.dumps(payload).encode(),headers={"Content-Type":"application/json"})
try:
    r=json.load(urllib.request.urlopen(req,timeout=1800))
    print("OK completion_tokens=",r.get("usage",{}).get("completion_tokens"),
          " prompt_tokens=",r.get("usage",{}).get("prompt_tokens"))
    print("timings=",json.dumps({k:r.get("timings",{}).get(k) for k in
          ("prompt_ms","prompt_per_second","predicted_per_second")},ensure_ascii=False))
except Exception as e:
    print("FAIL",repr(e))
PYEOF

echo "--- 长请求后显存峰值 ---"
nvidia-smi --query-gpu=memory.used,memory.total --format=csv,noheader
echo "--- /props ---"
curl -s --noproxy '*' "http://127.0.0.1:$PORT/props" | python3 -c "import json,sys; d=json.load(sys.stdin); print({k:d.get(k) for k in ('n_ctx','total_slots','model_path','chat_template')if k!='chat_template'})" 2>/dev/null || echo "(props 解析失败)"

echo
echo "--- 质量对比提示词（端口 $PORT）---"
python3 /home/zyw/bonsai2/quality_test.py "$PORT" bonsai2-27b "$OUT/quality-$TAG"

echo
echo "--- 关闭服务 ---"
pkill -x llama-server 2>/dev/null
sleep 2
echo "================ $TAG 完成 ================"