#!/bin/bash
# Serve a model on one PACE ICE GPU.
#
#   mkdir -p logs && sbatch serve.sbatch
#   sbatch --export=ALL,MODEL=qwen3:14b serve.sbatch
#
# The server lives and dies with this job. Put your own script at the bottom.

#SBATCH --job-name=serve
#SBATCH --partition=ice-gpu
#SBATCH --gres=gpu:v100:1
#SBATCH --cpus-per-task=8
#SBATCH --mem=32G
#SBATCH --time=02:00:00
#SBATCH --output=./logs/%x-%j.out
#SBATCH --mail-type=NONE

cd "$SLURM_SUBMIT_DIR"

MODEL=${MODEL:-qwen3:8b}
EMBED_MODEL=${EMBED_MODEL:-nomic-embed-text}
# A port of its own: two jobs can share a node, and both would otherwise use 11434.
PORT=$((20000 + SLURM_JOB_ID % 20000))
export OLLAMA_HOST=127.0.0.1:$PORT                              # 0.0.0.0:$PORT to allow an ssh tunnel
export OLLAMA_NUM_PARALLEL=${OLLAMA_NUM_PARALLEL:-4}            # requests answered at once
export OLLAMA_CONTEXT_LENGTH=${OLLAMA_CONTEXT_LENGTH:-8192}      # tokens per request; longer prompts are cut silently

echo "node $(hostname) port $PORT"
nvidia-smi --query-gpu=name,memory.total --format=csv,noheader

# The ollama module fails in batch jobs unless HOSTNAME is set.
export HOSTNAME=$(hostname)
module load ollama
OLLAMA_BIN=$(command -v ollama)
ollama serve &> "./logs/server-$SLURM_JOB_ID.log" &
for i in $(seq 1 60); do curl -s "http://$OLLAMA_HOST" >/dev/null && break; sleep 1; done
echo "server ready on $OLLAMA_HOST"

ollama pull "$MODEL"
ollama pull "$EMBED_MODEL"

curl -s "http://$OLLAMA_HOST/v1/models"; echo

# ---- your code goes here -------------------------------------------------
# module load llm-python-apis          # python + openai; this unloads the ollama module,
# export PATH="$(dirname "$OLLAMA_BIN"):$PATH"   # so put the ollama binary back on PATH
# python my_rag.py --base-url "http://$OLLAMA_HOST/v1" --model "$MODEL"

# Warn if any prompt was longer than OLLAMA_CONTEXT_LENGTH and got cut.
echo "prompts truncated: $(grep -c 'truncating input prompt' "./logs/server-$SLURM_JOB_ID.log")"

# Uncomment to hold the GPU so you can reach it from your laptop by ssh tunnel.
# sleep infinity