| 1 | #!/usr/bin/env bash |
| 2 | # Run WM cases across the GPUs you have. |
| 3 | # |
| 4 | # Cases are split round-robin across GPU_LIST, and each GPU works through its own |
| 5 | # share one at a time. The GPU count does not have to match the case count: with |
| 6 | # 3 cases on 2 GPUs, GPU 0 runs cases 1 and 3 in sequence while GPU 1 runs case 2. |
| 7 | # Only one inference process per GPU is ever live, which also keeps host memory use |
| 8 | # bounded (loading a checkpoint is memory-hungry). |
| 9 | # |
| 10 | # Environment: |
| 11 | # GPU_LIST comma-separated GPU indices (default: 0,1,2) |
| 12 | # CASES space/comma-separated case names (default: every case dir found) |
| 13 | # PYTHON_BIN interpreter to use (default: python3) |
| 14 | # ACTION_OVERLAY write the HUD copies (default: 1; set to 0 to skip them) |
| 15 | # |
| 16 | # Positional: [checkpoint] [gemma_path] [output_root] |
| 17 | set -euo pipefail |
| 18 | |
| 19 | wm_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" |
| 20 | checkpoint="${1:-$wm_root/checkpoints/echo-wm-base.safetensors}" |
| 21 | gemma_path="${2:-$wm_root/checkpoints/gemma-3}" |
| 22 | output_root="${3:-$wm_root/outputs/wm_cases_multigpu}" |
| 23 | gpu_list="${GPU_LIST:-0,1,2}" |
| 24 | python_bin="${PYTHON_BIN:-python3}" |
| 25 | case "${ACTION_OVERLAY-1}" in |
| 26 | 0|false|no|off|"") overlay_flag="--no-action-overlay" ;; |
| 27 | *) overlay_flag="--action-overlay" ;; |
| 28 | esac |
| 29 | |
| 30 | IFS=',' read -r -a gpus <<< "$gpu_list" |
| 31 | |
| 32 | # Default to every checked-in case; override with CASES="0004 0010". |
| 33 | if [[ -n "${CASES:-}" ]]; then |
| 34 | IFS=', ' read -r -a cases <<< "$CASES" |
| 35 | else |
| 36 | cases=() |
| 37 | for case_dir in "$wm_root"/examples/wm_cases/*/; do |
| 38 | [[ -f "$case_dir/case.json" ]] && cases+=("$(basename "$case_dir")") |
| 39 | done |
| 40 | fi |
| 41 | if (( ${#cases[@]} == 0 )); then |
| 42 | echo "No cases found under $wm_root/examples/wm_cases" >&2 |
| 43 | exit 2 |
| 44 | fi |
| 45 | |
| 46 | echo "Cases: ${cases[*]}" |
| 47 | echo "GPUs: ${gpus[*]}" |
| 48 | echo "Output: $output_root" |
| 49 | mkdir -p "$output_root" |
| 50 | |
| 51 | # Walk this GPU's share of the case list, one case at a time. |
| 52 | run_share() { |
| 53 | local gpu="$1" start="$2" stride="$3" rc=0 |
| 54 | for (( i = start; i < ${#cases[@]}; i += stride )); do |
| 55 | local case_name="${cases[$i]}" |
| 56 | local case_output="$output_root/$case_name" |
| 57 | mkdir -p "$case_output" |
| 58 | echo "[$case_name] -> GPU $gpu" |
| 59 | if ( |
| 60 | cd "$wm_root" |
| 61 | CUDA_VISIBLE_DEVICES="$gpu" "$python_bin" scripts/run_wm_case.py \ |
| 62 | --case "examples/wm_cases/$case_name" \ |
| 63 | --checkpoint "$checkpoint" \ |
| 64 | --gemma-path "$gemma_path" \ |
| 65 | $overlay_flag \ |
| 66 | --output-dir "$output_root" |
| 67 | ) >"$case_output/run_gpu${gpu}.log" 2>&1; then |
| 68 | echo "[$case_name] completed (GPU $gpu)" |
| 69 | else |
| 70 | echo "[$case_name] failed; see $case_output/run_gpu${gpu}.log" >&2 |
| 71 | rc=1 |
| 72 | fi |
| 73 | done |
| 74 | return "$rc" |
| 75 | } |
| 76 | |
| 77 | pids=() |
| 78 | for idx in "${!gpus[@]}"; do |
| 79 | run_share "${gpus[$idx]}" "$idx" "${#gpus[@]}" & |
| 80 | pids+=("$!") |
| 81 | done |
| 82 | |
| 83 | status=0 |
| 84 | for pid in "${pids[@]}"; do |
| 85 | wait "$pid" || status=1 |
| 86 | done |
| 87 | |
| 88 | exit "$status" |
| 89 |