#!/usr/bin/env bash
# GPU 0, colocated with STT. CAG/RAG live in the LiveKit agent — this script
# only serves granite-4.1-8b. Mem fraction is reduced so Cohere STT fits.
set -euo pipefail
DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$DIR/common.sh"

MODEL="${MODEL:-ibm-granite/granite-4.1-8b}"
PORT="${PORT:-8000}"
# ~0.68 of 32GB ≈ 21.8GB for weights+KV; leaves ~0.18 for STT on the same card.
MEM_FRACTION="${MEM_FRACTION:-0.68}"
MAX_REQ="${MAX_REQ:-24}"
LOG_DIR="${LOG_DIR:-$HOME/agent-logs}"
mkdir -p "$LOG_DIR"

ensure_toolchain
# Shares ~/sglang-env with TTS (omni installs sglang).
ensure_sglang_omni

export CUDA_VISIBLE_DEVICES=0

if ! python3 -c "import sglang" 2>/dev/null; then
  echo "[llm] FATAL: sglang failed to install/import."
  exit 1
fi

echo "[llm] launching $MODEL on :$PORT [GPU0 w/ STT, mem $MEM_FRACTION]"
# On RTX 5090, FP8 KV without scales + ngram sampling triggered a device-side
# assert and killed the server. Use bf16 KV and no speculative decoding until
# those paths are validated on Blackwell.
python3 -m sglang.launch_server \
  --model-path "$MODEL" \
  --host 0.0.0.0 \
  --port "$PORT" \
  --attention-backend flashinfer \
  --mem-fraction-static "$MEM_FRACTION" \
  --chunked-prefill-size 512 \
  --max-running-requests "$MAX_REQ" \
  --schedule-policy lpm \
  > "$LOG_DIR/llm.log" 2>&1 &
LLM_PID=$!

waited=0; max_wait=600
until curl -sf "http://localhost:${PORT}/health" > /dev/null 2>&1; do
  if [ "$waited" -ge "$max_wait" ]; then
    echo "[llm] TIMEOUT"; tail -n 30 "$LOG_DIR/llm.log"; exit 1
  fi
  if ! kill -0 "$LLM_PID" 2>/dev/null; then
    echo "[llm] process died"; tail -n 40 "$LOG_DIR/llm.log"; exit 1
  fi
  sleep 2; waited=$((waited + 2))
done

echo "$LLM_PID" > /tmp/llm.pid
echo "[llm] up. pid=$LLM_PID"