#!/usr/bin/env bash
# Serve ornith-ai/Ornith-1.5-35B-A3B-FP8 as frank-default-small via
# FreeToken on Dankserver (RTX 3090, 62 GB RAM).
#
# Runtime: FreeToken v0.1.2 (edge-native MoE serving engine).
#   - Experts live in host RAM; an LRU cache of expert slots is resident
#     on the GPU. Misses are streamed over PCIe.
#   - Paged KV cache (radix prefix sharing).  Only the 10 full-attention
#     layers page KV (~20 KiB/token); the 30 GDN/linear-attention layers
#     hold constant-size per-sequence state (mamba slots).
#   - OpenAI-compatible API (/v1/chat/completions, /v1/models).
#
# Model: ornith-ai/Ornith-1.5-35B-A3B-FP8 — Ornith-1.5 (self-improvement
# RL-trained coding/agentic model over Qwen3.5 + Gemma4), 35B total / 3B
# active, in compressed-tensors per-channel FP8 (float-quantized, dynamic
# token activations).  Same qwen3_5_moe arch and ~35 GB FP8 mass as the
# Huihui FP8 it replaces: 40 layers (10 full-attention + 30 GDN), 256
# experts / 8 used, 2 KV heads × 128 dim, native 262144 context.  FreeToken
# loads a Qwen-Sharp chat template over it (see CHAT_TEMPLATE_FILE in the
# env file); the checkpoint's own stock template mis-serializes reasoning.
#
# Pool sizing: the GPU budget is split MoE-first between expert residency
# and KV pages.  --moe-cache-size pins the number of resident expert
# slots; --num-tokens pins the KV pool capacity.  Live resize is possible
# via `ft ctl cache --kv N --moe M` (grows KV by shrinking MoE; shrinking
# KV requires a restart due to the swap-during-rebuild memory cost).
#
# Tunables come from frank-default-small.env unless FRANK_SMALL_ENV_FILE
# overrides that path.
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ENV_FILE="${FRANK_SMALL_ENV_FILE:-$SCRIPT_DIR/frank-default-small.env}"
if [[ -f "$ENV_FILE" ]]; then
    set -a
    # shellcheck disable=SC1090
    source "$ENV_FILE"
    set +a
fi

# --- defaults (used when the env file does not set them) -----------------
: "${MODEL_DIR:=/home/dank/models/ornith-1.5-35b-a3b-fp8}"
: "${ALIAS:=frank-default-small}"
: "${BIND_ADDR:=0.0.0.0}"
: "${LISTEN_PORT:=7770}"
: "${FT_VENV:=/home/dank/freetoken/.venv}"
: "${CUDA_HOME_DIR:=${FT_VENV}/lib/python3.12/site-packages/nvidia/cu13}"

# Pool sizing (expert residency vs KV cache tradeoff).
: "${MOE_CACHE_SIZE:=2048}"
: "${NUM_TOKENS:=453005}"
: "${MAX_RUNNING_REQUESTS:=4}"

# Paths
FT_BIN="${FT_VENV}/bin/ft"
MODEL="${MODEL_DIR}"

if [[ ! -x "$FT_BIN" ]]; then
    echo "missing FreeToken CLI: $FT_BIN" >&2
    exit 1
fi
if [[ ! -d "$MODEL" ]]; then
    echo "missing model directory: $MODEL" >&2
    exit 1
fi

export PATH="${FT_VENV}/bin:${PATH}"
export CUDA_HOME="${CUDA_HOME_DIR}"

exec "$FT_BIN" serve \
    --model "$MODEL" \
    --served-model-name "$ALIAS" \
    --host "$BIND_ADDR" \
    --port "$LISTEN_PORT" \
    --moe-cache-size "$MOE_CACHE_SIZE" \
    --num-tokens "$NUM_TOKENS" \
    --max-running-requests "$MAX_RUNNING_REQUESTS"