Mooncake/mooncake-integration/ollama/scripts/llama_start.sh

56 lines
2.0 KiB
Bash

#!/usr/bin/env bash
# Start one llama.cpp server instance ("agent worker").
# Usage: llama_start.sh <name> <gpu|auto> <port> [model_gguf] [ctx] [np]
# auto-pick only scans the devices listed in $OMB_GPUS.
set -euo pipefail
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/env.sh"
NAME="${1:?name}"; GPU="${2:-auto}"; PORT="${3:?port}"
MODEL="${4:-$MODELS_DIR/qwen2.5-coder-1.5b-instruct-q4_k_m.gguf}"
CTX="${5:-40960}"; NP="${6:-1}"
pick_free_gpu() {
# pick a GPU from $OMB_GPUS that is currently near-idle (< 2 GiB used)
local used line idx mem
for idx in ${OMB_GPUS//,/ }; do
mem=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i "$idx" 2>/dev/null | tr -d ' ')
if [ -n "$mem" ] && [ "$mem" -lt 2048 ]; then echo "$idx"; return 0; fi
done
echo "${OMB_GPUS%%,*}" # fallback: first allowed gpu
}
if [ "$GPU" = "auto" ]; then GPU="$(pick_free_gpu)"; fi
if ss -ltn 2>/dev/null | grep -q "127.0.0.1:$PORT "; then
echo "[llama:$NAME] already on $PORT"; exit 0
fi
LOG="$OMB_LOGS/llama_${NAME}.log"
KVTYPE="${OMB_LLAMA_KVTYPE:-f16}" # f16 | q8_0 | q4_0 (smaller KV => faster store transfer)
EXTRA=()
if [ "$KVTYPE" != "f16" ]; then
# quantized KV cache requires flash attention in llama.cpp
EXTRA+=(--cache-type-k "$KVTYPE" --cache-type-v "$KVTYPE" -fa on)
elif [ "${OMB_LLAMA_FA:-0}" = "1" ]; then
EXTRA+=(-fa on)
fi
echo "[llama:$NAME] gpu=$GPU port=$PORT ctx=$CTX np=$NP kv=$KVTYPE model=$(basename "$MODEL")"
CUDA_VISIBLE_DEVICES="$GPU" nohup "$LLAMA_BUILD/bin/llama-server" \
-m "$MODEL" \
--host 127.0.0.1 --port "$PORT" \
-ngl 99 -c "$CTX" -np "$NP" -t 8 \
--slot-save-path "$OMB_SLOTS/" \
--no-webui \
"${EXTRA[@]}" \
> "$LOG" 2>&1 &
echo $! > "$OMB_RUN/llama_${NAME}.pid"
# wait for health
for _ in $(seq 1 240); do
if curl -fsS "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
echo "[llama:$NAME] healthy on $PORT (gpu $GPU)"; exit 0
fi
sleep 0.5
done
echo "[llama:$NAME] FAILED to become healthy"; tail -30 "$LOG"; exit 1