#!/bin/bash
# span-02 launcher: locates (or downloads) the weights and the MTP draft, then starts dllm-server
# (/v1/systemone, /health) on $PORT with the Decision Index submission settings.
#
#   bash launch.sh                    # from a clone of this repo (weights next to this script) or anywhere (downloads)
#   PORT=8000 HF_TOKEN=... bash launch.sh
#
# Env:
#   DLLM_MODEL_DIR   directory with model.safetensors, config.json, tokenizer.json, draft_vocab_16384.json, think_cost.json
#                    (default: this script's directory if it holds model.safetensors, else $DLLM_DATA/model, downloaded from $DLLM_HF_REPO)
#   DLLM_HF_REPO     model repo to download (respanai/span-02)
#   DLLM_DRAFT_REPO  MTP draft for speculative decoding (google/gemma-4-26B-A4B-it-assistant; public)
#   DLLM_DATA        download directory (/data/dllm)
#   DLLM_SERVER_BIN  server binary (default: dllm-server on PATH, else engine/target/release/dllm-server next to this script)
#   PORT / HOST      bind address (0.0.0.0:8000)
#   AR_API_KEY       optional bearer token the server then requires
#   HF_TOKEN         needed while the model repo is private
# Extra arguments are passed to dllm-server (`dllm-server --help` lists them).
set -euo pipefail
HERE=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
REPO=${DLLM_HF_REPO:-respanai/span-02}
DRAFT=${DLLM_DRAFT_REPO:-google/gemma-4-26B-A4B-it-assistant}
D=${DLLM_DATA:-/data/dllm}
if [ -n "${DLLM_MODEL_DIR:-}" ]; then M=$DLLM_MODEL_DIR
elif [ -f "$HERE/model.safetensors" ]; then M=$HERE
else M=$D/model; fi
if [ -n "${DLLM_SERVER_BIN:-}" ]; then BIN=$DLLM_SERVER_BIN
elif command -v dllm-server >/dev/null; then BIN=dllm-server
else BIN=$HERE/engine/target/release/dllm-server; fi
dl() {  # repo dir
  if [ -f "$2/config.json" ] && ls "$2"/*.safetensors >/dev/null 2>&1; then echo "[launch] $1 found in $2"; return; fi
  mkdir -p "$2"
  if command -v hf >/dev/null; then hf download "$1" --local-dir "$2" --exclude 'engine/*'
  else huggingface-cli download "$1" --local-dir "$2" --exclude 'engine/*'; fi
}
dl "$REPO" "$M"
dl "$DRAFT" "$D/draft"
# Submission settings (do not change for a board run):
#  --budget-ms 0          no wall-clock limit: the clock is never read for thinking decisions
#  --think-ms 750 --think-cost think_cost.json   deterministic thought budget: per-question thought cap = a pure function of
#                         the request under the fixed cost model in think_cost.json (targets 750 ms on an RTX PRO 6000)
#  --cap 1024             hard maximum of that cap; a thought ends at the end-of-thought token, EOS or the cap
#  --tau 0.8              think when the no-think answer's top probability < 0.8; final distribution = 0.5 no-think + 0.5 post-thought
#  --spec-draft           MTP speculative decoding of thoughts and number answers (lossless: token-identical to greedy)
#  --image-soft-tokens 560  Gemma-4 image budget (max_soft_tokens 560, about 1.3 MP)
#  --serial --no-prefix-cache --spec-max-branches 16   determinism: answers do not depend on load or earlier requests
#  --max-len 131072       context limit per question (prompt + images + thought)
#  --pages 12000 --kv-gb 24   KV pools (full attention 192K tokens, 3.9 GB; sliding window ~20 GB from --kv-gb 24) sized so every DI request
#                         fits on its own under --serial (largest DI 0.3 request: ~3.7K full pages); the pool size does not change answers
exec "$BIN" --bind "${HOST:-0.0.0.0}:${PORT:-8000}" --model "$M" --spec-draft "$D/draft" \
  --max-len 131072 --budget-ms 0 --tau 0.8 --cap 1024 --think-ms 750 --think-cost "$M/think_cost.json" \
  --image-soft-tokens 560 --serial --no-prefix-cache --spec-max-branches 16 \
  --pages 12000 --kv-gb 24 "$@"