390 lines
13 KiB
Bash
Executable File
390 lines
13 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# Deploy Kimi-K3 standard PD on 601-608. Speculative decoding stays disabled.
|
|
set -Eeuo pipefail
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
# shellcheck source=/dev/null
|
|
source "${SCRIPT_DIR}/config.env"
|
|
|
|
ACTION="${1:-status}"
|
|
RUN_ID="${RUN_ID:-kimi3-pd-standard-$(date '+%Y%m%d-%H%M%S')}"
|
|
RESULT_ROOT="${RESULT_ROOT:-${SCRIPT_DIR}/results/${RUN_ID}}"
|
|
SUDO_PASSWORD_FILE="${SUDO_PASSWORD_FILE:-/data/hzy/.sudo_password}"
|
|
|
|
P_CONTAINER_PREFIX="${EXPERIMENT}_prefill"
|
|
D_CONTAINER_PREFIX="${EXPERIMENT}_decode"
|
|
MC_CONTAINER="${EXPERIMENT}_mc_master"
|
|
ROUTER_CONTAINER="${EXPERIMENT}_router"
|
|
|
|
mkdir -p "${RESULT_ROOT}"/{commands,logs,metadata,bench}
|
|
|
|
log() {
|
|
printf '[%s] %s\n' "$(date '+%F %T')" "$*" | tee -a "${RESULT_ROOT}/orchestrator.log"
|
|
}
|
|
|
|
require_password() {
|
|
if [[ -z "${SUDO_PASSWORD:-}" ]]; then
|
|
[[ -r "${SUDO_PASSWORD_FILE}" ]] || {
|
|
echo "ERROR: set SUDO_PASSWORD or create ${SUDO_PASSWORD_FILE}" >&2
|
|
exit 2
|
|
}
|
|
IFS= read -r SUDO_PASSWORD <"${SUDO_PASSWORD_FILE}"
|
|
fi
|
|
}
|
|
|
|
is_local() {
|
|
[[ "$1" == "${P_HEAD}" ]]
|
|
}
|
|
|
|
remote() {
|
|
local host="$1"
|
|
shift
|
|
if is_local "$host"; then
|
|
"$@"
|
|
return
|
|
fi
|
|
local command
|
|
printf -v command '%q ' "$@"
|
|
ssh "${SSH_OPTIONS[@]}" "${SSH_USER}@${host}" "${command}"
|
|
}
|
|
|
|
sudo_host() {
|
|
local host="$1"
|
|
shift
|
|
require_password
|
|
if is_local "$host"; then
|
|
printf '%s\n' "${SUDO_PASSWORD}" | sudo -S -p '' -- "$@"
|
|
return
|
|
fi
|
|
local command
|
|
printf -v command '%q ' "$@"
|
|
printf '%s\n' "${SUDO_PASSWORD}" | ssh "${SSH_OPTIONS[@]}" \
|
|
"${SSH_USER}@${host}" "sudo -S -p '' -- ${command}"
|
|
}
|
|
|
|
node_suffix() {
|
|
printf '%s' "${1##*.}"
|
|
}
|
|
|
|
container_name() {
|
|
local group="$1" rank="$2"
|
|
if [[ "$group" == "p" ]]; then
|
|
printf '%s_node%s' "${P_CONTAINER_PREFIX}" "$rank"
|
|
else
|
|
printf '%s_node%s' "${D_CONTAINER_PREFIX}" "$rank"
|
|
fi
|
|
}
|
|
|
|
build_image_on_host() {
|
|
local host="$1" suffix build_dir
|
|
suffix="$(node_suffix "$host")"
|
|
build_dir="/tmp/${EXPERIMENT}_build"
|
|
log "staging PD image build context on node ${suffix}"
|
|
|
|
if is_local "$host"; then
|
|
rm -rf "$build_dir"
|
|
mkdir -p "$build_dir/patches"
|
|
cp "${SCRIPT_DIR}/Dockerfile.pd" "$build_dir/Dockerfile"
|
|
cp "${SCRIPT_DIR}/patches/31869.patch" "$build_dir/patches/31869.patch"
|
|
cp "${SCRIPT_DIR}/patches/32797.patch" "$build_dir/patches/32797.patch"
|
|
cp "${MOONCAKE_WHEEL}" "$build_dir/mooncake.whl"
|
|
else
|
|
ssh "${SSH_OPTIONS[@]}" "${SSH_USER}@${host}" \
|
|
"rm -rf ${build_dir} && mkdir -p ${build_dir}/patches"
|
|
scp "${SSH_OPTIONS[@]}" "${SCRIPT_DIR}/Dockerfile.pd" \
|
|
"${SSH_USER}@${host}:${build_dir}/Dockerfile"
|
|
scp "${SSH_OPTIONS[@]}" "${SCRIPT_DIR}/patches/31869.patch" \
|
|
"${SSH_USER}@${host}:${build_dir}/patches/31869.patch"
|
|
scp "${SSH_OPTIONS[@]}" "${SCRIPT_DIR}/patches/32797.patch" \
|
|
"${SSH_USER}@${host}:${build_dir}/patches/32797.patch"
|
|
scp "${SSH_OPTIONS[@]}" "${MOONCAKE_WHEEL}" \
|
|
"${SSH_USER}@${host}:${build_dir}/mooncake.whl"
|
|
fi
|
|
|
|
sudo_host "$host" docker build \
|
|
--build-arg "BASE_IMAGE=${BASE_IMAGE}" \
|
|
-t "${PD_IMAGE}" "$build_dir" \
|
|
>"${RESULT_ROOT}/logs/build_node${suffix}.log" 2>&1
|
|
log "PD image built on node ${suffix}"
|
|
}
|
|
|
|
build_image_all() {
|
|
local host
|
|
local -a build_nodes
|
|
read -r -a build_nodes <<<"${BUILD_NODES_OVERRIDE:-${ALL_NODES[*]}}"
|
|
for host in "${build_nodes[@]}"; do
|
|
build_image_on_host "$host"
|
|
done
|
|
}
|
|
|
|
preflight() {
|
|
require_password
|
|
local host suffix
|
|
: >"${RESULT_ROOT}/metadata/preflight.tsv"
|
|
printf 'node\timage\tmodel\trdma\tgpu_processes\tports\n' \
|
|
>>"${RESULT_ROOT}/metadata/preflight.tsv"
|
|
for host in "${ALL_NODES[@]}"; do
|
|
suffix="$(node_suffix "$host")"
|
|
log "preflight node ${suffix}"
|
|
sudo_host "$host" docker image inspect "$PD_IMAGE" >/dev/null
|
|
remote "$host" test -d "$MODEL_PATH"
|
|
remote "$host" test -e /dev/infiniband/uverbs0
|
|
local gpu_pids ports
|
|
gpu_pids="$(remote "$host" nvidia-smi --query-compute-apps=pid --format=csv,noheader | xargs || true)"
|
|
ports="$(remote "$host" sh -c "ss -lnt | grep -E ':(20000|28800|30000|31000|50051)[[:space:]]' || true")"
|
|
if [[ -n "$gpu_pids" ]]; then
|
|
echo "ERROR: node ${suffix} still has GPU processes: ${gpu_pids}" >&2
|
|
return 1
|
|
fi
|
|
if [[ -n "$ports" ]]; then
|
|
echo "ERROR: node ${suffix} has occupied PD ports: ${ports}" >&2
|
|
return 1
|
|
fi
|
|
printf '%s\tok\tok\tok\t0\tfree\n' "$suffix" \
|
|
>>"${RESULT_ROOT}/metadata/preflight.tsv"
|
|
done
|
|
|
|
local hash_cmd
|
|
hash_cmd="sha256sum /sgl-workspace/sglang/python/sglang/srt/disaggregation/prefill.py /sgl-workspace/sglang/python/sglang/srt/disaggregation/decode.py /sgl-workspace/sglang/python/sglang/srt/disaggregation/utils.py /sgl-workspace/sglang/python/sglang/srt/managers/scheduler_pp_mixin.py /sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py; python3 -m pip list | grep -E 'sglang|flashinfer|mooncake-transfer-engine'"
|
|
sudo_host "$P_HEAD" docker run --rm --entrypoint bash "$PD_IMAGE" -lc "$hash_cmd" \
|
|
>"${RESULT_ROOT}/metadata/pd_image_versions.txt"
|
|
grep -q 'mooncake-transfer-engine-cuda13.*0.3.12.post1' \
|
|
"${RESULT_ROOT}/metadata/pd_image_versions.txt"
|
|
log "preflight passed on all 8 nodes"
|
|
}
|
|
|
|
common_docker_args() {
|
|
local host="$1"
|
|
local suffix
|
|
suffix="$(node_suffix "$host")"
|
|
COMMON_DOCKER_ARGS=(
|
|
--gpus all --network host --ipc=host --ulimit memlock=-1
|
|
--device /dev/infiniband --shm-size 32g --entrypoint bash
|
|
-v "${MODEL_PATH}:${MODEL_PATH}:ro"
|
|
-e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
|
|
-e "SGLANG_HOST_IP=174.1.60.${suffix}"
|
|
-e GLOO_SOCKET_IFNAME=bond0
|
|
-e NCCL_SOCKET_IFNAME=bond1
|
|
-e "NCCL_IB_HCA=${NCCL_IB_HCAS}"
|
|
-e NCCL_IB_GID_INDEX=3
|
|
-e NCCL_IB_TIMEOUT=22
|
|
-e NCCL_IB_RETRY_CNT=7
|
|
-e NCCL_CUMEM_ENABLE=1
|
|
-e "MOONCAKE_MASTER=${MC_MASTER}"
|
|
-e MOONCAKE_PROTOCOL=rdma
|
|
-e SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK=0
|
|
-e SGLANG_MOE_FUSED_GATE_RADIX=1
|
|
-e FLASHINFER_DISABLE_JIT=1
|
|
-e FLASHINFER_DISABLE_VERSION_CHECK=1
|
|
)
|
|
}
|
|
|
|
start_node() {
|
|
local group="$1" rank="$2" host name bootstrap
|
|
local -a nodes launch_args cmd
|
|
if [[ "$group" == "p" ]]; then
|
|
nodes=("${P_NODES[@]}")
|
|
host="${nodes[$rank]}"
|
|
name="$(container_name p "$rank")"
|
|
launch_args=(
|
|
--model-path "$MODEL_PATH" --served-model-name "$SERVED_MODEL_NAME"
|
|
--tp-size "$P_TP" --pp-size "$P_PP" --ep-size "$P_EP"
|
|
--nnodes 4 --node-rank "$rank" --dist-init-addr "${P_HEAD}:${P_DIST_PORT}"
|
|
--trust-remote-code --moe-runner-backend flashinfer_mxfp4
|
|
--chunked-prefill-size "$P_CHUNKED_PREFILL_SIZE"
|
|
--page-size "$P_PAGE_SIZE"
|
|
--mem-fraction-static "$P_MEM_FRACTION_STATIC"
|
|
--cuda-graph-max-bs-decode 16
|
|
--mamba-radix-cache-strategy extra_buffer_lazy
|
|
--disable-radix-cache --dist-timeout 3600
|
|
--mamba-full-memory-ratio "$P_MAMBA_FULL_MEMORY_RATIO"
|
|
--disaggregation-transfer-backend mooncake
|
|
--disaggregation-bootstrap-port "$P_BOOTSTRAP_PORT"
|
|
--disaggregation-ib-device "$DISAGG_IB_DEVICES"
|
|
--disaggregation-mode prefill
|
|
--host 0.0.0.0 --port "$P_PORT"
|
|
)
|
|
else
|
|
nodes=("${D_NODES[@]}")
|
|
host="${nodes[$rank]}"
|
|
name="$(container_name d "$rank")"
|
|
launch_args=(
|
|
--model-path "$MODEL_PATH" --served-model-name "$SERVED_MODEL_NAME"
|
|
--tp-size "$D_TP" --pp-size "$D_PP" --ep-size "$D_EP"
|
|
--nnodes 4 --node-rank "$rank" --dist-init-addr "${D_HEAD}:${D_DIST_PORT}"
|
|
--trust-remote-code --moe-runner-backend marlin
|
|
--mem-fraction-static "$D_MEM_FRACTION_STATIC"
|
|
--cuda-graph-max-bs-decode 16
|
|
--mamba-radix-cache-strategy extra_buffer_lazy
|
|
--disable-radix-cache --dist-timeout 3600
|
|
--mamba-full-memory-ratio "$D_MAMBA_FULL_MEMORY_RATIO"
|
|
--page-size "$D_PAGE_SIZE"
|
|
--disaggregation-transfer-backend mooncake
|
|
--disaggregation-bootstrap-port "$P_BOOTSTRAP_PORT"
|
|
--disaggregation-ib-device "$DISAGG_IB_DEVICES"
|
|
--disaggregation-mode decode
|
|
--host 0.0.0.0 --port "$D_PORT"
|
|
)
|
|
fi
|
|
|
|
printf -v bootstrap '%q ' python3 -m sglang.launch_server "${launch_args[@]}"
|
|
bootstrap="exec ${bootstrap}"
|
|
common_docker_args "$host"
|
|
cmd=(docker run -d --name "$name" "${COMMON_DOCKER_ARGS[@]}" "$PD_IMAGE" -lc "$bootstrap")
|
|
printf '%q ' "${cmd[@]}" >"${RESULT_ROOT}/commands/${group}_node${rank}.cmd.txt"
|
|
printf '\n' >>"${RESULT_ROOT}/commands/${group}_node${rank}.cmd.txt"
|
|
sudo_host "$host" docker rm -f "$name" >/dev/null 2>&1 || true
|
|
sudo_host "$host" "${cmd[@]}" >/dev/null
|
|
log "started ${group} rank ${rank} on node $(node_suffix "$host")"
|
|
}
|
|
|
|
wait_health() {
|
|
local label="$1" host="$2" port="$3" elapsed
|
|
for ((elapsed = 1; elapsed <= HEALTH_WAIT_S; elapsed++)); do
|
|
if curl --fail --silent --max-time 5 "http://${host}:${port}/health" >/dev/null 2>&1; then
|
|
log "${label} healthy after ${elapsed}s"
|
|
return 0
|
|
fi
|
|
if (( elapsed % 30 == 0 )); then
|
|
log "waiting for ${label}, elapsed=${elapsed}s"
|
|
collect_group_logs "${label}_starting"
|
|
if grep -Eiq 'Traceback|CUDA out of memory|NCCL.*(error|failed)|AssertionError|RuntimeError' \
|
|
"${RESULT_ROOT}/logs/${label}_starting_"*.log 2>/dev/null; then
|
|
log "ERROR: fatal pattern found while starting ${label}"
|
|
return 1
|
|
fi
|
|
fi
|
|
sleep 1
|
|
done
|
|
return 1
|
|
}
|
|
|
|
start_group() {
|
|
local group="$1" rank
|
|
for rank in 1 2 3; do
|
|
start_node "$group" "$rank"
|
|
done
|
|
sleep 5
|
|
start_node "$group" 0
|
|
if [[ "$group" == "p" ]]; then
|
|
wait_health p "$P_HEAD" "$P_PORT"
|
|
else
|
|
wait_health d "$D_HEAD" "$D_PORT"
|
|
fi
|
|
collect_group_logs "${group}_healthy"
|
|
}
|
|
|
|
start_mc_master() {
|
|
sudo_host "$P_HEAD" docker rm -f "$MC_CONTAINER" >/dev/null 2>&1 || true
|
|
sudo_host "$P_HEAD" docker run -d --name "$MC_CONTAINER" --network host --gpus all \
|
|
--entrypoint /usr/local/bin/mooncake_master "$PD_IMAGE" >/dev/null
|
|
for _ in $(seq 1 30); do
|
|
remote "$P_HEAD" sh -c "ss -lnt | grep -q ':50051 '" && {
|
|
log "Mooncake master listening on 50051"
|
|
return 0
|
|
}
|
|
sleep 1
|
|
done
|
|
return 1
|
|
}
|
|
|
|
start_router() {
|
|
local -a cmd=(
|
|
docker run -d --name "$ROUTER_CONTAINER" --network host
|
|
--entrypoint python3 "$PD_IMAGE" -m sglang_router.launch_router
|
|
--pd-disaggregation --mini-lb
|
|
--prefill "http://${P_HEAD}:${P_PORT}" "$P_BOOTSTRAP_PORT"
|
|
--decode "http://${D_HEAD}:${D_PORT}"
|
|
--host 0.0.0.0 --port "$ROUTER_PORT"
|
|
)
|
|
printf '%q ' "${cmd[@]}" >"${RESULT_ROOT}/commands/router.cmd.txt"
|
|
printf '\n' >>"${RESULT_ROOT}/commands/router.cmd.txt"
|
|
sudo_host "$ROUTER_HOST" docker rm -f "$ROUTER_CONTAINER" >/dev/null 2>&1 || true
|
|
sudo_host "$ROUTER_HOST" "${cmd[@]}" >/dev/null
|
|
wait_health router "$ROUTER_HOST" "$ROUTER_PORT"
|
|
}
|
|
|
|
collect_group_logs() {
|
|
local label="$1" group rank host name
|
|
for group in p d; do
|
|
for rank in 0 1 2 3; do
|
|
if [[ "$group" == "p" ]]; then
|
|
host="${P_NODES[$rank]}"
|
|
else
|
|
host="${D_NODES[$rank]}"
|
|
fi
|
|
name="$(container_name "$group" "$rank")"
|
|
sudo_host "$host" docker logs "$name" \
|
|
>"${RESULT_ROOT}/logs/${label}_${group}_node${rank}.log" 2>&1 || true
|
|
done
|
|
done
|
|
sudo_host "$P_HEAD" docker logs "$MC_CONTAINER" \
|
|
>"${RESULT_ROOT}/logs/${label}_mc_master.log" 2>&1 || true
|
|
sudo_host "$ROUTER_HOST" docker logs "$ROUTER_CONTAINER" \
|
|
>"${RESULT_ROOT}/logs/${label}_router.log" 2>&1 || true
|
|
}
|
|
|
|
status_all() {
|
|
local host
|
|
for host in "${ALL_NODES[@]}"; do
|
|
echo "===== node $(node_suffix "$host") ====="
|
|
sudo_host "$host" docker ps --format '{{.Names}}|{{.Image}}|{{.Status}}' \
|
|
--filter "name=${EXPERIMENT}"
|
|
remote "$host" nvidia-smi --query-gpu=index,memory.used,utilization.gpu \
|
|
--format=csv,noheader
|
|
done
|
|
for endpoint in \
|
|
"p=http://${P_HEAD}:${P_PORT}/health" \
|
|
"d=http://${D_HEAD}:${D_PORT}/health" \
|
|
"router=http://${ROUTER_HOST}:${ROUTER_PORT}/health"; do
|
|
local name="${endpoint%%=*}" url="${endpoint#*=}"
|
|
if curl --fail --silent --max-time 5 "$url" >/dev/null 2>&1; then
|
|
echo "${name}=healthy"
|
|
else
|
|
echo "${name}=down"
|
|
fi
|
|
done
|
|
}
|
|
|
|
stop_all() {
|
|
local rank host name
|
|
sudo_host "$ROUTER_HOST" docker rm -f "$ROUTER_CONTAINER" >/dev/null 2>&1 || true
|
|
for rank in 0 1 2 3; do
|
|
host="${D_NODES[$rank]}"
|
|
name="$(container_name d "$rank")"
|
|
sudo_host "$host" docker rm -f "$name" >/dev/null 2>&1 || true
|
|
done
|
|
for rank in 0 1 2 3; do
|
|
host="${P_NODES[$rank]}"
|
|
name="$(container_name p "$rank")"
|
|
sudo_host "$host" docker rm -f "$name" >/dev/null 2>&1 || true
|
|
done
|
|
sudo_host "$P_HEAD" docker rm -f "$MC_CONTAINER" >/dev/null 2>&1 || true
|
|
log "standard PD service stopped"
|
|
}
|
|
|
|
start_all() {
|
|
preflight
|
|
start_mc_master
|
|
start_group p
|
|
start_group d
|
|
start_router
|
|
collect_group_logs ready
|
|
status_all | tee "${RESULT_ROOT}/metadata/status_ready.txt"
|
|
log "standard PD service ready at http://${ROUTER_HOST}:${ROUTER_PORT}"
|
|
}
|
|
|
|
case "$ACTION" in
|
|
build-image) build_image_all ;;
|
|
preflight) preflight ;;
|
|
start) start_all ;;
|
|
stop) stop_all ;;
|
|
restart) stop_all; start_all ;;
|
|
status) status_all ;;
|
|
logs) collect_group_logs manual ;;
|
|
*)
|
|
echo "Usage: $0 {build-image|preflight|start|stop|restart|status|logs}" >&2
|
|
exit 2
|
|
;;
|
|
esac
|