#!/bin/bash
# Start 1 vLLM TP=8 service on 8x H200.
# Uses all 8 GPUs: (0,1,2,3,4,5,6,7).
# This maximizes single-service throughput for highest-concurrency workloads.
set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/config.env"

VENV="${VENV_SERVER}"
export PATH="$VENV/bin:$PATH"
VLLM="$VENV/bin/vllm"

TP=8
LOG_DIR="${SCRIPT_DIR}/logs"
mkdir -p "$LOG_DIR"

# 1 service on all 8 GPUs.
# Port and GPU mapping.
SERVICES=(
  "30005:0,1,2,3,4,5,6,7"
)

# Kill any existing vLLM processes for this model.
pkill -9 -f "vllm serve.*${MODEL_NAME}" 2>/dev/null || true
sleep 2

# Remove old PID files.
rm -f "${SCRIPT_DIR}"/*.pid

for svc in "${SERVICES[@]}"; do
  PORT="${svc%%:*}"
  GPUS="${svc##*:}"
  LOG="${LOG_DIR}/server_port${PORT}_$(date +%Y%m%d_%H%M%S).log"
  PID_FILE="${SCRIPT_DIR}/server_port${PORT}.pid"

  echo "=== Starting service on port ${PORT}, GPUs ${GPUS} ==="
  echo "Log: ${LOG}"

  CUDA_VISIBLE_DEVICES="${GPUS}" \
  nohup "$VLLM" serve "$MODEL_PATH" \
      --trust-remote-code \
      --tensor-parallel-size "$TP" \
      --kv-cache-dtype fp8 \
      --block-size 256 \
      --served-model-name "$SERVED_MODEL_NAME" \
      --port "$PORT" \
      > "$LOG" 2>&1 &
  PID=$!
  echo $PID > "$PID_FILE"
  echo "PID: $PID"
done

echo ""
echo "Waiting for service to become healthy..."

MAX_WAIT=240
all_healthy=0
for i in $(seq 1 $MAX_WAIT); do
  all_healthy=1
  for svc in "${SERVICES[@]}"; do
    PORT="${svc%%:*}"
    if ! curl -s "http://127.0.0.1:${PORT}/health" > /dev/null 2>&1; then
      all_healthy=0
      break
    fi
  done

  if [[ "$all_healthy" -eq 1 ]]; then
    echo "Service is healthy!"
    for svc in "${SERVICES[@]}"; do
      PORT="${svc%%:*}"
      echo "  - http://127.0.0.1:${PORT}"
    done
    exit 0
  fi

  # Check if any process died early.
  for svc in "${SERVICES[@]}"; do
    PORT="${svc%%:*}"
    PID_FILE="${SCRIPT_DIR}/server_port${PORT}.pid"
    if [[ -f "$PID_FILE" ]]; then
      PID="$(cat "$PID_FILE")"
      if ! kill -0 "$PID" 2>/dev/null; then
        echo "ERROR: Service on port ${PORT} (PID ${PID}) exited early"
        LOG="${LOG_DIR}/server_port${PORT}_*.log"
        tail -200 $LOG 2>/dev/null || true
        exit 1
      fi
    fi
  done

  if (( i % 10 == 0 )); then
    echo "Waiting... (${i}/${MAX_WAIT})"
  fi
  sleep 5
done

echo "ERROR: Service did not become healthy after ${MAX_WAIT} retries"
for svc in "${SERVICES[@]}"; do
  PORT="${svc%%:*}"
  LOG="${LOG_DIR}/server_port${PORT}_*.log"
  echo "--- Last 50 lines of port ${PORT} ---"
  tail -50 $LOG 2>/dev/null || true
done
exit 1