scripts/server_orchestrator.sh

#!/usr/bin/env bash
# Server orchestrator — runs INSIDE the GPU job. Drives llama-server configs,
# marks readiness, and signals the remote benchmark driver via a shared volume.
#
# Env:
#   BENCH_PLAN   e.g. "A B C D"   (subset of A B C D E)
#   NMAX         dflash n-max value (default 7)
#   CTX          context size (default 32768)
#   EXTRA_ARGS   appended to every llama-server invocation
#   LOG_VERBOSE  1 = server logs to files with speculation stats (default 1)
set -euo pipefail

MODELS=/artifacts/models
BIN_DIR=/llamacpp/bin
RUN_DIR=/bench/run
LOG_DIR=/bench/logs
READY=/bench/ready.flag
mkdir -p "$RUN_DIR" "$LOG_DIR"

Q5="$MODELS/Qwen3.8-27B-OBLITERATED-Q5_K_M.gguf"
# bucket no longer holds the plain name when variants exist; only require it
# as readable if neither variant file is present
Q5_V2="$MODELS/Qwen3.8-27B-OBLITERATED-Q5_K_M-V2.gguf"
Q5_V3="$MODELS/Qwen3.8-27B-OBLITERATED-Q5_K_M-V3.gguf"
Q4="$MODELS/Qwen3.8-27B-OBLITERATED-Q4_K_M.gguf"
DFT="$MODELS/Qwen3.8-27B-DFlash2-Q8_0.gguf"
MTP="$MODELS/mtp-Qwen3.8-27B-Q8_0.gguf"
DSP="$MODELS/Qwen3.8-27B-DSpark-Q8_0.gguf"

echo "waiting for model files to be readable via bucket mount ..."
for i in $(seq 1 60); do
  ok=1
  for f in "$Q4" "$DFT" "$MTP"; do
    if ! head -c 16 "$f" >/dev/null 2>&1; then ok=0; echo "not readable yet: $f"; fi
  done
  # Q5 target: accept whichever variant exists (plain, V2, or V3)
  if ! head -c 16 "$Q5" >/dev/null 2>&1 && \
     ! head -c 16 "$Q5_V2" >/dev/null 2>&1 && \
     ! head -c 16 "$Q5_V3" >/dev/null 2>&1; then
    ok=0; echo "no readable Q5 variant (plain/V2/V3)"
  fi
  [ "$ok" = "1" ] && break
  sleep 3
done
for f in "$Q4" "$DFT" "$MTP"; do
  head -c 16 "$f" >/dev/null 2>&1 || { echo "MISSING/unreadable $f"; exit 1; }
done
head -c 16 "$Q5" >/dev/null 2>&1 || head -c 16 "$Q5_V2" >/dev/null 2>&1 || head -c 16 "$Q5_V3" >/dev/null 2>&1 \
  || { echo "MISSING all Q5 variants"; exit 1; }
echo "all models readable"
ls -la "$MODELS/" | grep -E "gguf|V2|V3"

mark_ready() { echo "$1" > "$READY"; }

launch() {
  local name="$1"; shift   # remaining "$@" = binary + its arguments
  echo "=== LAUNCH $name : $(date -u +%H:%M:%S) ==="
  rm -f "$READY"
  local logfile="$LOG_DIR/server_${name}.log"
  # stderr+stdout -> logfile; server timings + draft-acceptance stats land here
  nohup "$@" >"$logfile" 2>&1 &
  SERVER_PID=$!
  echo "$SERVER_PID" > "$RUN_DIR/server.pid"
  echo "launched pid=$SERVER_PID cmd: $*"
}

http_ok() {
  # dependency-free HTTP probe (runtime image has no curl/wget/python)
  local path="$1" head
  (exec 3<>"/dev/tcp/127.0.0.1/${PORT:-8080}" || return 1
   printf 'GET %s HTTP/1.0\r\nHost: localhost\r\n\r\n' "$path" >&3
   head=$(head -1 <&3 2>/dev/null)
   exec 3<&- 3>&-
   [[ "$head" == *200* ]]) 2>/dev/null
}

wait_healthy() {
  local name="$1"
  for i in $(seq 1 240); do
    if http_ok "/health"; then
      echo "healthy after ~${i}x5s"
      mark_ready "$name"
      return 0
    fi
    if ! kill -0 "$SERVER_PID" 2>/dev/null; then
      echo "SERVER DIED during load:"; tail -30 "$LOG_DIR/server_${name}.log" || true
      exit 1
    fi
    sleep 5
  done
  echo "TIMEOUT waiting healthy"; tail -30 "$LOG_DIR/server_${name}.log" || true
  exit 1
}

stop_server() {
  if [ -f "$RUN_DIR/server.pid" ]; then
    local pid; pid=$(cat "$RUN_DIR/server.pid")
    kill "$pid" 2>/dev/null || true
    for i in $(seq 1 20); do kill -0 "$pid" 2>/dev/null || break; sleep 1; done
    kill -9 "$pid" 2>/dev/null || true
    sleep 3
  fi
  rm -f "$RUN_DIR/server.pid"
}

BASE_ARGS=(--host 0.0.0.0 --port 8080 -c "${CTX:-32768}" -fa on --jinja \
           --threads 8 -b 4096 -ub 2048 -np 1 --no-webui)

echo "plan: ${BENCH_PLAN:-A C}"

# Variant GGUF resolution (pre-staged in bucket by hf buckets cp):
#   M_VARIANT=V2 -> *-V2.gguf | V3 -> *-V3.gguf | blank -> bucket default name
#   falls back to the default Q5_K_M.gguf if the variant file is absent
Q5_V2="$MODELS/Qwen3.8-27B-OBLITERATED-Q5_K_M-V2.gguf"
Q5_V3="$MODELS/Qwen3.8-27B-OBLITERATED-Q5_K_M-V3.gguf"
select_q5() {
  local v="${M_VARIANT:-}"
  case "$v" in
    V2) [ -f "$Q5_V2" ] && { echo "$Q5_V2"; return; } ;;
    V3) [ -f "$Q5_V3" ] && { echo "$Q5_V3"; return; } ;;
  esac
  echo "$Q5"
}

IFS=' ' read -ra PLAN <<< "${BENCH_PLAN:-A C}"

for cfg in "${PLAN[@]}"; do
  Q5_THIS=$(select_q5)
  echo "config $cfg uses target: $Q5_THIS"
  case "$cfg" in
    B)  # Q5 + native MTP (stock-trained head grafted onto OBLITERATUS target)
      launch B "$BIN_DIR/llama-server" -m "$Q5_THIS" "${BASE_ARGS[@]}" \
        -md "$MTP" --spec-type draft-mtp --spec-draft-n-max "${NMAX_MTP:-4}" \
        -ngld 999 ${EXTRA_ARGS:-} 2>"$LOG_DIR/server_B.log"
      ;;
    A)  # Q5 baseline
      launch A "$BIN_DIR/llama-server" -m "$Q5_THIS" "${BASE_ARGS[@]}" \
        ${EXTRA_ARGS:-} 2>"$LOG_DIR/server_A.log"
      ;;
    C3)  # Q5 + DFlash2, short drafts (n-max 3)
      NMAX=3
      launch C "$BIN_DIR/llama-server" -m "$Q5" "${BASE_ARGS[@]}" \
        -md "$DFT" --spec-type draft-dflash --spec-draft-n-max "${NMAX}" \
        -ngld 999 ${EXTRA_ARGS:-}
      ;;
    S)  # Q5 + DSpark (Markov-head drafter, block size 7, dynamic length)
      launch S "$BIN_DIR/llama-server" -m "$Q5_THIS" "${BASE_ARGS[@]}" \
        -md "$DSP" --spec-type draft-dspark --spec-draft-n-max 7 \
        -ngld 999 ${EXTRA_ARGS:-}
      ;;
    S3)  # Q5 + DSpark, capped at 3 draft tokens
      launch S "$BIN_DIR/llama-server" -m "$Q5_THIS" "${BASE_ARGS[@]}" \
        -md "$DSP" --spec-type draft-dspark --spec-draft-n-max 3 \
        -ngld 999 ${EXTRA_ARGS:-}
      ;;
    C2)  # Q5 + DFlash2, shortest drafts (n-max 2)
      NMAX=2
      launch C "$BIN_DIR/llama-server" -m "$Q5" "${BASE_ARGS[@]}" \
        -md "$DFT" --spec-type draft-dflash --spec-draft-n-max "${NMAX}" \
        -ngld 999 ${EXTRA_ARGS:-}
      ;;
    B3)  # Q5 + native MTP, n-max 3
      NMAX_MTP=3
      launch B "$BIN_DIR/llama-server" -m "$Q5" "${BASE_ARGS[@]}" \
        -md "$MTP" --spec-type draft-mtp --spec-draft-n-max "${NMAX_MTP}" \
        -ngld 999 ${EXTRA_ARGS:-}
      ;;
    E3)  # Q4 + native MTP, n-max 3
      NMAX_MTP=3
      launch E "$BIN_DIR/llama-server" -m "$Q4" "${BASE_ARGS[@]}" \
        -md "$MTP" --spec-type draft-mtp --spec-draft-n-max "${NMAX_MTP}" \
        -ngld 999 ${EXTRA_ARGS:-}
      ;;
    C5)  # Q5 + DFlash2, medium drafts (n-max 5)
      NMAX=5
      launch C "$BIN_DIR/llama-server" -m "$Q5" "${BASE_ARGS[@]}" \
        -md "$DFT" --spec-type draft-dflash --spec-draft-n-max "${NMAX}" \
        -ngld 999 ${EXTRA_ARGS:-}
      ;;
    C)  # Q5 + DFlash2 (NMAX controls draft length)
      launch C "$BIN_DIR/llama-server" -m "$Q5" "${BASE_ARGS[@]}" \
        -md "$DFT" --spec-type draft-dflash --spec-draft-n-max "${NMAX:-7}" \
        -ngld 999 ${EXTRA_ARGS:-}
      ;;
    D)  # Q4 + DFlash2
      launch D "$BIN_DIR/llama-server" -m "$Q4" "${BASE_ARGS[@]}" \
        -md "$DFT" --spec-type draft-dflash --spec-draft-n-max "${NMAX:-7}" \
        -ngld 999 ${EXTRA_ARGS:-} 2>"$LOG_DIR/server_D.log"
      ;;
    E)  # Q4 + native MTP
      launch E "$BIN_DIR/llama-server" -m "$Q4" "${BASE_ARGS[@]}" \
        -md "$MTP" --spec-type draft-mtp --spec-draft-n-max "${NMAX_MTP:-4}" \
        -ngld 999 ${EXTRA_ARGS:-} 2>"$LOG_DIR/server_E.log"
      ;;
    *)  echo "unknown config '$cfg'"; exit 1;;
  esac
  wait_healthy "$cfg"
  echo "=== READY $cfg — waiting for driver signal to move on ==="
  # IDLE WATCHDOG: if no driver action within IDLE_TIMEOUT (default 20 min),
  # self-terminate instead of burning metered GPU time overnight.
  IDLE_TIMEOUT="${IDLE_TIMEOUT:-1200}"
  idle_start=$(date +%s)
  moved_on=false
  while true; do
    if [ -f /bench/STOP ]; then echo "STOP received"; moved_on="stop"; break; fi
    if [ ! -f "$READY" ] && [ -f "/bench/next" ]; then
      nxt=$(cat /bench/next); rm -f /bench/next
      echo "driver requested next: $nxt"
      moved_on=true
      break
    fi
    now=$(date +%s)
    if [ $((now - idle_start)) -ge "$IDLE_TIMEOUT" ]; then
      echo "IDLE TIMEOUT (${IDLE_TIMEOUT}s) with no driver activity — self-terminating to stop billing"
      touch /bench/STOP
      moved_on="stop"
      break
    fi
    sleep 3
  done
  # dump this config's speculation/timing evidence into RETAINED job logs
  echo "=== SPEC STATS $cfg ==="
  grep -aE "draft acceptance|acc per pos" "$LOG_DIR/server_${cfg}.log" | tail -8 || true
  grep -aE "prompt eval time|eval time" "$LOG_DIR/server_${cfg}.log" | tail -6 || true
  echo "=== END SPEC STATS $cfg ==="
  if [ "$moved_on" = "stop" ]; then break; fi
  stop_server
done

stop_server || true
# archive all server logs to the bucket (rw mount) for post-mortem analysis
if mkdir -p /artifacts/logs 2>/dev/null; then
  cp "$LOG_DIR"/server_*.log /artifacts/logs/ 2>/dev/null || true
  echo "logs archived to bucket"
fi
nvidia-smi --query-gpu=name,memory.total,memory.used --format=csv > /tmp/gpu_final.txt 2>&1 || true
cat /tmp/gpu_final.txt || true
echo "=== ORCHESTRATOR DONE ==="