#!/usr/bin/env bash
# Start vLLM server for a given TP×DP configuration.
# Usage: start_vllm_dp.sh <TP> <DP>
#
# By default this delegates to the Docker start script because the experiment
# is intended to run vLLM inside a container. Set USE_DOCKER=0 to use the
# local VENV_VLLM environment instead.
set -e

TP="${1}"
DP="${2}"

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

if [[ "${USE_DOCKER:-1}" == "1" ]]; then
  exec "${SCRIPT_DIR}/start_vllm_docker.sh" "$@"
fi

RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}"
mkdir -p "${RUNTIME_BASE}/logs" "${RUNTIME_BASE}/tmp"

VENV="${VENV_VLLM}"
export PATH="$VENV/bin:$PATH"
export PYTHONUNBUFFERED=1
export TMPDIR="${RUNTIME_BASE}/tmp"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}"
export VLLM_ENGINE_READY_TIMEOUT_S="${VLLM_ENGINE_READY_TIMEOUT_S:-3600}"

LOG="${RUNTIME_BASE}/logs/${EXPERIMENT}_vllm_tp${TP}_dp${DP}_$(date +%Y%m%d_%H%M%S).log"
PID_FILE="${RUNTIME_BASE}/${EXPERIMENT}_vllm_tp${TP}_dp${DP}.pid"

rm -f "$PID_FILE"

SERVER_ARGS=(
  vllm serve "$MODEL_PATH"
  --trust-remote-code
  --kv-cache-dtype "$KV_CACHE_DTYPE"
  --block-size "$BLOCK_SIZE"
  --tensor-parallel-size "$TP"
  --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION"
  --max-model-len "$MAX_MODEL_LEN"
  --max-num-seqs "$MAX_NUM_SEQS"
  --no-enable-flashinfer-autotune
  --tokenizer-mode "$TOKENIZER_MODE"
  --reasoning-parser "$REASONING_PARSER"
  --spec-method "$SPEC_METHOD"
  --spec-model "$MODEL_PATH"
  --spec-tokens "$SPEC_TOKENS"
  --no-disable-hybrid-kv-cache-manager
  --disable-uvicorn-access-log
  --host 0.0.0.0
  --port "$VLLM_PORT"
)

if [[ "$DP" -gt 1 ]]; then
  SERVER_ARGS+=(
    --data-parallel-size "$DP"
  )
fi

SERVER_ARGS_STR="${SERVER_ARGS[*]}"

echo "=== Starting vLLM server (TP=${TP}, DP=${DP}) ==="
echo "Model: $MODEL_PATH"
echo "Port: $VLLM_PORT"
echo "Command: $SERVER_ARGS_STR"
echo "Log: $LOG"

nohup "${SERVER_ARGS[@]}" > "$LOG" 2>&1 &

PID=$!
echo $PID > "$PID_FILE"
echo "PID: $PID"
echo "Waiting for health on port ${VLLM_PORT}..."

for i in $(seq 1 "${HEALTH_CHECK_RETRIES:-720}"); do
  if curl --fail --silent --show-error --max-time 5 "http://127.0.0.1:${VLLM_PORT}/health" >/dev/null 2>&1; then
    echo "vLLM server is ready at http://127.0.0.1:${VLLM_PORT}"
    echo "Log: $LOG"
    exit 0
  fi
  if ! kill -0 $PID 2>/dev/null; then
    echo "ERROR: vLLM server exited early"
    tail -200 "$LOG"
    exit 1
  fi
  echo "Waiting... ($i/${HEALTH_CHECK_RETRIES:-720})"
  sleep 5
done

echo "ERROR: vLLM server not healthy after ${HEALTH_CHECK_RETRIES:-720} retries"
tail -200 "$LOG"
exit 1
