85 lines
3.4 KiB
Bash
Executable File
85 lines
3.4 KiB
Bash
Executable File
#!/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 "$@"
|