#!/usr/bin/env bash
# serve-frank-large.sh — frank-large: MiMo-V2.6-Distill-Qwen-9B (BF16) under
# vLLM on the Strix Halo gfx1151 iGPU, :7770.
#
# 2026-10-01 cutover from gufo (Qwen3.8-27B GGUF + DFlash2). The runtime is
# the digest-pinned kyuz0 TheRock image (the fork that carries the gfx1151
# build: ROCm 10 / torch 2.11 / vLLM 0.30 with the Triton Gated DeltaNet
# kernels); like frank-primary, the service drives `podman run` directly and
# then blocks on the container, so a container death is a non-zero ExecStart
# and the unit's Restart=always is what self-heals.
#
# The engine flags are the ones measured in the 2026-10-01 bench window on
# this host (see frank-large.env for the numbers). Every tunable arrives via
# frank-large.env; keys are FRANK_LARGE_-prefixed because HIP reads bare
# TEMP/HOST/GPU_MAX_HEAP_SIZE from the environment.
set -euo pipefail

ENV_FILE="${FRANK_LARGE_ENV_FILE:-/home/dank/serverAppConfigs/scripts/frank-large.env}"
if [[ -f "$ENV_FILE" ]]; then
    set -a
    # shellcheck disable=SC1090
    source "$ENV_FILE"
    set +a
fi

# --- runtime switch -----------------------------------------------------------
# llama: bare llama-server (master build with qwen35 GDN kernels) on GGUF
# weights. vllm (default): the digest-pinned TheRock container below.
FRANK_LARGE_RUNTIME="${FRANK_LARGE_RUNTIME:-vllm}"
if [[ "$FRANK_LARGE_RUNTIME" == "llama" ]]; then
    : "${FRANK_LARGE_LLAMA_BIN:?FRANK_LARGE_LLAMA_BIN not set}"
    : "${FRANK_LARGE_GGUF:?FRANK_LARGE_GGUF not set (target GGUF)}"
    : "${FRANK_LARGE_DRAFT_GGUF:=}"
    : "${FRANK_LARGE_LLAMA_CTX:=1048576}"   # 4 slots x 262,144, f16 KV
    for f in "$FRANK_LARGE_LLAMA_BIN" "$FRANK_LARGE_GGUF"; do
        [[ -x "$f" || -f "$f" ]] || { echo "ERROR: missing $f" >&2; exit 1; }
    done
    DRAFT_ARGS=()
    if [[ -n "$FRANK_LARGE_DRAFT_GGUF" ]]; then
        [[ -f "$FRANK_LARGE_DRAFT_GGUF" ]] || { echo "ERROR: missing draft $FRANK_LARGE_DRAFT_GGUF" >&2; exit 1; }
        DRAFT_ARGS=(--model-draft "$FRANK_LARGE_DRAFT_GGUF" --spec-type draft-dflash)
    fi
    # ComfyUI co-tenancy on this iGPU: gufo-era law, 2 resident queues.
    export GPU_MAX_HW_QUEUES=2
    FRANK_LARGE_PORT="${FRANK_LARGE_PUBLISH:-7770:7770}"
    FRANK_LARGE_PORT="${FRANK_LARGE_PORT##*:}"
    exec "$FRANK_LARGE_LLAMA_BIN" \
        --model "$FRANK_LARGE_GGUF" \
        "${DRAFT_ARGS[@]}" \
        --alias "${FRANK_LARGE_SERVED_NAME:=frank-large}" \
        --host 0.0.0.0 --port "$FRANK_LARGE_PORT" \
        --parallel 4 --ctx-size "$FRANK_LARGE_LLAMA_CTX" \
        --n-gpu-layers 99 --flash-attn on --jinja \
        --reasoning auto --reasoning-format deepseek \
        --cont-batching
fi

# --- below: vllm runtime ------------------------------------------------------
: "${FRANK_LARGE_IMAGE:?FRANK_LARGE_IMAGE not set (digest-pinned kyuz0 vllm-therock image)}"
: "${FRANK_LARGE_MODEL_DIR:?FRANK_LARGE_MODEL_DIR not set}"
: "${FRANK_LARGE_MODEL_SHARDS:?FRANK_LARGE_MODEL_SHARDS not set}"
: "${FRANK_LARGE_CONTAINER:=frank-large-mimo}"
: "${FRANK_LARGE_LEGACY_CONTAINERS:=vllm-mimo}"
: "${FRANK_LARGE_SERVED_NAME:=frank-large}"
: "${FRANK_LARGE_PUBLISH:=7770:7770}"
: "${FRANK_LARGE_MAXLEN:=262144}"
: "${FRANK_LARGE_SEQS:=4}"
: "${FRANK_LARGE_UTIL:=0.45}"
: "${FRANK_LARGE_REASONING_PARSER:=qwen3}"
: "${FRANK_LARGE_TOOL_PARSER:=qwen3_coder}"
: "${FRANK_LARGE_SPEC:=}"
: "${FRANK_LARGE_SPEC_TOKENS:=8}"
: "${FRANK_LARGE_MODELS_DIR:=${FRANK_LARGE_MODEL_DIR%/*}}"
FRANK_LARGE_QUANT="${FRANK_LARGE_QUANT:-}"
FRANK_LARGE_QUANT_CONFIG="${FRANK_LARGE_QUANT_CONFIG:-}"

# --- gates (fail FAST, before touching the running stack) ------------------
# Image gate by REPO DIGEST, not the human-readable ref: rootless podman
# resolves the tag from the network, and the whole point of the pin is to
# boot exactly these bytes offline. Capture first, grep second — piping
# straight into `grep -q` dies of the SIGPIPE race under pipefail.
if ! digests="$(podman images --format '{{.RepoDigests}}')" ||
   ! grep -qF "${FRANK_LARGE_IMAGE#*@}" <<<"$digests"; then
    echo "ERROR: pinned runtime image absent from rootless storage: $FRANK_LARGE_IMAGE" >&2
    echo "frank-large: fetch it (podman pull $FRANK_LARGE_IMAGE)" >&2
    exit 1
fi

# Payload gate: config + weight-map index + the exact shard count. A partial
# payload fails the load minutes later, after the port has been freed and
# the old server stopped.
if [[ ! -f "$FRANK_LARGE_MODEL_DIR/config.json" ]]; then
    echo "ERROR: payload config missing: $FRANK_LARGE_MODEL_DIR/config.json" >&2
    echo "frank-large: re-run scripts/download-frank-large-mimo.sh" >&2
    exit 1
fi
if [[ ! -f "$FRANK_LARGE_MODEL_DIR/model.safetensors.index.json" ]]; then
    echo "ERROR: weight-map index missing: $FRANK_LARGE_MODEL_DIR/model.safetensors.index.json" >&2
    echo "frank-large: re-run scripts/download-frank-large-mimo.sh" >&2
    exit 1
fi
shards_found=$(find "$FRANK_LARGE_MODEL_DIR" -maxdepth 1 -name 'model-*-of-*.safetensors' | wc -l)
if [[ "$shards_found" != "$FRANK_LARGE_MODEL_SHARDS" ]]; then
    echo "ERROR: payload shard count wrong: $FRANK_LARGE_MODEL_DIR (found $shards_found, want $FRANK_LARGE_MODEL_SHARDS)" >&2
    echo "frank-large: re-run scripts/download-frank-large-mimo.sh" >&2
    exit 1
fi
if [[ -n "$FRANK_LARGE_SPEC" ]]; then
    : "${FRANK_LARGE_SPEC_DIR:?FRANK_LARGE_SPEC_DIR not set (host dir holding the draft weights)}"
    # dflash ships ONE model.safetensors; the Qwen3.5-9B MTP head ships as
    # shards behind model.safetensors.index.json. Either shape is a real
    # draft; config.json must be present in both, and a directory with
    # neither weight layout is not a draft at all.
    if [[ ! -f "$FRANK_LARGE_SPEC_DIR/config.json" ]] ||
       { [[ ! -f "$FRANK_LARGE_SPEC_DIR/model.safetensors" ]] &&
         [[ ! -f "$FRANK_LARGE_SPEC_DIR/model.safetensors.index.json" ]]; }; then
        echo "ERROR: spec draft absent or partial: $FRANK_LARGE_SPEC_DIR (FRANK_LARGE_SPEC=$FRANK_LARGE_SPEC)" >&2
        echo "frank-large: re-run scripts/download-frank-large-dflash.sh (or the Qwen3.5-9B fetch for mtp)" >&2
        exit 1
    fi
    # The models tree is mounted read-only at /models; the engine must be
    # given the CONTAINER path to the draft, not the host path.
    FRANK_LARGE_SPEC_CONTAINER_DIR="/models/${FRANK_LARGE_SPEC_DIR##*/}"
fi

# The engine needs the RENDER nodes too, not just /dev/kfd: mounting the
# whole /dev/dri directory (never individual renderD nodes) is the
# verified-working shape (Dankserver 2026-09-29: individual render nodes
# enumerate the device but cannot create a context).
for dev in /dev/kfd /dev/dri; do
    if [[ ! -e "$dev" ]]; then
        echo "ERROR: $dev missing; gfx1151 cannot serve without it" >&2
        exit 1
    fi
done

SPEC_ARGS=()
if [[ -n "$FRANK_LARGE_SPEC" ]]; then
    SPEC_ARGS=(--speculative-config "{\"method\":\"$FRANK_LARGE_SPEC\",\"model\":\"$FRANK_LARGE_SPEC_CONTAINER_DIR\",\"num_speculative_tokens\":$FRANK_LARGE_SPEC_TOKENS}")
fi

QUANT_ARGS=()
if [[ -n "$FRANK_LARGE_QUANT" ]]; then
    # Online quantization: vLLM quantizes the BF16 checkpoint at load time
    # (weights >= 6-bit only; KV is untouched — the context stays bf16).
    QUANT_ARGS=(--quantization "$FRANK_LARGE_QUANT")
    if [[ -n "$FRANK_LARGE_QUANT_CONFIG" ]]; then
        QUANT_ARGS+=(--quantization-config "$FRANK_LARGE_QUANT_CONFIG")
    fi
fi


# Remove stale container identities BEFORE starting: a crashed unit can
# leave the name registered forever after, and the bench-window identity
# (vllm-mimo) holds the GPU and port when a window ends without teardown.
for name in $FRANK_LARGE_LEGACY_CONTAINERS "$FRANK_LARGE_CONTAINER"; do
    if podman ps -a --format '{{.Names}}' | grep -qx "$name"; then
        echo "removing stale container: $name"
        podman rm -f "$name" >/dev/null 2>&1 || true
    fi
done

# --device /dev/mem exists for the fleet's other GPU containers; the GTT
# aperture (114 GiB) needs no size knob. Models mount read-only wholesale so
# spec drafts resolve from disk. THE PATH HANDED TO vllm IS THE CONTAINER
# PATH: the engine runs inside this container, and a host path that is not
# mounted at the same location dies as "Repo id must be in the form
# 'namespace/repo_name'" (2026-10-01 first boot, 51 crash-loop restarts).
MODEL_CONTAINER_DIR="/models/${FRANK_LARGE_MODEL_DIR##*/}"
podman run -d \
    --name "$FRANK_LARGE_CONTAINER" \
    --device /dev/kfd --device /dev/dri --device /dev/mem \
    --security-opt seccomp=unconfined --group-add keep-groups \
    --ulimit memlock=-1:-1 \
    -v "$FRANK_LARGE_MODELS_DIR:/models:ro" \
    -v /home/dank/models/hf-cache:/root/.cache/huggingface \
    -e HSA_ENABLE_SDMA=0 \
    -p "$FRANK_LARGE_PUBLISH" \
    --entrypoint vllm \
    "$FRANK_LARGE_IMAGE" \
    serve "$MODEL_CONTAINER_DIR" \
    --served-model-name "$FRANK_LARGE_SERVED_NAME" \
    --host 0.0.0.0 --port "${FRANK_LARGE_PUBLISH##*:}" \
    --dtype bfloat16 \
    --max-model-len "$FRANK_LARGE_MAXLEN" \
    --max-num-seqs "$FRANK_LARGE_SEQS" \
    --gpu-memory-utilization "$FRANK_LARGE_UTIL" \
    --reasoning-parser "$FRANK_LARGE_REASONING_PARSER" \
    --enable-auto-tool-choice \
    --tool-call-parser "$FRANK_LARGE_TOOL_PARSER" \
    --language-model-only \
    --trust-remote-code \
    "${QUANT_ARGS[@]}" \
    "${SPEC_ARGS[@]}"

# Block on the container so a container death returns non-zero from ExecStart
# and the unit's Restart=always fires. systemd stop goes through the unit's
# ExecStop (podman stop -t); podman wait unblocks when the container is gone.
exec podman wait "$FRANK_LARGE_CONTAINER"