多机
This commit is contained in:
qihuanye
255
scripts/launch_multinode_eval.sh
Executable file
255
scripts/launch_multinode_eval.sh
Executable file
@@ -0,0 +1,255 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# Launch 2-node LeWM evaluation from node-3.
|
||||
#
|
||||
# Defaults match the current cluster layout:
|
||||
# node-3: 10.16.200.9, node_rank=0
|
||||
# node-2: 10.16.200.8, node_rank=1
|
||||
# Each node runs two local torchrun processes for two visible GPUs.
|
||||
|
||||
REPO_ROOT="${REPO_ROOT:-/home/lewm/lewm}"
|
||||
REMOTE_HOST="${REMOTE_HOST:-lewm@10.16.200.8}"
|
||||
MASTER_ADDR="${MASTER_ADDR:-10.16.200.9}"
|
||||
MASTER_PORT="${MASTER_PORT:-29500}"
|
||||
|
||||
NNODES="${NNODES:-2}"
|
||||
NPROC_PER_NODE="${NPROC_PER_NODE:-2}"
|
||||
CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1}"
|
||||
STABLEWM_HOME="${STABLEWM_HOME:-/home/lewm/.stable-wm}"
|
||||
|
||||
CONFIG_NAME="${CONFIG_NAME:-pusht.yaml}"
|
||||
POLICY="${POLICY:-pusht/lewm}"
|
||||
OUTPUT_FILENAME="${OUTPUT_FILENAME:-pusht_multinode_results.txt}"
|
||||
EXTRA_ARGS="${EXTRA_ARGS:-}"
|
||||
DRY_RUN="${DRY_RUN:-0}"
|
||||
TAIL_LOGS="${TAIL_LOGS:-1}"
|
||||
PRELOAD_WAIT="${PRELOAD_WAIT:-0}"
|
||||
PRELOAD_SIGNAL_FILE="${PRELOAD_SIGNAL_FILE:-/tmp/lewm_preload_start}"
|
||||
PRELOAD_CLEAR_SIGNAL="${PRELOAD_CLEAR_SIGNAL:-1}"
|
||||
|
||||
LOG_DIR="${LOG_DIR:-${REPO_ROOT}/logs/multinode}"
|
||||
mkdir -p "${LOG_DIR}"
|
||||
RUN_ID="$(date +%Y%m%d_%H%M%S)"
|
||||
LOCAL_LOG="${LOG_DIR}/${RUN_ID}_node3_rank0.log"
|
||||
REMOTE_LOG="${LOG_DIR}/${RUN_ID}_node2_rank1.log"
|
||||
|
||||
SSH_OPTS=(
|
||||
-F /dev/null
|
||||
-o StrictHostKeyChecking=no
|
||||
-o ServerAliveInterval=30
|
||||
-o ServerAliveCountMax=20
|
||||
)
|
||||
|
||||
COMMON_ARGS=(
|
||||
"--config-name=${CONFIG_NAME}"
|
||||
"policy=${POLICY}"
|
||||
"multi_node.enabled=true"
|
||||
"output.filename=${OUTPUT_FILENAME}"
|
||||
)
|
||||
|
||||
if [[ "${PRELOAD_WAIT}" == "1" ]]; then
|
||||
COMMON_ARGS+=(
|
||||
"preload_wait.enabled=true"
|
||||
"preload_wait.file=${PRELOAD_SIGNAL_FILE}"
|
||||
)
|
||||
fi
|
||||
|
||||
if [[ -n "${EXTRA_ARGS}" ]]; then
|
||||
# shellcheck disable=SC2206
|
||||
COMMON_ARGS+=(${EXTRA_ARGS})
|
||||
fi
|
||||
|
||||
make_command() {
|
||||
local node_rank="$1"
|
||||
local repo_q cuda_q stablewm_q arg_q eval_args
|
||||
printf -v repo_q '%q' "${REPO_ROOT}"
|
||||
printf -v cuda_q '%q' "${CUDA_VISIBLE_DEVICES}"
|
||||
printf -v stablewm_q '%q' "${STABLEWM_HOME}"
|
||||
|
||||
eval_args=""
|
||||
for arg in "${COMMON_ARGS[@]}"; do
|
||||
printf -v arg_q '%q' "${arg}"
|
||||
eval_args+=" ${arg_q}"
|
||||
done
|
||||
|
||||
printf 'cd %s && source .venv/bin/activate && export CUDA_VISIBLE_DEVICES=%s && export STABLEWM_HOME=%s && torchrun --nnodes=%q --nproc_per_node=%q --node_rank=%q --master_addr=%q --master_port=%q eval.py%s' \
|
||||
"${repo_q}" \
|
||||
"${cuda_q}" \
|
||||
"${stablewm_q}" \
|
||||
"${NNODES}" \
|
||||
"${NPROC_PER_NODE}" \
|
||||
"${node_rank}" \
|
||||
"${MASTER_ADDR}" \
|
||||
"${MASTER_PORT}" \
|
||||
"${eval_args}"
|
||||
}
|
||||
|
||||
REMOTE_CMD="$(make_command 1)"
|
||||
LOCAL_CMD="$(make_command 0)"
|
||||
printf -v REMOTE_CMD_Q '%q' "${REMOTE_CMD}"
|
||||
|
||||
REMOTE_PID=""
|
||||
LOCAL_PID=""
|
||||
LOCAL_TAIL_PID=""
|
||||
REMOTE_TAIL_PID=""
|
||||
REMOTE_CLEANUP_CMD=""
|
||||
REMOTE_CLEANUP_CMD_Q=""
|
||||
|
||||
start_log_tail() {
|
||||
local label="$1"
|
||||
local log_file="$2"
|
||||
local label_q log_q
|
||||
|
||||
printf -v label_q '%q' "${label}"
|
||||
printf -v log_q '%q' "${log_file}"
|
||||
setsid bash -lc "tail -n +1 -F ${log_q} 2>/dev/null | sed -u 's/^/[${label_q}] /'" &
|
||||
}
|
||||
|
||||
stop_log_tails() {
|
||||
local pid
|
||||
for pid in "${LOCAL_TAIL_PID}" "${REMOTE_TAIL_PID}"; do
|
||||
if [[ -n "${pid}" ]] && kill -0 "${pid}" 2>/dev/null; then
|
||||
kill -TERM "-${pid}" 2>/dev/null || kill -TERM "${pid}" 2>/dev/null || true
|
||||
fi
|
||||
done
|
||||
}
|
||||
|
||||
remote_cleanup_command() {
|
||||
local pattern_q
|
||||
local patterns=(
|
||||
"torchrun .*--master_addr=${MASTER_ADDR} .*--master_port=${MASTER_PORT} .*eval.py"
|
||||
"torchrun .*--master_port=${MASTER_PORT} .*eval.py"
|
||||
"python.*eval.py .*output.filename=${OUTPUT_FILENAME}"
|
||||
)
|
||||
|
||||
printf 'set +e; '
|
||||
for pattern in "${patterns[@]}"; do
|
||||
printf -v pattern_q '%q' "${pattern}"
|
||||
printf 'pkill -TERM -f %s 2>/dev/null; ' "${pattern_q}"
|
||||
done
|
||||
printf 'sleep 2; '
|
||||
for pattern in "${patterns[@]}"; do
|
||||
printf -v pattern_q '%q' "${pattern}"
|
||||
printf 'pkill -KILL -f %s 2>/dev/null; ' "${pattern_q}"
|
||||
done
|
||||
printf 'true'
|
||||
}
|
||||
|
||||
cleanup() {
|
||||
local status="$?"
|
||||
trap - INT TERM EXIT
|
||||
|
||||
if [[ "${status}" -eq 0 ]]; then
|
||||
return 0
|
||||
fi
|
||||
|
||||
echo
|
||||
echo "Stopping multi-node eval..."
|
||||
stop_log_tails
|
||||
|
||||
if [[ -n "${LOCAL_PID}" ]] && kill -0 "${LOCAL_PID}" 2>/dev/null; then
|
||||
kill -TERM "-${LOCAL_PID}" 2>/dev/null || kill -TERM "${LOCAL_PID}" 2>/dev/null || true
|
||||
fi
|
||||
|
||||
if [[ -n "${REMOTE_PID}" ]] && kill -0 "${REMOTE_PID}" 2>/dev/null; then
|
||||
kill -TERM "${REMOTE_PID}" 2>/dev/null || true
|
||||
fi
|
||||
|
||||
ssh "${SSH_OPTS[@]}" "${REMOTE_HOST}" "bash -lc ${REMOTE_CLEANUP_CMD_Q}" >/dev/null 2>&1 || true
|
||||
|
||||
if [[ -n "${LOCAL_PID}" ]] && kill -0 "${LOCAL_PID}" 2>/dev/null; then
|
||||
sleep 2
|
||||
kill -KILL "-${LOCAL_PID}" 2>/dev/null || kill -KILL "${LOCAL_PID}" 2>/dev/null || true
|
||||
fi
|
||||
|
||||
echo "Cleanup requested. Check logs if any process was already exiting:"
|
||||
echo " local: ${LOCAL_LOG}"
|
||||
echo " remote: ${REMOTE_LOG}"
|
||||
exit "${status}"
|
||||
}
|
||||
|
||||
trap cleanup INT TERM EXIT
|
||||
|
||||
REMOTE_CLEANUP_CMD="$(remote_cleanup_command)"
|
||||
printf -v REMOTE_CLEANUP_CMD_Q '%q' "${REMOTE_CLEANUP_CMD}"
|
||||
|
||||
echo "Launching multi-node eval"
|
||||
echo " master: ${MASTER_ADDR}:${MASTER_PORT}"
|
||||
echo " remote: ${REMOTE_HOST}"
|
||||
echo " repo: ${REPO_ROOT}"
|
||||
echo " stablewm: ${STABLEWM_HOME}"
|
||||
echo " config: ${CONFIG_NAME}"
|
||||
echo " policy: ${POLICY}"
|
||||
echo " output: ${OUTPUT_FILENAME}"
|
||||
echo " extra: ${EXTRA_ARGS:-<none>}"
|
||||
echo " tail logs: ${TAIL_LOGS}"
|
||||
echo " preload wait: ${PRELOAD_WAIT}"
|
||||
if [[ "${PRELOAD_WAIT}" == "1" ]]; then
|
||||
echo " preload signal: ${PRELOAD_SIGNAL_FILE}"
|
||||
echo " start command: touch ${PRELOAD_SIGNAL_FILE}"
|
||||
fi
|
||||
echo " local log: ${LOCAL_LOG}"
|
||||
echo " remote log: ${REMOTE_LOG}"
|
||||
|
||||
if [[ "${DRY_RUN}" == "1" ]]; then
|
||||
echo
|
||||
echo "Remote command:"
|
||||
echo "ssh ${SSH_OPTS[*]} ${REMOTE_HOST} bash -lc ${REMOTE_CMD_Q}"
|
||||
echo
|
||||
echo "Local command:"
|
||||
printf -v LOCAL_CMD_Q '%q' "${LOCAL_CMD}"
|
||||
echo "bash -lc ${LOCAL_CMD_Q}"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
if [[ "${PRELOAD_WAIT}" == "1" && "${PRELOAD_CLEAR_SIGNAL}" == "1" ]]; then
|
||||
rm -f "${PRELOAD_SIGNAL_FILE}"
|
||||
fi
|
||||
|
||||
echo "Starting remote node_rank=1..."
|
||||
ssh "${SSH_OPTS[@]}" "${REMOTE_HOST}" "bash -lc ${REMOTE_CMD_Q}" >"${REMOTE_LOG}" 2>&1 &
|
||||
REMOTE_PID="$!"
|
||||
|
||||
if [[ "${TAIL_LOGS}" == "1" ]]; then
|
||||
start_log_tail "node2" "${REMOTE_LOG}"
|
||||
REMOTE_TAIL_PID="$!"
|
||||
fi
|
||||
|
||||
sleep 3
|
||||
|
||||
echo "Starting local node_rank=0..."
|
||||
set +e
|
||||
setsid bash -lc "${LOCAL_CMD}" >"${LOCAL_LOG}" 2>&1 &
|
||||
LOCAL_PID="$!"
|
||||
|
||||
if [[ "${TAIL_LOGS}" == "1" ]]; then
|
||||
start_log_tail "node3" "${LOCAL_LOG}"
|
||||
LOCAL_TAIL_PID="$!"
|
||||
fi
|
||||
|
||||
wait "${LOCAL_PID}"
|
||||
LOCAL_STATUS="$?"
|
||||
|
||||
wait "${REMOTE_PID}"
|
||||
REMOTE_STATUS="$?"
|
||||
set -e
|
||||
|
||||
stop_log_tails
|
||||
trap - INT TERM EXIT
|
||||
|
||||
echo "Local status: ${LOCAL_STATUS}"
|
||||
echo "Remote status: ${REMOTE_STATUS}"
|
||||
echo "Local log: ${LOCAL_LOG}"
|
||||
echo "Remote log: ${REMOTE_LOG}"
|
||||
|
||||
if [[ "${LOCAL_STATUS}" -ne 0 || "${REMOTE_STATUS}" -ne 0 ]]; then
|
||||
echo "Multi-node eval failed. Tail logs:"
|
||||
echo "===== local tail ====="
|
||||
tail -80 "${LOCAL_LOG}" || true
|
||||
echo "===== remote tail ====="
|
||||
tail -80 "${REMOTE_LOG}" || true
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Multi-node eval complete."
|
||||
@@ -44,6 +44,13 @@ mkdir -p "${OUTPUT_DIR}"
|
||||
# ROCR_VISIBLE_DEVICES=2,3 HIP_VISIBLE_DEVICES=0,1 MULTI_GPU=1 MULTI_GPU_DEVICES='[0,1]'
|
||||
MULTI_GPU="${MULTI_GPU:-0}"
|
||||
MULTI_GPU_DEVICES="${MULTI_GPU_DEVICES:-[0,1]}"
|
||||
MULTI_NODE="${MULTI_NODE:-0}"
|
||||
|
||||
# Multi-node warmup uses the same eval.py entrypoint under torchrun.
|
||||
# Example:
|
||||
# torchrun --nnodes=2 --nproc_per_node=2 --node_rank=0 --master_addr=<ip> --master_port=29500 \
|
||||
# eval.py --config-name=pusht.yaml policy=pusht/lewm multi_node.enabled=true
|
||||
# This script leaves multi-node launch to the caller.
|
||||
|
||||
COMMON_ARGS=(
|
||||
"eval.num_eval=${WARMUP_NUM_EVAL}"
|
||||
@@ -57,6 +64,12 @@ if [[ "${MULTI_GPU}" == "1" ]]; then
|
||||
)
|
||||
fi
|
||||
|
||||
if [[ "${MULTI_NODE}" == "1" ]]; then
|
||||
COMMON_ARGS+=(
|
||||
"multi_node.enabled=true"
|
||||
)
|
||||
fi
|
||||
|
||||
run_warmup() {
|
||||
local config_name="$1"
|
||||
local policy="$2"
|
||||
@@ -80,10 +93,11 @@ echo " HIP_VISIBLE_DEVICES: ${HIP_VISIBLE_DEVICES}"
|
||||
echo " CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES}"
|
||||
echo " WARMUP_NUM_EVAL: ${WARMUP_NUM_EVAL}"
|
||||
echo " INFERENCE_PRECISION: ${INFERENCE_PRECISION}"
|
||||
echo " MULTI_GPU: ${MULTI_GPU}"
|
||||
if [[ "${MULTI_GPU}" == "1" ]]; then
|
||||
echo " MULTI_GPU_DEVICES: ${MULTI_GPU_DEVICES}"
|
||||
fi
|
||||
echo " MULTI_GPU: ${MULTI_GPU}"
|
||||
if [[ "${MULTI_GPU}" == "1" ]]; then
|
||||
echo " MULTI_GPU_DEVICES: ${MULTI_GPU_DEVICES}"
|
||||
fi
|
||||
echo " MULTI_NODE: ${MULTI_NODE}"
|
||||
|
||||
# Defaults match the checkpoint names used in this repo. If onsite checkpoint
|
||||
# folders differ, override by editing these calls or passing the equivalent
|
||||
|
||||
Reference in New Issue
Block a user