#!/usr/bin/env bash
# n-cpu-moe Sweep Benchmark Template
# Usage: bash scripts/ncm-sweep.sh
# Tests a range of --n-cpu-moe values, records VRAM + generation speed
# Edit the variables at the top to match your model and desired test values.
set -uo pipefail

# === CONFIGURE THESE ===
MODEL="/models/downloads/Ornith-1.0-35B-GGUF/ornith-1.0-35b-Q8_0.gguf"
PORT=18001
BINARY="/usr/local/bin/llama-server-sm75"
VALUES=(40 38 36 32 28 24)   # highest first (fastest to load, gradual VRAM increase)
PROMPT="Write a Python function that checks if a number is prime."
# =======================

echo "===== n-cpu-moe sweep: $(basename $MODEL) ====="
echo "GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader)"
nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | xargs -I{} echo "VRAM total: {} MiB"

for NCMOE in "${VALUES[@]}"; do
    echo ""
    echo "--- --n-cpu-moe $NCMOE ---"

    kill $(lsof -t -i:$PORT) 2>/dev/null || true
    sleep 1

    $BINARY \
        -m "$MODEL" \
        --port $PORT --host 127.0.0.1 \
        --jinja --reasoning off \
        -ngl 99 --no-mmap --mlock \
        --prio 2 --poll 30 --no-cont-batching \
        --timeout 60 --threads 16 --threads-batch 32 \
        --parallel 1 -c 8192 \
        --batch-size 2048 --ubatch-size 512 \
        --cache-type-k q8_0 --cache-type-v q8_0 \
        --flash-attn on --no-host -fitt 1024 \
        --n-cpu-moe $NCMOE --no-kv-offload \
        --temp 0.6 --top-p 0.95 --min-p 0.05 --top-k 20 \
        --metrics \
        > /tmp/ncm-server-$NCMOE.log 2>&1 &
    SERVER_PID=$!

    echo "  Loading model..."
    for i in $(seq 1 120); do
        if curl -sf http://127.0.0.1:$PORT/v1/models > /dev/null 2>&1; then
            echo "  Ready after ${i}s"
            break
        fi
        sleep 1
        if ! kill -0 $SERVER_PID 2>/dev/null; then
            echo "  FAIL: Server crashed during load (OOM?)"
            tail -3 /tmp/ncm-server-$NCMOE.log
            continue 2
        fi
    done

    VRAM=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits)
    echo "  VRAM: ${VRAM}MiB"

    for run in 1 2 3; do
        RESULT=$(curl -s http://127.0.0.1:$PORT/v1/chat/completions \
            -H 'Content-Type: application/json' \
            -d "{\"model\":\"test\",\"messages\":[{\"role\":\"user\",\"content\":\"$PROMPT\"}],\"max_tokens\":100,\"temperature\":0.6}" 2>/dev/null)
        TG=$(echo "$RESULT" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('timings',{}).get('predicted_per_second', 'N/A'))" 2>/dev/null)
        PP=$(echo "$RESULT" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('timings',{}).get('prompt_per_second', 'N/A'))" 2>/dev/null)
        echo "  Run $run: PP=${PP} t/s, TG=${TG} t/s"
    done

    kill $SERVER_PID 2>/dev/null || true
    wait $SERVER_PID 2>/dev/null || true
    sleep 2
done

echo ""
echo "===== Done ====="