#!/usr/bin/env bash
# =============================================================================
# setup_vllm.sh — Install & run vLLM on Ubuntu 26 for ConflictChecker hybrid mode
#
# Target box: 2x RTX 3080 (12GB each, 24GB total, Ampere sm_86, NO NVLink).
# Serves an OpenAI-compatible endpoint for the TEXT stages; ConflictChecker
# points LOCAL_BASE_URL at it. Vision stays on OpenRouter.
#
# Usage:
#   chmod +x setup_vllm.sh
#   ./setup_vllm.sh install     # one-time: drivers (if needed), uv, venv, vllm
#   ./setup_vllm.sh serve       # launch the server (foreground)
#   ./setup_vllm.sh service     # install + start a systemd service (persistent)
#   ./setup_vllm.sh test        # smoke-test the running endpoint
#
# Tested against vLLM v0.23 on 2x3080. Fixes baked in:
#   - no --disable-log-requests (removed in newer vLLM)
#   - NCCL_P2P_DISABLE=1            (3080s have no working GPU P2P)
#   - --disable-custom-all-reduce  (same reason; silences the P2P warning)
#   - VLLM_USE_FLASHINFER_SAMPLER=0 (FlashInfer JIT sampler crashes w/o CUDA toolkit)
#   - ENFORCE_EAGER=1 default       (less VRAM + faster startup on the tight 24GB)
# =============================================================================
set -euo pipefail

# ---- Config (override via env) ----------------------------------------------
VENV_DIR="${VENV_DIR:-$HOME/vllm-env}"
PY_VERSION="${PY_VERSION:-3.12}"                 # vLLM-supported Python
MODEL="${MODEL:-Qwen/Qwen3-14B-AWQ}"            # 14B-AWQ ~9GB; fits 24GB w/ TP2
TP_SIZE="${TP_SIZE:-2}"                          # 2 = shard across both 3080s
MAX_LEN="${MAX_LEN:-32768}"                      # context cap (KV-cache bound)
PORT="${PORT:-8000}"
HOST="${HOST:-0.0.0.0}"                          # 0.0.0.0 so the app host can reach it
GPU_UTIL="${GPU_UTIL:-0.90}"                     # fraction of VRAM vLLM may use
QUANT="${QUANT:-awq_marlin}"                     # AWQ Int4 (Ampere has no FP8)
DTYPE="${DTYPE:-float16}"                        # Ampere: fp16
ENFORCE_EAGER="${ENFORCE_EAGER:-1}"             # 1 = skip CUDA-graph capture
# Accuracy alternative (both cards):  MODEL=Qwen/Qwen3-32B-AWQ MAX_LEN=16384 ./setup_vllm.sh serve
# Simplicity alternative (one card):  CUDA_VISIBLE_DEVICES=0 TP_SIZE=1 MAX_LEN=8192 ./setup_vllm.sh serve
# If it OOMs:                         MAX_LEN=16384 (or 8192) and/or GPU_UTIL=0.85
# -----------------------------------------------------------------------------

log(){ printf '\n\033[1;36m== %s ==\033[0m\n' "$*"; }

# Env that makes vLLM behave on no-NVLink Ampere + no CUDA toolkit.
export_runtime_env() {
  export NCCL_P2P_DISABLE="${NCCL_P2P_DISABLE:-1}"
  export NCCL_SHM_DISABLE="${NCCL_SHM_DISABLE:-0}"
  export VLLM_USE_FLASHINFER_SAMPLER="${VLLM_USE_FLASHINFER_SAMPLER:-0}"
}

install_system() {
  log "System packages"
  sudo apt-get update -y
  sudo apt-get install -y build-essential git curl ca-certificates

  log "NVIDIA driver"
  if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi >/dev/null 2>&1; then
    nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv
  else
    echo "Installing recommended NVIDIA driver (reboot required afterwards)..."
    sudo apt-get install -y ubuntu-drivers-common
    sudo ubuntu-drivers autoinstall
    echo "!! Reboot now, then re-run: ./setup_vllm.sh install"
    exit 0
  fi
}

install_uv() {
  if ! command -v uv >/dev/null 2>&1; then
    log "Installing uv (Python/venv manager)"
    curl -LsSf https://astral.sh/uv/install.sh | sh
    export PATH="$HOME/.local/bin:$PATH"
  fi
}

install_vllm() {
  log "Creating venv ($VENV_DIR, Python $PY_VERSION)"
  uv venv "$VENV_DIR" --python "$PY_VERSION"
  log "Installing vLLM (pulls matching CUDA torch wheels)"
  uv pip install --python "$VENV_DIR/bin/python" --upgrade pip
  uv pip install --python "$VENV_DIR/bin/python" vllm
  "$VENV_DIR/bin/python" -c "import vllm, torch; print('vLLM', vllm.__version__, '| CUDA avail', torch.cuda.is_available(), torch.version.cuda)"
}

build_args() {
  ARGS=( serve "$MODEL"
    --host "$HOST" --port "$PORT"
    --tensor-parallel-size "$TP_SIZE"
    --quantization "$QUANT"
    --dtype "$DTYPE"
    --max-model-len "$MAX_LEN"
    --gpu-memory-utilization "$GPU_UTIL"
    --served-model-name "$MODEL" )
  [ "$TP_SIZE" -gt 1 ] && ARGS+=(--disable-custom-all-reduce)
  [ "$ENFORCE_EAGER" = "1" ] && ARGS+=(--enforce-eager)
}

serve() {
  export_runtime_env
  build_args
  log "Serving $MODEL  (TP=$TP_SIZE, max_len=$MAX_LEN, eager=$ENFORCE_EAGER, port=$PORT)"
  exec "$VENV_DIR/bin/vllm" "${ARGS[@]}"
}

install_service() {
  export_runtime_env
  build_args
  log "Installing systemd service (vllm.service)"
  # Render the resolved arg list into the unit.
  local exec_args=""; for a in "${ARGS[@]}"; do exec_args+=" $a"; done
  sudo tee /etc/systemd/system/vllm.service >/dev/null <<EOF
[Unit]
Description=vLLM OpenAI server (ConflictChecker hybrid text backend)
After=network-online.target
Wants=network-online.target

[Service]
User=$USER
Environment=PATH=$VENV_DIR/bin:/usr/bin:/bin
Environment=NCCL_P2P_DISABLE=$NCCL_P2P_DISABLE
Environment=NCCL_SHM_DISABLE=$NCCL_SHM_DISABLE
Environment=VLLM_USE_FLASHINFER_SAMPLER=$VLLM_USE_FLASHINFER_SAMPLER
ExecStart=$VENV_DIR/bin/vllm$exec_args
Restart=on-failure
RestartSec=10
TimeoutStartSec=900

[Install]
WantedBy=multi-user.target
EOF
  sudo systemctl daemon-reload
  sudo systemctl enable --now vllm.service
  echo "Started. Logs: journalctl -u vllm -f"
}

smoke_test() {
  log "Endpoint check (http://localhost:$PORT)"
  curl -s "http://localhost:$PORT/v1/models" | head -c 400; echo
  echo "--- chat completion (JSON mode) ---"
  curl -s "http://localhost:$PORT/v1/chat/completions" \
    -H 'Content-Type: application/json' \
    -d "{\"model\":\"$MODEL\",\"messages\":[{\"role\":\"user\",\"content\":\"Reply with JSON {\\\"ok\\\":true}\"}],\"response_format\":{\"type\":\"json_object\"},\"max_tokens\":50}" \
    | head -c 600; echo
}

case "${1:-}" in
  install) install_system; install_uv; install_vllm
           echo; echo "Done. Run: ./setup_vllm.sh serve   (or: ./setup_vllm.sh service)";;
  serve)   serve;;
  service) install_service;;
  test)    smoke_test;;
  *) echo "usage: $0 {install|serve|service|test}"; exit 1;;
esac