#!/bin/bash
# One vLLM server per GPU, to use as the teacher for gen_data.py.
#
# Usage: teachers.sh start | wait | stop
#   start  launch a server on each GPU in the background (own process group)
#   wait   block until every server answers /health (TIMEOUT seconds)
#   stop   terminate the servers started by this script
#
# Env:   TEACHER_MODEL  model path or Hugging Face id (required for start)
#        GPUS           space-separated GPU ids              (default: "0 1")
#        PORT_BASE      first port; GPU i uses PORT_BASE+i   (default: 19400)
#        SERVED_NAME    model name the servers answer to     (default: teacher)
#        MAX_SEQS       concurrent sequences per server      (default: 24)
#        MAX_LEN        context length                       (default: 4096)
#        GPU_UTIL       fraction of GPU memory to use        (default: 0.92)
#        VLLM           vllm executable                      (default: vllm)
#        VLLM_ARGS      extra arguments for `vllm serve`, e.g. quantization or
#                       dtype flags your model needs
#        LOGS           directory for logs and pid files     (default: ./logs)
#        TIMEOUT        seconds for `wait`                   (default: 900)
set -euo pipefail
LOGS=${LOGS:-./logs}
GPUS=${GPUS:-"0 1"}
PORT_BASE=${PORT_BASE:-19400}
cmd=${1:-}

ports() { local i=0; for _ in $GPUS; do echo $((PORT_BASE + i)); i=$((i + 1)); done; }

case "$cmd" in
  start)
    : "${TEACHER_MODEL:?set TEACHER_MODEL to a model path or id}"
    mkdir -p "$LOGS"
    i=0
    for g in $GPUS; do
      port=$((PORT_BASE + i)); i=$((i + 1))
      [ -e "$LOGS/teacher$g.pid" ] && kill -0 "$(cat "$LOGS/teacher$g.pid")" 2>/dev/null \
        && { echo "teacher on GPU $g already running" >&2; exit 1; }
      # setsid makes the server its own process-group leader, so $! is the group
      # id that `stop` signals.
      # shellcheck disable=SC2086
      CUDA_VISIBLE_DEVICES=$g setsid nohup "${VLLM:-vllm}" serve "$TEACHER_MODEL" \
        --host 127.0.0.1 --port "$port" --served-model-name "${SERVED_NAME:-teacher}" \
        --max-model-len "${MAX_LEN:-4096}" --max-num-seqs "${MAX_SEQS:-24}" \
        --gpu-memory-utilization "${GPU_UTIL:-0.92}" --enable-prefix-caching \
        ${VLLM_ARGS:-} > "$LOGS/teacher$g.log" 2>&1 &
      echo $! > "$LOGS/teacher$g.pid"
      echo "GPU $g -> http://127.0.0.1:$port (log $LOGS/teacher$g.log)"
    done
    ;;
  wait)
    deadline=$((SECONDS + ${TIMEOUT:-900}))
    i=0
    for g in $GPUS; do
      port=$((PORT_BASE + i)); i=$((i + 1))
      until curl -fsS -o /dev/null "http://127.0.0.1:$port/health" 2>/dev/null; do
        if ! kill -0 "$(cat "$LOGS/teacher$g.pid" 2>/dev/null || echo 0)" 2>/dev/null; then
          echo "teacher on GPU $g exited; last log lines:" >&2
          tail -n 15 "$LOGS/teacher$g.log" >&2
          exit 1
        fi
        [ $SECONDS -lt $deadline ] || { echo "port $port not healthy in time" >&2; exit 1; }
        sleep 5
      done
      echo "GPU $g: port $port healthy"
    done
    ;;
  stop)
    for f in "$LOGS"/teacher*.pid; do
      [ -e "$f" ] || continue
      kill -TERM -- "-$(cat "$f")" 2>/dev/null || true
      rm -f "$f"
    done
    ;;
  *)
    echo "usage: $0 start|wait|stop" >&2; exit 2 ;;
esac