#!/usr/bin/env bash
# Purpose: start TabbyAPI, the server the ExLlamaV3 project names as its recommended backend,
#          after checking that its config.yml loads the lab's EXL3 model with the lab's
#          context length, so that the model can be measured through the same
#          OpenAI-compatible interface as every other engine here
# Platform: nvidia (ExLlamaV3 requires CUDA; its README lists ROCm support as still to do)
# Minimum memory: 12 GB
# Assumes: a TabbyAPI checkout in $TABBY_DIR, a config.yml written as this lab's Track N tab
#          shows, the EXL3 model directory it names, an NVIDIA GPU with a working driver, the
#          port free; the first start creates a venv and installs PyTorch and ExLlamaV3 (a
#          long, unattended download); the server runs in the foreground
#
# Usage: bash serve-exllamav3.sh
#        TABBY_DIR=~/tabbyAPI GPU_LIB=cu13 bash serve-exllamav3.sh
#
# Environment:
#   TABBY_DIR  the TabbyAPI checkout                        (default: $HOME/tabbyAPI)
#   GPU_LIB    dependency set for the first start: cu12, cu13 (default: let start.py choose)
#   PIDFILE    where to write the process id                (default: ./tabbyapi.pid)
#
# TabbyAPI is configured through config.yml rather than command-line options: model_dir and
# model_name choose the model, max_seq_len and cache_size (in tokens, a multiple of 256) set
# the context, and api_tokens.yml holds the generated API key unless network.disable_auth is
# true. The pid recorded is start.sh's shell; the Python server is its child, so for memory
# use nvidia-smi's per-process figure, which the measuring script reads. TabbyAPI is licensed
# AGPL-3.0; read it before building anything on it.

set -euo pipefail

TABBY_DIR="${TABBY_DIR:-$HOME/tabbyAPI}"
GPU_LIB="${GPU_LIB:-}"
PIDFILE="${PIDFILE:-./tabbyapi.pid}"

die() { echo "serve-exllamav3: $*" >&2; exit 1; }

command -v nvidia-smi >/dev/null || die "nvidia-smi not found; ExLlamaV3 needs an NVIDIA GPU"
[ -d "$TABBY_DIR" ] || die "$TABBY_DIR does not exist; clone TabbyAPI there or set TABBY_DIR"
[ -f "$TABBY_DIR/start.sh" ] || die "no start.sh in $TABBY_DIR"
CONFIG="$TABBY_DIR/config.yml"
[ -f "$CONFIG" ] || die "no config.yml in $TABBY_DIR; write it as the lab page shows"

value() { sed -n "s/^[[:space:]]*$1:[[:space:]]*//p" "$CONFIG" | head -n 1 | tr -d '"'; }
MODEL_DIR="$(value model_dir)"
MODEL_NAME="$(value model_name)"
HOST="$(value host)"
PORT="$(value port)"
[ -n "$MODEL_NAME" ] || die "config.yml has no model_name; TabbyAPI would start with no model loaded"
[ -f "$MODEL_DIR/$MODEL_NAME/config.json" ] || die "$MODEL_DIR/$MODEL_NAME/config.json not found"
grep -q '"quant_method": "exl3"' "$MODEL_DIR/$MODEL_NAME/config.json" \
  || echo "serve-exllamav3: warning: $MODEL_NAME/config.json does not say quant_method exl3" >&2
[ "$(value max_seq_len)" = "8192" ] || echo "serve-exllamav3: warning: max_seq_len is not 8192" >&2
[ "$(value cache_size)" = "8192" ] || echo "serve-exllamav3: warning: cache_size is not 8192" >&2
if command -v curl >/dev/null && curl --silent --max-time 2 "http://${HOST:-127.0.0.1}:${PORT:-5000}/health" >/dev/null 2>&1; then
  die "something is already answering on ${HOST:-127.0.0.1}:${PORT:-5000}; stop it first"
fi

echo "==> TabbyAPI with the ExLlamaV3 backend"
echo "    checkout $TABBY_DIR ($(git -C "$TABBY_DIR" rev-parse --short HEAD 2>/dev/null || echo 'not a git checkout'))"
echo "    model    $MODEL_DIR/$MODEL_NAME"
echo "    context  max_seq_len $(value max_seq_len), cache_size $(value cache_size) tokens"
echo "    listen   http://${HOST:-127.0.0.1}:${PORT:-5000}/v1"
echo "    pid written to $PIDFILE"
if [ -f "$TABBY_DIR/api_tokens.yml" ]; then
  echo "    API key: grep '^api_key:' $TABBY_DIR/api_tokens.yml"
else
  echo "    api_tokens.yml is created on this first start; read the key from it afterwards"
fi

echo $$ > "$PIDFILE"
cd "$TABBY_DIR"
if [ -n "$GPU_LIB" ]; then
  exec ./start.sh --gpu-lib "$GPU_LIB"
fi
exec ./start.sh