#!/usr/bin/env bash
# Start SGLang server for a given TP×DP configuration.
# Usage: start_sglang_dp.sh <TP> <DP>
#
# By default this delegates to the Docker start script because the experiment
# is intended to run SGLang inside a container. Set USE_DOCKER=0 to use the
# local VENV_CLIENT 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_sglang_docker.sh" "$@"
fi

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

VENV="${VENV_CLIENT}"
export PATH="$VENV/bin:$PATH"
export PYTHONUNBUFFERED=1
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
export TMPDIR="${RUNTIME_BASE}/tmp"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}"

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

rm -f "$PID_FILE"

SERVER_ARGS=(
  python3 -m sglang.launch_server
  --model-path "$MODEL_PATH"
  --trust-remote-code
  --tp-size "$TP"
  --moe-runner-backend "$MOE_RUNNER_BACKEND"
  --mem-fraction-static "$MEM_FRACTION_STATIC"
  --context-length "$CONTEXT_LENGTH"
  --max-running-requests "$MAX_RUNNING_REQUESTS"
  --host 0.0.0.0
  --port "$SGLANG_PORT"
)

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

SERVER_ARGS_STR="${SERVER_ARGS[*]}"

echo "=== Starting SGLang server (TP=${TP}, DP=${DP}) ==="
echo "Model: $MODEL_PATH"
echo "Port: $SGLANG_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 ${SGLANG_PORT}..."

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

echo "ERROR: SGLang server not healthy after 240 retries"
tail -200 "$LOG"
exit 1