| 1 | #!/usr/bin/env bash |
| 2 | # Run checked-in causal WM cases across the available GPUs. |
| 3 | # |
| 4 | # Cases are assigned round-robin. Each GPU processes its assigned cases |
| 5 | # sequentially, so at most one inference process is active per GPU. |
| 6 | # |
| 7 | # Environment: |
| 8 | # GPU_LIST comma-separated GPU indices (default: 0,1,2) |
| 9 | # CASES space/comma-separated case names (default: every case dir) |
| 10 | # PYTHON_BIN interpreter to use (default: python from the active environment) |
| 11 | # ACTION_OVERLAY write the HUD copies (default: 1; set to 0 to skip them) |
| 12 | # |
| 13 | # Positional: [checkpoint] [gemma_path] [output_root] |
| 14 | set -euo pipefail |
| 15 | |
| 16 | wm_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" |
| 17 | checkpoint="${1:-$wm_root/checkpoints/echo-wm-flash.safetensors}" |
| 18 | gemma_path="${2:-$wm_root/checkpoints/gemma-3}" |
| 19 | output_root="${3:-$wm_root/outputs/wm_causal_cases_multigpu}" |
| 20 | gpu_list="${GPU_LIST:-0,1,2}" |
| 21 | python_bin="${PYTHON_BIN:-python}" |
| 22 | case "${ACTION_OVERLAY-1}" in |
| 23 | 0|false|no|off|"") overlay_flag="--no-action-overlay" ;; |
| 24 | *) overlay_flag="--action-overlay" ;; |
| 25 | esac |
| 26 | |
| 27 | IFS=',' read -r -a gpus <<< "$gpu_list" |
| 28 | |
| 29 | if [[ -n "${CASES:-}" ]]; then |
| 30 | IFS=', ' read -r -a cases <<< "$CASES" |
| 31 | else |
| 32 | cases=() |
| 33 | for case_dir in "$wm_root"/examples/wm_causal_cases/*/; do |
| 34 | [[ -f "$case_dir/case.json" ]] && cases+=("$(basename "$case_dir")") |
| 35 | done |
| 36 | fi |
| 37 | if (( ${#cases[@]} == 0 )); then |
| 38 | echo "No cases found under $wm_root/examples/wm_causal_cases" >&2 |
| 39 | exit 2 |
| 40 | fi |
| 41 | |
| 42 | echo "Cases: ${cases[*]}" |
| 43 | echo "GPUs: ${gpus[*]}" |
| 44 | echo "Output: $output_root" |
| 45 | mkdir -p "$output_root" |
| 46 | |
| 47 | run_share() { |
| 48 | local gpu="$1" start="$2" stride="$3" rc=0 |
| 49 | for (( i = start; i < ${#cases[@]}; i += stride )); do |
| 50 | local case_name="${cases[$i]}" |
| 51 | local case_output="$output_root/$case_name" |
| 52 | mkdir -p "$case_output" |
| 53 | echo "[$case_name] -> GPU $gpu" |
| 54 | if ( |
| 55 | cd "$wm_root" |
| 56 | CUDA_VISIBLE_DEVICES="$gpu" "$python_bin" scripts/run_wm_case_causal.py \ |
| 57 | --case "examples/wm_causal_cases/$case_name" \ |
| 58 | --checkpoint "$checkpoint" \ |
| 59 | --gemma-path "$gemma_path" \ |
| 60 | $overlay_flag \ |
| 61 | --output-dir "$output_root" |
| 62 | ) >"$case_output/run_gpu${gpu}.log" 2>&1; then |
| 63 | echo "[$case_name] completed (GPU $gpu)" |
| 64 | else |
| 65 | echo "[$case_name] failed; see $case_output/run_gpu${gpu}.log" >&2 |
| 66 | rc=1 |
| 67 | fi |
| 68 | done |
| 69 | return "$rc" |
| 70 | } |
| 71 | |
| 72 | pids=() |
| 73 | for idx in "${!gpus[@]}"; do |
| 74 | run_share "${gpus[$idx]}" "$idx" "${#gpus[@]}" & |
| 75 | pids+=("$!") |
| 76 | done |
| 77 | |
| 78 | status=0 |
| 79 | for pid in "${pids[@]}"; do |
| 80 | wait "$pid" || status=1 |
| 81 | done |
| 82 | |
| 83 | exit "$status" |
| 84 |