| #!/usr/bin/env bash |
| |
| set -euo pipefail |
|
|
| SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) |
| PROJECT_DIR=$(cd -- "${SCRIPT_DIR}/../.." && pwd) |
| RUNNER="${SCRIPT_DIR}/revsi/run_eval_vllm.sh" |
| MERGER="${SCRIPT_DIR}/revsi/merge_multinode_shards.py" |
| PYTHON_BIN=${PYTHON_BIN:-$(command -v python)} |
| HOSTFILE=${HOSTFILE:-} |
| NUM_NODES=${NUM_NODES:-4} |
| GPUS_PER_NODE=${GPUS_PER_NODE:-8} |
|
|
| MODEL=${MODEL:-} |
| MODEL=${MODEL%/} |
| REVSI_ROOT=${REVSI_ROOT:-} |
| FRAME_BUDGET=${REVSI_FRAME_BUDGET:-all} |
| FRAME_DIR="${REVSI_ROOT:+${REVSI_ROOT}/${FRAME_BUDGET}_frame}" |
| QA_FILE=${REVSI_QA_FILE:-"${FRAME_DIR:+${FRAME_DIR}/test-00000-of-00001.parquet}"} |
| VIDEO_ROOT=${REVSI_VIDEO_ROOT:-${REVSI_ROOT}} |
| TASK_FILTER=${REVSI_TASK_FILTER:-} |
| MAX_SAMPLES=${REVSI_MAX_SAMPLES:-0} |
| REMOTE_SETUP=${REMOTE_SETUP:-} |
| SSH_STRICT_HOST_KEY_CHECKING=${SSH_STRICT_HOST_KEY_CHECKING:-yes} |
|
|
| case "${FRAME_BUDGET}" in |
| 16|32|64|all) ;; |
| *) |
| echo "[FATAL] REVSI_FRAME_BUDGET must be 16, 32, 64, or all." >&2 |
| exit 2 |
| ;; |
| esac |
| [[ -n "${MODEL}" ]] || { |
| echo "[FATAL] Set MODEL to the model directory." >&2 |
| exit 2 |
| } |
| [[ -n "${QA_FILE}" ]] || { |
| echo "[FATAL] Set REVSI_QA_FILE or REVSI_ROOT." >&2 |
| exit 2 |
| } |
| [[ -n "${HOSTFILE}" ]] || { |
| echo "[FATAL] Set HOSTFILE to one host per line." >&2 |
| exit 2 |
| } |
| for path in \ |
| "${RUNNER}" "${MERGER}" "${HOSTFILE}" "${MODEL}/config.json" "${QA_FILE}"; do |
| [[ -f "${path}" ]] || { |
| echo "[FATAL] Required input is unavailable: ${path}" >&2 |
| exit 1 |
| } |
| done |
| if [[ "${QA_FILE}" == *.parquet ]]; then |
| [[ -n "${FRAME_DIR}" ]] || { |
| echo "[FATAL] Parquet input requires REVSI_ROOT." >&2 |
| exit 2 |
| } |
| compgen -G "${FRAME_DIR}/*.mp4" >/dev/null || { |
| echo "[FATAL] No ReVSI videos found under ${FRAME_DIR}." >&2 |
| exit 1 |
| } |
| fi |
|
|
| if [[ -n "${REVSI_EXPECTED_SAMPLES:-}" ]]; then |
| EXPECTED_SAMPLES=${REVSI_EXPECTED_SAMPLES} |
| elif [[ -z "${TASK_FILTER}" && "${MAX_SAMPLES}" == "0" ]]; then |
| EXPECTED_SAMPLES=6808 |
| else |
| EXPECTED_SAMPLES=0 |
| fi |
|
|
| mapfile -t AVAILABLE_HOSTS < <(awk 'NF && !seen[$1]++ {print $1}' "${HOSTFILE}") |
| [[ "${NUM_NODES}" =~ ^[1-9][0-9]*$ ]] || { |
| echo "[FATAL] NUM_NODES must be a positive integer." >&2 |
| exit 2 |
| } |
| ((NUM_NODES <= ${#AVAILABLE_HOSTS[@]})) || { |
| echo "[FATAL] Requested ${NUM_NODES} nodes, only ${#AVAILABLE_HOSTS[@]} available." >&2 |
| exit 2 |
| } |
| HOSTS=("${AVAILABLE_HOSTS[@]:0:NUM_NODES}") |
| GLOBAL_SHARDS=$((NUM_NODES * GPUS_PER_NODE)) |
|
|
| checkpoint_dir=$(dirname "$(dirname "${MODEL%/}")") |
| checkpoint_tag=$(basename "${checkpoint_dir}") |
| experiment_tag=$(basename "$(dirname "${checkpoint_dir}")") |
| timestamp=$(date +%Y%m%d_%H%M%S) |
| EVAL_MAX_FRAMES=${REVSI_MAX_FRAMES:-${FRAME_BUDGET}} |
| EVAL_EXACT_NFRAMES=${REVSI_EXACT_NFRAMES:-1} |
| if [[ "${FRAME_BUDGET}" == "all" ]]; then |
| EVAL_MAX_FRAMES=${REVSI_MAX_FRAMES:-128} |
| fi |
| setting_tag="revsi-${FRAME_BUDGET}frame-multinode${NUM_NODES}x${GPUS_PER_NODE}-f${EVAL_MAX_FRAMES}-exact${EVAL_EXACT_NFRAMES}-fps${REVSI_FPS:-2}-total${REVSI_VIDEO_TOTAL_PIXELS:-16777216}" |
| OUTPUT_DIR=${OUTPUT_DIR:-"${PROJECT_DIR}/outputs/${experiment_tag}/spatial_intelligence/revsi/${checkpoint_tag}/${setting_tag}/${timestamp}"} |
| LOG_DIR=${LOG_DIR:-"${PROJECT_DIR}/logs/revsi_multinode/${timestamp}"} |
| mkdir -p "${OUTPUT_DIR}" "${LOG_DIR}" |
|
|
| SSH_OPTS=(-o "StrictHostKeyChecking=${SSH_STRICT_HOST_KEY_CHECKING}" -o BatchMode=yes -o ServerAliveInterval=60) |
| WORKER_PATTERN='eval/task/revsi/eval_revsi_vllm.py' |
|
|
| run_on_host() { |
| local node=$1 |
| shift |
| if ((node == 0)); then |
| bash -s <<<"$*" |
| else |
| ssh "${SSH_OPTS[@]}" "${HOSTS[node]}" bash -s <<<"$*" |
| fi |
| } |
|
|
| kill_stale_workers() { |
| local node reap_pids=() command |
| command="pkill -f $(printf %q "${WORKER_PATTERN}") >/dev/null 2>&1 || true; sleep 5; pkill -9 -f $(printf %q "${WORKER_PATTERN}") >/dev/null 2>&1 || true; sleep 3" |
| for ((node = 0; node < NUM_NODES; node++)); do |
| run_on_host "${node}" "${command}" >/dev/null 2>&1 & |
| reap_pids+=("$!") |
| done |
| wait "${reap_pids[@]}" 2>/dev/null || true |
| } |
|
|
| require_free_gpus() { |
| local node free bad=0 |
| local need_gib=${REVSI_MIN_FREE_GIB:-75} |
| for ((node = 0; node < NUM_NODES; node++)); do |
| free=$(run_on_host "${node}" \ |
| "nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits | sort -n | awk 'NR==1 {print}'" 2>/dev/null | |
| tr -d '\r' | awk '/^[0-9]+$/ {value=$0} END {print value}') |
| if [[ -z "${free}" ]]; then |
| echo "[FATAL] node=${node} host=${HOSTS[node]}: cannot query GPU memory." >&2 |
| bad=1 |
| elif ((free / 1024 < need_gib)); then |
| echo "[FATAL] node=${node} host=${HOSTS[node]}: only $((free / 1024)) GiB free (need ${need_gib} GiB)." >&2 |
| bad=1 |
| fi |
| done |
| ((bad == 0)) || exit 1 |
| } |
|
|
| pids=() |
| cleanup() { |
| local status=$? |
| trap - EXIT INT TERM |
| for pid in "${pids[@]:-}"; do |
| [[ -n "${pid}" ]] && kill "${pid}" >/dev/null 2>&1 || true |
| done |
| ((status == 0)) || kill_stale_workers |
| exit "${status}" |
| } |
| trap cleanup EXIT INT TERM |
|
|
| kill_stale_workers |
| require_free_gpus |
|
|
| dispatch() { |
| local node=$1 |
| local offset=$((node * GPUS_PER_NODE)) |
| local gpu_ids |
| gpu_ids=$(seq -s, 0 $((GPUS_PER_NODE - 1))) |
| local env_values=( |
| "CUDA_VISIBLE_DEVICES=${gpu_ids}" |
| "REVSI_ROOT=${REVSI_ROOT}" |
| "FRAME_BUDGET=${FRAME_BUDGET}" |
| "QA_FILE=${QA_FILE}" |
| "VIDEO_ROOT=${VIDEO_ROOT}" |
| "OUTPUT_DIR=${OUTPUT_DIR}" |
| "GLOBAL_SHARD_COUNT=${GLOBAL_SHARDS}" |
| "GLOBAL_SHARD_OFFSET=${offset}" |
| "MERGE_SHARDS=0" |
| "EXPECTED_SAMPLES=0" |
| "TASK_FILTER=${TASK_FILTER}" |
| "MAX_SAMPLES=${MAX_SAMPLES}" |
| "MEDIA_MODE=video" |
| "TP_SIZE=1" |
| "BATCH_SIZE=${REVSI_BATCH_SIZE:-16}" |
| "MAX_FRAMES=${EVAL_MAX_FRAMES}" |
| "EXACT_NFRAMES=${EVAL_EXACT_NFRAMES}" |
| "FPS=${REVSI_FPS:-2}" |
| "VIDEO_MIN_PIXELS=${REVSI_VIDEO_MIN_PIXELS:-65536}" |
| "VIDEO_MAX_PIXELS=${REVSI_VIDEO_MAX_PIXELS:-}" |
| "VIDEO_TOTAL_PIXELS=${REVSI_VIDEO_TOTAL_PIXELS:-16777216}" |
| "MAX_MODEL_LEN=${REVSI_MAX_MODEL_LEN:-32768}" |
| "MAX_NEW_TOKENS=${REVSI_MAX_NEW_TOKENS:-64}" |
| "GPU_MEM_UTIL=${REVSI_GPU_MEMORY_UTILIZATION:-0.90}" |
| "STRICT_NUMERIC_PROMPT=${REVSI_STRICT_NUMERIC_PROMPT:-0}" |
| "ENABLE_THINKING=${REVSI_ENABLE_THINKING:-false}" |
| "SCORE_LOG_INTERVAL=${REVSI_SCORE_LOG_INTERVAL:-50}" |
| "LAUNCH_DELAY=2" |
| "VLLM_BASE_PORT=49000" |
| ) |
| local prefix="" value |
| for value in "${env_values[@]}"; do |
| prefix+="$(printf %q "${value}") " |
| done |
| local setup_prefix="" |
| if [[ -n "${REMOTE_SETUP}" ]]; then |
| setup_prefix="${REMOTE_SETUP} && " |
| fi |
| local command="${setup_prefix}cd $(printf %q "${PROJECT_DIR}") && ${prefix}bash $(printf %q "${RUNNER}") $(printf %q "${MODEL}")" |
| local log="${LOG_DIR}/node${node}.log" |
| echo "[revsi] node=${node}/${NUM_NODES} host=${HOSTS[node]} shards=${offset}-$((offset + GPUS_PER_NODE - 1)) log=${log}" |
| if ((node == 0)); then |
| bash -lc "${command}" 2>&1 | tee "${log}" |
| else |
| ssh "${SSH_OPTS[@]}" "${HOSTS[node]}" \ |
| "bash -lc $(printf %q "${command}")" >"${log}" 2>&1 & |
| pids+=("$!") |
| fi |
| } |
|
|
| for ((node = 1; node < NUM_NODES; node++)); do |
| dispatch "${node}" |
| done |
| dispatch 0 |
|
|
| status=0 |
| for pid in "${pids[@]}"; do |
| if ! wait "${pid}"; then status=1; fi |
| done |
| pids=() |
| ((status == 0)) || { |
| echo "[FATAL] One or more ReVSI nodes failed; inspect ${LOG_DIR}." >&2 |
| exit "${status}" |
| } |
|
|
| "${PYTHON_BIN}" "${MERGER}" \ |
| --output-dir "${OUTPUT_DIR}" \ |
| --num-shards "${GLOBAL_SHARDS}" \ |
| --expected-samples "${EXPECTED_SAMPLES}" |
| chmod -R a+rX "${OUTPUT_DIR}" 2>/dev/null || true |
| echo "[revsi] Multi-node evaluation complete: ${OUTPUT_DIR}/summary.json" |
|
|
