#!/usr/bin/env bash
# Serve a model on ana-ml2:8017 as `probe`, optionally without speculative decoding.
# Usage: serve_probe.sh <model-dir> [nospec]
#
# PROBE_IMAGE must match the image the LIVE seat runs, or the probe measures a
# different vLLM than the one that will serve the model. The Qwen3.8-27B gen seat
# is pinned to a nightly carrying the #51113 qwen3_5_mtp x GDN partial-accept fix;
# probing that model on :latest reproduces the multi-turn corruption we already
# diagnosed and reads as a candidate failure. Default stays :latest for older
# seats; export PROBE_IMAGE to match whatever GEN_IMAGE is.
set -euo pipefail
MODEL="$1"; MODE="${2:-spec}"
IMAGE="${PROBE_IMAGE:-vllm/vllm-openai:latest}"
H=infra-ops@10.250.50.54
# Two separate quoting traps, both of which produced a dead container:
#   1. This string is interpolated into a double-quoted ssh "..." command, so the
#      inner double quotes need escaping or the remote sees {method:qwen3_5_mtp,...}.
#   2. The JSON must ALSO be single-quoted at the REMOTE shell, or bash BRACE
#      EXPANSION splits {"a":1,"b":2} into two words on the comma -- which surfaces
#      as: Value "method":"qwen3_5_mtp" cannot be converted.
SPEC="--speculative-config '{\"method\":\"qwen3_5_mtp\",\"num_speculative_tokens\":3}'"
[ "$MODE" = "nospec" ] && SPEC=""

# The probe must reproduce the LIVE seat's serving surface, not just its weights.
# Without these, surface_test reports failures that are probe-config artifacts:
#   * no --tool-call-parser/--enable-auto-tool-choice -> every tools-bearing
#     request 400s, reading as "tool calling broken in the candidate"
#   * no --reasoning-parser -> the think trace lands in `content` instead of
#     `reasoning_content`, so the thinking-split check measures nothing
# Mirrors stacks/gen-seat/compose.yaml. JSON args are single-quoted for the REMOTE
# shell (brace expansion, see above).
EXTRA="${PROBE_EXTRA_ARGS:---enable-prefix-caching --reasoning-parser qwen3 --enable-auto-tool-choice --tool-call-parser qwen3_coder --limit-mm-per-prompt '{\"image\": 4}' --default-chat-template-kwargs '{\"reasoning_effort\": \"medium\"}'}"

ssh $H "sudo docker rm -f vllm-probe 2>/dev/null >/dev/null || true
sudo docker run -d --name vllm-probe --gpus '\"device=0\"' --ipc host \
  -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
  -v /tank/aimodels:/tank/aimodels -p 8017:8000 \
  $IMAGE \
  $MODEL --served-model-name probe --host 0.0.0.0 --port 8000 \
  --quantization compressed-tensors --gpu-memory-utilization 0.40 \
  --max-model-len ${PROBE_MAXLEN:-32768} --max-num-seqs 16 --max-num-batched-tokens 16384 \
  --trust-remote-code --dtype auto --mamba-cache-dtype float32 \
  --kv-cache-dtype fp8 --enable-chunked-prefill $EXTRA $SPEC >/dev/null"

for i in $(seq 1 60); do
  s=$(curl -s -o /dev/null -w '%{http_code}' -m 3 http://10.250.50.54:8017/health 2>/dev/null || true)
  [ "$s" = "200" ] && { echo "probe healthy after $((i*10))s ($MODEL, $MODE, $IMAGE)"; exit 0; }
  if ! ssh $H 'sudo docker ps -q -f name=vllm-probe' | grep -q .; then
    echo "CONTAINER DIED"; ssh $H 'sudo docker logs --tail 25 vllm-probe 2>&1 | tail -25'; exit 1; fi
  sleep 10
done
echo "TIMEOUT" >&2; exit 1
