#!/usr/bin/env bash
# ---------------------------------------------------------------------------
# serve_ornith_vllm.sh (versión DETACHED)
# Levanta Ornith-1.0 con vLLM en segundo plano (OpenAI-compatible).
#
# Igual que el serve de Gemma 4: nohup + disown -> el servidor sobrevive aunque
# cierres la shell o termine este script. Guarda el PID y espera a que el
# endpoint /health responda antes de devolverte el control.
#
# Ornith es un modelo de razonamiento (<think>...</think>) con tool-calling
# estilo Qwen3, así que activamos --reasoning-parser qwen3 y --tool-call-parser
# qwen3_xml para obtener reasoning_content y tool_calls.
#
# Uso:
# chmod +x serve_ornith_vllm.sh stop_ornith.sh
# ./serve_ornith_vllm.sh # arranca en background y espera al health
# ./stop_ornith.sh # lo apaga
#
# El número de GPUs (tensor-parallel-size) se DETECTA AUTOMÁTICAMENTE.
#
# Variables de entorno configurables (con sus valores por defecto):
# MODEL deepreinforce-ai/Ornith-1.0-9B
# SERVED_NAME Ornith-1.0-9B
# TP_SIZE (auto-detectado) -> fuerza un valor para sobreescribir
# PORT 8000
# HOST 0.0.0.0
# MAX_LEN 262144 (bajalo si te quedas sin VRAM/KV cache, p.ej. 16384)
# GPU_UTIL 0.90
# API_KEY EMPTY (token que exigira el server; cliente debe enviarlo)
# LOG ornith_vllm.log
# PIDFILE ornith_server.pid
# SHOW_PROGRESS 1 (=1 muestra el log en vivo al cargar; =0 silencioso)
#
# Ejemplos:
# ./serve_ornith_vllm.sh # usa TODAS las GPUs visibles
# TP_SIZE=2 ./serve_ornith_vllm.sh # fuerza solo 2 GPUs
# MAX_LEN=16384 ./serve_ornith_vllm.sh # recorta contexto (recomendado en 1 GPU)
# SHOW_PROGRESS=0 ./serve_ornith_vllm.sh # arranque silencioso
# ---------------------------------------------------------------------------
set -euo pipefail
# ---------------------------------------------------------------------------
# Deteccion automatica de GPUs
# Prioridad: 1) TP_SIZE manual 2) CUDA_VISIBLE_DEVICES 3) nvidia-smi
# ---------------------------------------------------------------------------
detect_gpus() {
if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
echo "${CUDA_VISIBLE_DEVICES}" | tr ',' '\n' | grep -cE '^[0-9]+$' || echo 0
return
fi
if command -v nvidia-smi >/dev/null 2>&1; then
nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | grep -c . || echo 0
return
fi
echo 0
}
if [[ -n "${TP_SIZE:-}" ]]; then
echo "[info] TP_SIZE forzado manualmente: ${TP_SIZE}"
else
NUM_GPUS="$(detect_gpus)"
if [[ "${NUM_GPUS}" -lt 1 ]]; then
echo "[error] No se detecto ninguna GPU NVIDIA." >&2
echo " vLLM requiere GPU. Verifica los drivers con 'nvidia-smi'." >&2
echo " Si quieres forzar un valor de todos modos: TP_SIZE=1 ./serve_ornith_vllm.sh" >&2
exit 1
fi
TP_SIZE="${NUM_GPUS}"
echo "[info] GPUs detectadas: ${NUM_GPUS} -> tensor-parallel-size=${TP_SIZE}"
if (( TP_SIZE > 1 )) && (( (TP_SIZE & (TP_SIZE - 1)) != 0 )); then
echo "[warn] ${TP_SIZE} no es potencia de 2; si vLLM falla al cargar," >&2
echo " fuerza una potencia de 2 (p.ej. TP_SIZE=2 o TP_SIZE=4)." >&2
fi
fi
# ---------------------------------------------------------------------------
# Configuracion del modelo / servidor
# ---------------------------------------------------------------------------
# MODEL="${MODEL:-deepreinforce-ai/Ornith-1.0-9B}"
# SERVED_NAME="${SERVED_NAME:-Ornith-1.0-9B}"
MODEL="${MODEL:-tepirale/Ornith-Agents-A1-3.6-35B-A3B-dare_ties}"
SERVED_NAME="${SERVED_NAME:-Ornith-Agents-A1-3.6-35B-A3B-dare_ties}"
PORT="${PORT:-8000}"
HOST="${HOST:-0.0.0.0}"
MAX_LEN="${MAX_LEN:-50000}"
GPU_UTIL="${GPU_UTIL:-0.93}"
# GPU_UTIL="${GPU_UTIL:-0.90}"
API_KEY="${API_KEY:-EMPTY}"
LOG="${LOG:-ornith_vllm.log}"
PIDFILE="${PIDFILE:-ornith_server.pid}"
SHOW_PROGRESS="${SHOW_PROGRESS:-1}"
# Evita arrancar dos veces sobre el mismo PIDFILE
if [ -f "$PIDFILE" ] && kill -0 "$(cat "$PIDFILE")" 2>/dev/null; then
echo "[error] Ya hay un servidor corriendo (PID $(cat "$PIDFILE"), $PIDFILE)." >&2
echo " Apagalo primero con ./stop_ornith.sh" >&2
exit 1
fi
echo "============================================================"
echo " Sirviendo: ${MODEL}"
echo " Nombre expuesto: ${SERVED_NAME}"
echo " GPUs (TP): ${TP_SIZE} | Contexto max: ${MAX_LEN}"
echo " Endpoint: http://${HOST}:${PORT}/v1 | Log: ${LOG}"
echo "============================================================"
export VLLM_API_KEY="${API_KEY}"
# ─── Chat template oficial de Ornith (dispara el modo <think>) ───
TEMPLATE="${TEMPLATE:-ornith_chat_template.jinja}"
# TEMPLATE_URL="${TEMPLATE_URL:-https://huggingface.co/deepreinforce-ai/Ornith-1.0-9B/resolve/main/chat_template.jinja}"
TEMPLATE_URL="${TEMPLATE_URL:-https://huggingface.co/deepreinforce-ai/Ornith-1.0-35B/resolve/main/chat_template.jinja}"
if [ ! -f "$TEMPLATE" ]; then
echo "Descargando chat template de Ornith..."
curl -sfL -o "$TEMPLATE" "$TEMPLATE_URL" \
|| { echo "ERROR: no se pudo descargar el template"; exit 1; }
fi
TEMPLATE_PATH="$(pwd)/$TEMPLATE"
# ---------------------------------------------------------------------------
# Arranque DETACHED (sobrevive al cierre de la shell)
# ---------------------------------------------------------------------------
nohup vllm serve "${MODEL}" \
--served-model-name "${SERVED_NAME}" \
--tensor-parallel-size "${TP_SIZE}" \
--host "${HOST}" --port "${PORT}" \
--max-model-len "${MAX_LEN}" \
--gpu-memory-utilization "${GPU_UTIL}" \
--enable-auto-tool-choice \
--tool-call-parser qwen3_xml \
--reasoning-parser qwen3 \
--chat-template "${TEMPLATE_PATH}" \
--default-chat-template-kwargs '{"enable_thinking": true}' \
--trust-remote-code \
--limit-mm-per-prompt '{"image":4,"video":0}' \
--max-num-seqs 4 \
--kv-cache-dtype bfloat16 \
--attention-backend FLASH_ATTN \
--max-num-batched-tokens 8192 \
--compilation-config '{"cudagraph_mode": "PIECEWISE"}' \
--api-key "${API_KEY}" \
> "${LOG}" 2>&1 &
SERVER_PID=$!
echo "${SERVER_PID}" > "${PIDFILE}"
disown
echo "Servidor lanzado (PID ${SERVER_PID}, guardado en ${PIDFILE}). Logs: ${LOG}"
# ---------------------------------------------------------------------------
# Progreso de carga en vivo (opcional). tqdm de vLLM escribe en el log.
# ---------------------------------------------------------------------------
TAIL_PID=""
if [ "$SHOW_PROGRESS" = "1" ]; then
echo "── Progreso de carga (SHOW_PROGRESS=1) ───────────────────────"
tail -f "$LOG" &
TAIL_PID=$!
fi
stop_tail() { [ -n "$TAIL_PID" ] && kill "$TAIL_PID" 2>/dev/null || true; TAIL_PID=""; }
# ---------------------------------------------------------------------------
# Espera a que cargue el modelo (/health no requiere API key)
# ---------------------------------------------------------------------------
[ "$SHOW_PROGRESS" = "1" ] || echo "Esperando a que el modelo cargue (puede tardar varios minutos)..."
until curl -sf "http://localhost:${PORT}/health" > /dev/null 2>&1; do
if ! kill -0 "$SERVER_PID" 2>/dev/null; then
stop_tail
echo "ERROR: el servidor murió durante el arranque. Últimas líneas:"
tail -n 40 "$LOG"
rm -f "$PIDFILE"
exit 1
fi
sleep 3
done
stop_tail
[ "$SHOW_PROGRESS" = "1" ] && echo "──────────────────────────────────────────────────────────────"
echo "OK: servidor listo en http://localhost:${PORT}/v1 (sigue corriendo en segundo plano)"
echo " Modelo expuesto: ${SERVED_NAME}"
echo " Apaga el servidor: ./stop_ornith.sh (o: kill \$(cat ${PIDFILE}))"