#!/bin/bash
# Fleet benchmark script for llama-swap models
# Usage: bash scripts/bench-fleet.sh
# Requires: llama-swap running on :9292, curl, nvidia-smi
#
# Benchmarks every model defined in llama-swap by:
# 1. Cold-start (load model from disk)
# 2. Send a 500-token generation request
# 3. Parse OpenAI-compatible timings from response
# 4. Report VRAM usage
#
# Modify the MODELS array to match your llama-swap aliases.

MODELS=(
  "gemma-12b:gemma-12b"
  "qwen36:qwen36-35b-mtp"
  "coder-9b:qwopus-9b"
  "uncensored:qwen35-9b-mtp"
  "gemma:gemma-26b"
)

LONG_PROMPT='Write a detailed technical blog post about the architecture and design principles of the model, including attention mechanisms, training methodology, and deployment considerations. Be thorough and specific.'

bench_model() {
  local alias="$1"
  local label="$2"
  echo "=== BENCH: $label ($alias) ==="
  echo "--- Cold start ---"
  curl -s http://127.0.0.1:9292/v1/chat/completions \
    -X POST -H "Content-Type: application/json" \
    -d "{\"model\":\"$alias\",\"messages\":[{\"role\":\"user\",\"content\":\"Respond with: warm\"}],\"max_tokens\":5,\"stream\":false}" \
    --connect-timeout 5 --max-time 120 > /dev/null 2>&1 && echo "Done" || echo "FAILED"
  sleep 2

  echo "--- VRAM (MiB) ---"
  nvidia-smi -q -d MEMORY | grep 'Used' | head -1 | awk '{print $3" "$4}'

  echo "--- Gen 500t ---"
  curl -s http://127.0.0.1:9292/v1/chat/completions \
    -X POST -H "Content-Type: application/json" \
    -d "{\"model\":\"$alias\",\"messages\":[{\"role\":\"user\",\"content\":\"$LONG_PROMPT\"}],\"max_tokens\":500,\"stream\":false}" \
    --connect-timeout 5 --max-time 120 2>/dev/null | python3 -c "
import sys, json
d = json.load(sys.stdin)
t = d.get('timings', {})
u = d.get('usage', {})
print(json.dumps({
  'tokens/s': round(t.get('predicted_per_second', 0), 1),
  'prompt_t/s': round(t.get('prompt_per_second', 0), 1),
  'gen_tokens': u.get('completion_tokens', 0),
  'prompt_tokens': u.get('prompt_tokens', 0),
  'finish_reason': d.get('choices',[{}])[0].get('finish_reason',''),
}, indent=2))" 2>/dev/null || echo "PARSE FAIL"

  echo "--- VRAM (MiB) ---"
  nvidia-smi -q -d MEMORY | grep 'Used' | head -1 | awk '{print $3" "$4}'
  echo ""
}

# Check llama-swap is up
curl -sf http://127.0.0.1:9292/v1/models > /dev/null 2>&1
if [ $? -ne 0 ]; then
  echo "ERROR: llama-swap not running on :9292"
  echo "Start it with: /usr/local/bin/llama-swap --config ~/.config/llama-swap/config.yaml --listen :9292"
  exit 1
fi

# Run benchmarks sequentially (each cold-start evicts the previous model)
for entry in "${MODELS[@]}"; do
  IFS=':' read -r alias label <<< "$entry"
  bench_model "$alias" "$label"
done

echo "=== FLEET BENCHMARK COMPLETE ==="