#!/bin/bash
# Model Serving -- run each hand-picked model as its own vLLM container.
#
# Idempotent. Runs first from cloud-init runcmd, then (installed by itself, below) from
# model-serving.service, which retries on failure and after the post-driver reboot until
# every selected model is answering. /opt/model-serving/models.json is the base64 manifest
# cloud-init writes:  {"image": "...", "models": [{hf_id, served_name, port, vram_gb,
# min_compute_capability, args:[...]}, ...]}. This script is generic; the models are data.
set -u
LOG_DIR="/var/startup-script-logs"
LOG_FILE="${LOG_DIR}/model_serving.log"
ANALYSIS_LOG="${LOG_DIR}/analysis.log"
STACK_DIR="/opt/model-serving"
MANIFEST="${STACK_DIR}/models.json"
HF_CACHE="${STACK_DIR}/hf-cache"
PLAN="${STACK_DIR}/plan.sh"
UNIT="/etc/systemd/system/model-serving.service"
REBOOT_FLAG="${STACK_DIR}/.gpu-reboot-done"
mkdir -p "$LOG_DIR"
touch "$ANALYSIS_LOG"
log() { echo "$(date): $*" >> "$LOG_FILE"; }
note_once() { grep -qF "$1" "$ANALYSIS_LOG" || echo "$1" >> "$ANALYSIS_LOG"; }

if [ "${value:-0}" != "1" ]; then
    log "Model Serving not selected; skipping."
    systemctl disable model-serving.service >/dev/null 2>&1 || true
    rm -f "$UNIT"
    rm -rf "$STACK_DIR"
    exit 0
fi

if [ ! -s "$MANIFEST" ]; then
    log "manifest missing or empty at $MANIFEST"
    note_once "model-serving: Failed (no manifest)"
    exit 0
fi
mkdir -p "$HF_CACHE"

# First pass (from cloud-init runcmd): install model-serving.service and hand off to it,
# so there is exactly one worker. The service's ExecStart loop re-runs this script until
# it exits 0 (all models up) and also runs on every boot -- which is what carries the
# deployment across the post-driver reboot. Type=simple + an explicit loop rather than
# Type=oneshot+Restart= so it behaves the same on any systemd version.
if [ ! -f "$UNIT" ]; then
    printf '%s\n' \
        "[Unit]" \
        "Description=vLLM model serving (retry until all selected models are up)" \
        "Wants=network-online.target" \
        "After=network-online.target docker.service" \
        "" \
        "[Service]" \
        "Type=simple" \
        "Environment=value=1" \
        "ExecStart=/bin/bash -c 'until /usr/local/bin/model_serving.sh; do sleep 120; done'" \
        "TimeoutStartSec=0" \
        "" \
        "[Install]" \
        "WantedBy=multi-user.target" \
        > "$UNIT"
    systemctl daemon-reload >> "$LOG_FILE" 2>&1
    systemctl enable model-serving.service >> "$LOG_FILE" 2>&1
    log "installed model-serving.service; handing off"
    systemctl start --no-block model-serving.service >> "$LOG_FILE" 2>&1
    exit 0
fi

# --- wait for the GPU driver (nvidia-open, installed by cuda_toolkit.sh). Short wait: the
# service retries. If it never loads, reboot once -- that reliably brings it up, and this
# script runs again from the service on boot.
gpu_ready=false
for _ in $(seq 1 20); do
    if nvidia-smi -L >/dev/null 2>&1; then gpu_ready=true; break; fi
    sleep 15
done
if [ "$gpu_ready" != true ]; then
    log "GPU driver not ready yet"
    if [ ! -f "$REBOOT_FLAG" ]; then
        touch "$REBOOT_FLAG"
        note_once "model-serving: GPU driver not ready; rebooting once to load it"
        shutdown -r +1 "Rebooting to load the NVIDIA driver for model serving" >> "$LOG_FILE" 2>&1
    fi
    exit 1
fi

TOTAL_MB="$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null | head -n1 | tr -dc '0-9')"
COMPUTE_CAP="$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader,nounits 2>/dev/null | head -n1 | tr -d ' ')"
TOTAL_GB=0
[ -n "$TOTAL_MB" ] && TOTAL_GB=$(( TOTAL_MB / 1024 ))
log "gpu ready: total_gb=$TOTAL_GB compute_cap=$COMPUTE_CAP"

# --- manifest -> sourceable bash plan. Python owns JSON parsing, the gpu-fraction maths
# (vLLM sees the whole card, so concurrent containers need non-overlapping fractions) and
# shlex-quoting each vLLM arg so the --hf-overrides JSON survives.
python3 - "$MANIFEST" "$TOTAL_GB" > "$PLAN" <<'PY'
import json, shlex, sys
data = json.load(open(sys.argv[1]))
total_gb = float(sys.argv[2] or 0)
models = sorted(data["models"], key=lambda m: m["port"])
if total_gb > 0:
    frac = {m["served_name"]: m["vram_gb"] / total_gb for m in models}
    s = sum(frac.values())
    if s > 0.92:
        frac = {k: v * 0.92 / s for k, v in frac.items()}
else:
    frac = {m["served_name"]: 0.30 for m in models}
print("IMAGE=%s" % shlex.quote(data["image"]))
print("NMODELS=%d" % len(models))
for i, m in enumerate(models):
    f = min(max(round(frac[m["served_name"]], 2), 0.05), 0.92)
    argline = " ".join(shlex.quote(str(a)) for a in m["args"])
    print("M%d_NAME=%s"  % (i, shlex.quote(m["served_name"])))
    print("M%d_PORT=%s"  % (i, shlex.quote(str(m["port"]))))
    print("M%d_HF=%s"    % (i, shlex.quote(m["hf_id"])))
    print("M%d_MINCC=%s" % (i, shlex.quote(str(m["min_compute_capability"]))))
    print("M%d_FRAC=%s"  % (i, shlex.quote("%.2f" % f)))
    print("M%d_ARGS=%s"  % (i, shlex.quote(argline)))
PY

if ! bash -n "$PLAN" 2>>"$LOG_FILE"; then
    log "generated plan is not valid bash"
    note_once "model-serving: Failed (bad manifest)"
    exit 0
fi
# shellcheck disable=SC1090
source "$PLAN"
if [ "${NMODELS:-0}" -eq 0 ] || [ -z "${IMAGE:-}" ]; then
    log "empty plan"
    note_once "model-serving: Failed (empty plan)"
    exit 0
fi

for attempt in 1 2 3; do
    log "pulling $IMAGE (attempt $attempt)"
    docker pull "$IMAGE" >> "$LOG_FILE" 2>&1 && break
    sleep 15
done

cc_ok() { awk -v a="$COMPUTE_CAP" -v b="$1" 'BEGIN{exit !(a+0 >= b+0)}'; }

MD="http://169.254.169.254/latest/meta-data"
HOST_IP="$(curl -fsS --max-time 5 "$MD/public-ipv4" 2>/dev/null || true)"
[ -z "$HOST_IP" ] && HOST_IP="$(curl -fsS --max-time 5 "$MD/local-ipv4" 2>/dev/null || true)"
[ -z "$HOST_IP" ] && HOST_IP="$(hostname -I 2>/dev/null | awk '{print $1}')"
[ -z "$HOST_IP" ] && HOST_IP="localhost"

# --- bring up each model, one at a time (two vLLM engines profiling memory at once race
# and abort). Idempotent: a model already answering is left alone.
all_up=true
for i in $(seq 0 $((NMODELS - 1))); do
    v="M${i}_NAME";  NAME="${!v}"
    v="M${i}_PORT";  PORT="${!v}"
    v="M${i}_HF";    HF_ID="${!v}"
    v="M${i}_MINCC"; MIN_CC="${!v}"
    v="M${i}_FRAC";  FRAC="${!v}"
    v="M${i}_ARGS";  eval "MARGS=( ${!v} )"
    CN="vllm-${NAME}"

    if curl -fsS --max-time 5 "http://127.0.0.1:${PORT}/health" >/dev/null 2>&1; then
        note_once "model-serving ${NAME}: Success (http://${HOST_IP}:${PORT}/v1 , model ${NAME})"
        continue
    fi
    if ! cc_ok "$MIN_CC"; then
        log "$NAME needs compute capability $MIN_CC, card is $COMPUTE_CAP; skipping"
        note_once "model-serving ${NAME}: Failed (GPU compute capability ${COMPUTE_CAP} < required ${MIN_CC})"
        continue
    fi

    log "starting $CN on port $PORT (gpu-mem $FRAC)"
    docker rm -f "$CN" >> "$LOG_FILE" 2>&1 || true
    docker run -d --name "$CN" --restart unless-stopped --gpus all --ipc=host \
        -v "${HF_CACHE}:/root/.cache/huggingface" -p "${PORT}:8000" \
        "$IMAGE" "$HF_ID" \
        --served-model-name "$NAME" --host 0.0.0.0 --port 8000 \
        --gpu-memory-utilization "$FRAC" "${MARGS[@]}" >> "$LOG_FILE" 2>&1

    ready=false
    for _ in $(seq 1 96); do
        if curl -fsS --max-time 5 "http://127.0.0.1:${PORT}/health" >/dev/null 2>&1; then
            ready=true; break
        fi
        sleep 5
    done
    if [ "$ready" = true ]; then
        echo "model-serving ${NAME}: Success (http://${HOST_IP}:${PORT}/v1 , model ${NAME})" >> "$ANALYSIS_LOG"
    else
        all_up=false
        note_once "model-serving ${NAME}: Deployed, still starting (http://${HOST_IP}:${PORT}/v1 , model ${NAME})"
        docker logs --tail 20 "$CN" >> "$LOG_FILE" 2>&1 || true
    fi
done

[ "$all_up" = true ] && exit 0 || exit 1