#!/bin/bash
# Sweep llama.cpp runner configs for Qwen3.6-35B-A3B; report VRAM + throughput.
#
# Usage:   ./sweep.sh <ncmoe>:<tensor-split>[:<threads>] ...
# Example: ./sweep.sh 10:24,16:8 8:24,16:8
#
# Env: CTX (default 131072), DEPTHS (default 512,8192), PORT (default 18008)

HUB=${HUB:-/mnt/2TSAM990nvme/docker-volume-outsource/llm-models/hf/hub/models--unsloth--Qwen3.6-35B-A3B-GGUF}
SNAP=${SNAP:-snapshots/a483e9e6cbd595906af30beda3187c2663a1118c/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf}
CTX=${CTX:-131072}
PORT=${PORT:-18008}
NAME=qwen36-sweep
BENCH=${BENCH:-$(dirname "$0")/bench.py}

run_cfg() {
  local spec="$1"
  local ncmoe="${spec%%:*}"
  local rest="${spec#*:}"
  local ts="${rest%%:*}"
  local threads="${rest#*:}"
  [ "$threads" = "$ts" ] && threads=8

  echo "########## n-cpu-moe=$ncmoe  -ts $ts  threads=$threads  ctx=$CTX ##########"
  docker rm -f $NAME >/dev/null 2>&1
  docker run -d --name $NAME --runtime nvidia --network host \
    -e NVIDIA_VISIBLE_DEVICES=0,1 \
    -e NVIDIA_DRIVER_CAPABILITIES=compute,utility \
    -e CUDA_DEVICE_ORDER=PCI_BUS_ID \
    -v "$HUB":/models:ro \
    ghcr.io/ggml-org/llama.cpp:server-cuda \
    -m /models/$SNAP \
    --ctx-size $CTX -ctk q8_0 -ctv q8_0 -fa on \
    -ngl 99 --n-cpu-moe "$ncmoe" -sm layer -ts "$ts" \
    --threads "$threads" -np 1 --no-mmap $EXTRA \
    --host 0.0.0.0 --port $PORT >/dev/null 2>&1

  local ok=0
  for _ in $(seq 1 200); do
    if docker logs $NAME 2>&1 | grep -q "listening on"; then ok=1; break; fi
    if docker logs $NAME 2>&1 | grep -qiE "out of memory|CUDA error|failed to allocate|terminate called|error loading model"; then ok=2; break; fi
    if ! docker ps --format '{{.Names}}' | grep -q "^$NAME$"; then ok=3; break; fi
    sleep 3
  done

  if [ "$ok" != "1" ]; then
    echo "RESULT ncmoe=$ncmoe ts=$ts threads=$threads -> FAILED"
    docker logs $NAME 2>&1 | grep -iE "out of memory|failed to allocate" | head -2
    docker rm -f $NAME >/dev/null 2>&1; echo; return 1
  fi

  echo "--- VRAM used/free (MiB) ---"
  nvidia-smi --query-gpu=index,memory.used,memory.free --format=csv,noheader
  python3 "$BENCH" --port $PORT --label "n${ncmoe}_ts${ts}_t${threads}" \
    --depths "${DEPTHS:-512,8192}" --n-predict 128
  docker rm -f $NAME >/dev/null 2>&1
  echo
}

for c in "$@"; do run_cfg "$c"; done