#!/usr/bin/env bash
# Run SGLang server inside the already-running benchmark container.
# Usage: run_sglang_in_container.sh <TP> <DP>
set -e

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

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

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

PORT="${SGLANG_PORT:-30031}"
NAME="${EXPERIMENT}_sglang_container"
PID_FILE="${RUNTIME_BASE}/${EXPERIMENT}_sglang_tp${TP}_dp${DP}.pid"

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

# Build server args.
SERVER_ARGS=(
  -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 "$PORT"
)

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

SERVER_ARGS_STR="${SERVER_ARGS[*]}"

echo "=== Starting SGLang server inside container (TP=${TP}, DP=${DP}) ==="
echo "Container: $NAME"
echo "Command: python3 ${SERVER_ARGS_STR}"
echo "Log: $LOG"

# Kill any existing sglang process inside container first.
docker exec "$NAME" pkill -9 -f "sglang.launch_server" 2>/dev/null || true
sleep 3

# Start the SGLang launch module inside the container in background.
# We use nohup so it survives after docker exec returns.
docker exec -d \
  -e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" \
  -e PYTHONUNBUFFERED=1 \
  "$NAME" \
  bash -c "nohup python3 ${SERVER_ARGS_STR} > /tmp/sglang_server.log 2>&1 &"

# Wait for the server process to appear inside container.
sleep 2
SERVER_PID=$(docker exec "$NAME" pgrep -f "sglang.launch_server" | head -n 1 || true)
if [[ -z "$SERVER_PID" ]]; then
  echo "ERROR: sglang server process not found inside container"
  docker exec "$NAME" cat /tmp/sglang_server.log 2>/dev/null | tail -50 || true
  exit 1
fi

echo "Server PID inside container: $SERVER_PID"
echo "$SERVER_PID" > "$PID_FILE"

echo "Waiting for health on port ${PORT}..."
for i in $(seq 1 360); do
  if curl --fail --silent --show-error --max-time 5 "http://127.0.0.1:${PORT}/health" >/dev/null 2>&1; then
    echo "SGLang server is ready at http://127.0.0.1:${PORT}"
    echo "Log: $LOG"
    exit 0
  fi
  # Check if server process is still alive.
  if ! docker exec "$NAME" kill -0 "$SERVER_PID" 2>/dev/null; then
    echo "ERROR: SGLang server exited early"
    docker exec "$NAME" cat /tmp/sglang_server.log 2>/dev/null | tail -200 || true
    exit 1
  fi
  echo "Waiting... ($i/360)"
  sleep 5
done

echo "ERROR: SGLang server not healthy after 360 retries (30 mins)"
docker exec "$NAME" cat /tmp/sglang_server.log 2>/dev/null | tail -200 || true
exit 1
