#!/usr/bin/env bash
# serve-frank-primary.sh — frank-primary endpoint (vLLM 0.30,
# Qwen3.8-Flash-Next-FP8 on the 2x R9700 pair, dankserver :7770).
#
# Validates the image, the model payload and the env tunables, then execs the
# container. The unit (scripts/frank-primary.service, root systemd, User=dank,
# LimitMEMLOCK=infinity) is what makes the UVA pin budget legal — a plain
# shell here carries the 8 MiB session limit and the pinned offload fails.
#
# OFFLOAD SHAPE (the only legal fit for a ~125B-param FP8 model on 2x 32 GiB
# cards): --offload-backend uva pins the PLE n-gram table (51.2B params) and
# the routed experts in host RAM and reads them in place over PCIe. The
# first start MUST be watched for the UVA engagement log (not just /health):
# a silent fallback to the copy path turns the endpoint into a slideshow.
#
# NO --model-draft, no draft GGUF: vLLM's MTP drafting comes from the model's
# own mtp.* weights via --speculative-config, and only when SPEC_CONFIG is
# non-empty. Bring-up runs with it empty on purpose.
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ENV_FILE="${FRANK_PRIMARY_ENV_FILE:-$SCRIPT_DIR/frank-primary.env}"

if [ ! -f "$ENV_FILE" ]; then
    echo "serve-frank-primary: env file missing: $ENV_FILE" >&2
    exit 1
fi
set -a
# shellcheck disable=SC1090
source "$ENV_FILE"
set +a

: "${SERVED_NAME:?SERVED_NAME must be set}"
: "${BIND_HOST:?BIND_HOST must be set}"
: "${PORT:?PORT must be set}"
: "${MODEL_DIR:?MODEL_DIR must be set}"
: "${TP:?TP must be set}"
: "${GPU_UTIL:?GPU_UTIL must be set}"
: "${CTX_SIZE:?CTX_SIZE must be set}"
: "${MAX_SEQS:?MAX_SEQS must be set}"
: "${OFFLOAD_BACKEND:?OFFLOAD_BACKEND must be set}"
: "${OFFLOAD_GB:?OFFLOAD_GB must be set}"
: "${OFFLOAD_PARAMS:?OFFLOAD_PARAMS must be set}"
: "${VLLM_ROCM_USE_AITER:?VLLM_ROCM_USE_AITER must be set}"
: "${IMAGE:?IMAGE must be set}"
: "${CONTAINER_NAME:?CONTAINER_NAME must be set}"
: "${ROCR_VISIBLE_DEVICES:?ROCR_VISIBLE_DEVICES must be set}"
: "${TOOL_PARSER:?TOOL_PARSER must be set}"
: "${REASONING_PARSER:?REASONING_PARSER must be set}"
: "${CHAT_TEMPLATE_KWARGS:?CHAT_TEMPLATE_KWARGS must be set}"
: "${CACHE_ROOT:?CACHE_ROOT must be set}"

if [ "$OFFLOAD_BACKEND" != "uva" ]; then
    echo "serve-frank-primary: OFFLOAD_BACKEND=$OFFLOAD_BACKEND — this deployment is only qualified for uva" >&2
    exit 1
fi

if ! podman image inspect "$IMAGE" > /dev/null 2>&1; then
    echo "serve-frank-primary: image missing: $IMAGE — build it before serving" >&2
    exit 1
fi

# The model payload is verified by the fetch/verify step (131 shards, index
# cross-checked). The serve-time check is the load-bearing minimum: the
# wrapper config and the shard index must exist, or the container start is a
# guaranteed 20-minute timeout instead of a fast failure.
HOST_MODEL_DIR="/home/dank/models/${MODEL_DIR#/models/}"
if [ ! -f "$HOST_MODEL_DIR/config.json" ] || [ ! -f "$HOST_MODEL_DIR/model.safetensors.index.json" ]; then
    echo "serve-frank-primary: model payload incomplete at $HOST_MODEL_DIR (config.json or index missing)" >&2
    echo "serve-frank-primary: run scripts/download-frank-primary-fp8.sh on this host" >&2
    exit 1
fi

if [ ! -e /dev/kfd ] || [ ! -d /dev/dri ]; then
    echo "serve-frank-primary: /dev/kfd or /dev/dri missing — GPU stack not up" >&2
    exit 1
fi

mkdir -p /home/dank/.cache/frank-primary-vllm

SERVE_ARGS=("$MODEL_DIR"
    --served-model-name "$SERVED_NAME"
    --host "$BIND_HOST"
    --port "$PORT"
    --tensor-parallel-size "$TP"
    --gpu-memory-utilization "$GPU_UTIL"
    --max-model-len "$CTX_SIZE"
    --max-num-seqs "$MAX_SEQS"
    --offload-backend "$OFFLOAD_BACKEND"
    --cpu-offload-gb "$OFFLOAD_GB"
    --cpu-offload-params "$OFFLOAD_PARAMS"
    --enable-auto-tool-choice --tool-call-parser "$TOOL_PARSER"
    --reasoning-parser "$REASONING_PARSER"
    --default-chat-template-kwargs "$CHAT_TEMPLATE_KWARGS"
)

# SPEC_CONFIG (optional, JSON) is how MTP drafting is enabled after the
# bring-up measurement: --speculative-config '{"method":"mtp",...}'.
if [ -n "${SPEC_CONFIG:-}" ]; then
    SERVE_ARGS+=(--speculative-config "$SPEC_CONFIG")
fi
# A crashed/killed unit leaves the container name registered (systemd can
# SIGKILL conmon before the --rm cleanup runs); podman run would then die
# on name-in-use forever. Clear any orphan of OUR name first — idempotent.
podman rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true

exec podman run --rm \
    --name "$CONTAINER_NAME" \
    --userns=keep-id \
    --network=host \
    --ipc=host \
    --security-opt seccomp=unconfined \
    --cap-add SYS_PTRACE \
    --device /dev/kfd \
    --device /dev/dri \
    --group-add keep-groups \
    -e ROCR_VISIBLE_DEVICES="$ROCR_VISIBLE_DEVICES" \
    -e HIP_VISIBLE_DEVICES="$HIP_VISIBLE_DEVICES" \
    -e VLLM_ROCM_USE_AITER="$VLLM_ROCM_USE_AITER" \
    -e NCCL_P2P_DISABLE="$NCCL_P2P_DISABLE" \
    -e TRITON_CACHE_DIR="$CACHE_ROOT/triton" \
    -e TORCHINDUCTOR_CACHE_DIR="$CACHE_ROOT/inductor" \
    -e HOME="$CACHE_ROOT" \
    -e TRITON_ALLOW_NON_CONST_GLOBALS=1 \
    --volume /home/dank/models:/models \
    --volume /home/dank/.cache/frank-primary-vllm:"$CACHE_ROOT" \
    "$IMAGE" \
    "${SERVE_ARGS[@]}"