#!/usr/bin/env bash
# ============================================================
# Start SGLang multi-node TP=16 EP=2 across 2 nodes.
# Key: SGLANG_SHARED_EXPERT_TP1=1 + --ep-size 2
# Usage: ./start_sglang_multinode.sh [--dry-run]
# ============================================================
set -Eeuo pipefail

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

DRY_RUN="${DRY_RUN:-0}"
HEAD_CONTAINER="${EXPERIMENT}_head"
WORKER_CONTAINER="${EXPERIMENT}_worker"

log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"; }

build_docker_cmd() {
  local node_rank="$1" container_name="$2"
  cat <<DOCKER_CMD
docker rm -f ${container_name} 2>/dev/null;
docker run -d --name ${container_name} \\
  --gpus all --network host --ipc=host --ulimit memlock=-1 --shm-size=20g \\
  -v ${MODEL_PATH}:${MODEL_PATH}:ro \\
  -e NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME} \\
  -e NCCL_DEBUG=${NCCL_DEBUG} \\
  -e SGLANG_SHARED_EXPERT_TP1=1 \\
  ${DOCKER_IMAGE} \\
  python3 -m sglang.launch_server \\
    --model-path ${MODEL_PATH} \\
    --tp-size ${TP_SIZE} --ep-size ${EP_SIZE} --nnodes ${NNODES} \\
    --node-rank ${node_rank} --dist-init-addr ${HEAD_IP}:${DIST_INIT_PORT} \\
    --trust-remote-code --host 0.0.0.0 --port ${SGLANG_PORT} \\
    --mem-fraction-static ${MEM_FRACTION_STATIC} \\
    --cuda-graph-max-bs-decode ${CUDA_GRAPH_MAX_BS_DECODE} \\
    --max-running-requests ${MAX_RUNNING_REQUESTS}
DOCKER_CMD
}

main() {
  log "=== SGLang multi-node TP=${TP_SIZE} EP=${EP_SIZE} ==="

  if [[ "$DRY_RUN" == "0" ]]; then
    ssh -o StrictHostKeyChecking=no "${HEAD_NODE}" "docker rm -f ${HEAD_CONTAINER} 2>/dev/null" || true
    ssh -o StrictHostKeyChecking=no "${WORKER_NODE}" "docker rm -f ${WORKER_CONTAINER} 2>/dev/null" || true
    sleep 2
  fi

  log "Starting WORKER (${WORKER_NODE}, node_rank=1) ..."
  local cmd; cmd="$(build_docker_cmd 1 "${WORKER_CONTAINER}")"
  if [[ "$DRY_RUN" == "1" ]]; then log "[DRY] $cmd"; else
    ssh -o StrictHostKeyChecking=no "${WORKER_NODE}" "bash -c '${cmd}'" >/dev/null
    log "Worker started."
  fi
  sleep 5

  log "Starting HEAD (${HEAD_NODE}, node_rank=0) ..."
  cmd="$(build_docker_cmd 0 "${HEAD_CONTAINER}")"
  if [[ "$DRY_RUN" == "1" ]]; then log "[DRY] $cmd"; else
    ssh -o StrictHostKeyChecking=no "${HEAD_NODE}" "bash -c '${cmd}'" >/dev/null
    log "Head started."
  fi

  # Health check
  if [[ "$DRY_RUN" == "1" ]]; then log "[DRY] Skip health check"; return 0; fi
  log "Waiting for health on ${HEAD_IP}:${SGLANG_PORT} ..."
  for (( i=1; i<=${HEALTH_CHECK_RETRIES}; i++ )); do
    if curl --fail --silent --max-time 5 "http://${HEAD_IP}:${SGLANG_PORT}/health" >/dev/null 2>&1; then
      log "Server healthy! (${i}s)"
      return 0
    fi
    local hok=$(ssh -o StrictHostKeyChecking=no "${HEAD_NODE}" \
      "docker ps --filter name=${HEAD_CONTAINER} --format '{{.Names}}' 2>/dev/null" || true)
    local wok=$(ssh -o StrictHostKeyChecking=no "${WORKER_NODE}" \
      "docker ps --filter name=${WORKER_CONTAINER} --format '{{.Names}}' 2>/dev/null" || true)
    if [[ -z "$hok" ]]; then log "ERROR: Head died"; ssh "${HEAD_NODE}" "docker logs --tail 30 ${HEAD_CONTAINER}" || true; return 1; fi
    if [[ -z "$wok" ]]; then log "ERROR: Worker died"; ssh "${WORKER_NODE}" "docker logs --tail 30 ${WORKER_CONTAINER}" || true; return 1; fi
    if (( i % 12 == 0 )); then log "  waiting... ${i}s"; fi
    sleep "${HEALTH_CHECK_INTERVAL_S}"
  done
  log "ERROR: Timeout"; return 1
}

main "$@"