#!/usr/bin/env bash
# run-swift15-bf16.sh — RUNS ON DANKSERVER. Boot the no-recipe bf16 Swift1.5
# container for the two offline campaign phases. The engine owns the GPUs, so
# the production unit must be stopped first (this script does it and restores
# on exit, autoresearch pattern: stop, never mask — relinked unit).
#
#   kld    --kld-record the reference logits from the bf16 model
#          (runs the corpus and EXITS; engine writes DIR and leaves)
#   calib  serve with the calibration tap; a separate driver
#          (scripts/calib-push.py) pushes the corpus; the engine writes the
#          Hessians at RADIANCE_CALIB_TOKENS and the tap switches itself off.
#          Stop this script after calib-push prints TAP_DONE.
#
# Shape = the vendor's bf16 calibration lines (GUIDE.md §4.3:528 + MOE-W4.md
# "calibrate the checkpoint as it ships"): expert_tiered + weights-disk-tier,
# no speculation, deterministic. Large prefill steps because a disk-tier step
# costs the same at 1 token as at 32K. The container declared --max-ctx 4096;
# --max-model-len 16384 exceeds it on purpose (the KLD prompts run to 12,007
# tokens, calib to 8,242) — the boot ledger arbitrates the QSA buffers; if
# declare refuses, re-convert with a wider --max-ctx.
set -uo pipefail

MODE="${1:?usage: run-swift15-bf16.sh kld|calib}"
PIN_1013=sha256:4c3f484564bb2872784e46b0a59cbd71becc9c1e746e301a3596ffae253c322e
PIN_120=sha256:665775ed1d90de6dbb25558e07fc60d20328f1346f96f3c6e00845eb763f01fa
RAD=/home/dank/models/rad/qwen4exp-swift15-bf16.rad
C=swift15ref
LOG=/home/dank/logs/swift15_bf16.log
KLD_REF_DIR=/home/dank/corpora/kld-ref
CALIB_DIR=/home/dank/corpora/calib
KLD_CORPUS=/home/dank/corpora/out/kld.jsonl

restore() { sudo systemctl start frank-archon >/dev/null 2>&1 || true; }
trap restore EXIT
sudo systemctl stop frank-archon >/dev/null 2>&1 || true
podman rm -f "$C" >/dev/null 2>&1 || true

ARGS=(--model /models/rad/qwen4exp-swift15-bf16.rad
      --tp 2 --placement expert_tiered --weights-disk-tier
      --host-pool-mib 12288 --deterministic --gpu-headroom-mib 768
      --num-speculative-tokens 0
      --kv-cache-dtype fp8 --max-model-len 16384
      --ngram-placement ram
      --override-chat-template /models/tpl.jinja)

case "$MODE" in
  kld)
    # 1.0.13: the engine that will serve the candidate — the reference is
    # recorded on the same engine that scores it. The weights slab DERIVES as
    # card - headroom - kv, and the disk tier refuses under 9.57 GiB (4.21 of
    # layer buffers + 5.36 fixed), so kv is explicit SMALL: recording is
    # prefill-only, self-batched at MNBT 32768 = two 16k prompts at a time;
    # 4 slots x 16k fp8 (29 KiB/pos, 1.7 GiB) + headroom fits 3072.
    IMAGE="stilldeadcode/radiance@$PIN_1013"
    mkdir -p "$KLD_REF_DIR"
    ARGS+=(--kld-record /data/kld-ref --kld-corpus /data/kld.jsonl
           --vram-kv-mib 3072
           --max-num-seqs 4 --max-num-batched-tokens 32768)
    ;;
  calib)
    # 1.2.0: the calibration tap (RADIANCE_CALIB_DIR + RADIANCE_CALIB_TOKENS)
    # lives in the 1.2.0 qwen4exp plugin; 1.0.13 has only the DIR half with an
    # unparameterisable threshold. Grams are bf16 activations — engine-invariant.
    # Same KV ceiling as kld: the disk-tier weights floor (9.57 GiB) plus
    # ~18 GiB of fixed per-rank claims cap the pool near 3.5 GiB on a 32 GiB
    # card — pool minus the 2.17 GiB fixed holds ~237k fp8 tokens a rank, so
    # the vendor's big-step calibration shape is affordable: MOE-W4.md:244
    # measures 640 tok/s at a 32k step against 930 at 55k, and the doc's own
    # run drives MNBT 65536. 8 slots x 16k = 131k resident tokens fits the
    # pool; the step batches 65k. KLD's observed rate confirmed step size,
    # not slot count, is the throughput lever.
    IMAGE="stilldeadcode/radiance@$PIN_120"
    mkdir -p "$CALIB_DIR"
    ARGS+=(--port 7770 --vram-kv-mib 3072
           --max-num-seqs 8 --max-num-batched-tokens 65536
           --prefix-cache-dir /data/calib-kvcache)
    ;;
  probe)
    IMAGE="stilldeadcode/radiance@$PIN_1013"
    ARGS+=(--port 7770 --max-num-seqs 8 --max-num-batched-tokens 4096)
    ;;
  *) echo "unknown mode: $MODE" >&2; exit 2 ;;
esac

# The prod template copy, readable from the container: mount the models tree
# ro plus the state dir rw (KLD record + prefix tier land there).
cp -f /home/dank/serverAppConfigs/scripts/qwen38-froggeric-v22.5.jinja \
      /home/dank/models/tpl.jinja 2>/dev/null || true

KLD_MOUNT=()
if [ "$MODE" = kld ]; then
    # Nested under the /data mount, so these must come after it (line order
    # in podman run below): the corpus read-only, the record dir writable.
    KLD_MOUNT=(-v "$KLD_CORPUS:/data/kld.jsonl:ro"
               -v "$KLD_REF_DIR:/data/kld-ref")
fi
EXTRA_ENV=()
if [ "$MODE" = calib ]; then
    EXTRA_ENV=(-v "$CALIB_DIR:/data/calib"
               -e RADIANCE_CALIB_DIR=/data/calib
               -e RADIANCE_CALIB_TOKENS=10000000)
fi

podman run --rm --name "$C" \
    --network host --init \
    --security-opt seccomp=unconfined \
    --ulimit memlock=-1:-1 \
    --device /dev/kfd --device /dev/dri \
    -v /home/dank/models:/models:ro \
    -v /home/dank/radiance-state:/data \
    "${KLD_MOUNT[@]}" \
    "${EXTRA_ENV[@]}" \
    "$IMAGE" "${ARGS[@]}" 2>&1 | tee "$LOG"
rc=${PIPESTATUS[0]}
echo "swift15-bf16 $MODE exited rc=$rc; log: $LOG"
exit "$rc"