#!/usr/bin/env bash
# Start SGLang server in Docker for a given TPxDP configuration.
# Usage: start_sglang_docker.sh <TP> <DP>
#
# Uses the lmsysorg/sglang image and the same argument set as the
# bare-metal start script. The container is removed automatically on stop.
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" "${RUNTIME_BASE}/tmp"

IMAGE="${DOCKER_IMAGE:-lmsysorg/sglang:latest}"
PORT="${SGLANG_PORT:-30031}"
NAME="${EXPERIMENT}_sglang_tp${TP}_dp${DP}"
PID_FILE="${RUNTIME_BASE}/${EXPERIMENT}_sglang_tp${TP}_dp${DP}.pid"

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

# Clean up any stale container with the same name.
docker rm -f "$NAME" >/dev/null 2>&1 || true

SERVER_ARGS=(
  serve
  --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

if [[ -n "${CUDA_GRAPH_MAX_BS_DECODE:-}" ]]; then
  SERVER_ARGS+=(--cuda-graph-max-bs-decode "$CUDA_GRAPH_MAX_BS_DECODE")
fi

SERVER_ARGS_STR="${SERVER_ARGS[*]}"

echo "=== Starting SGLang server in Docker (TP=${TP}, DP=${DP}) ==="
echo "Image: $IMAGE"
echo "Model: $MODEL_PATH"
echo "Container name: $NAME"
echo "Host port: $PORT"
echo "Command: sglang ${SERVER_ARGS_STR}"
echo "Log: $LOG"

# Run docker in the foreground so that killing the host process stops the
# container (the --rm flag ensures cleanup). nohup lets us background it and
# capture the host PID in the same way as the bare-metal start script.
nohup docker run --rm \
  --name "$NAME" \
  --gpus all \
  --ipc host \
  --shm-size 16g \
  --entrypoint sglang \
  -p "${PORT}:${PORT}" \
  -v "${MODEL_PATH}:${MODEL_PATH}:ro" \
  -v "${RUNTIME_BASE}/tmp:/tmp" \
  -e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" \
  -e PYTHONUNBUFFERED=1 \
  "$IMAGE" \
  "${SERVER_ARGS[@]}" \
  > "$LOG" 2>&1 &

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

for i in $(seq 1 240); 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
  if ! kill -0 $PID 2>/dev/null; then
    echo "ERROR: Docker SGLang server exited early"
    tail -200 "$LOG"
    exit 1
  fi
  echo "Waiting... ($i/240)"
  sleep 5
done

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