wait_gpus.sh

#!/bin/bash
# wait_gpus.sh — Wait until all expected GPUs are fully initialized.
# Usage: ./wait_gpus.sh [expected_count] [timeout_seconds]
#   expected_count: number of GPUs to wait for (default: 4)
#   timeout_seconds: max wait time in seconds (default: 120)
# Exit 0 when all GPUs ready, 1 on timeout.

set -e

EXPECTED_GPUS="${1:-4}"
TIMEOUT="${2:-120}"
INTERVAL=3

log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
}

log "⏳ Waiting for ${EXPECTED_GPUS} GPU(s) to initialize (timeout: ${TIMEOUT}s)..."

ELAPSED=0
while (( ELAPSED < TIMEOUT )); do
    COUNT=$(nvidia-smi -L 2>/dev/null | grep -c "^GPU" || true)
    if (( COUNT >= EXPECTED_GPUS )); then
        log "✅ All ${EXPECTED_GPUS} GPU(s) ready (${COUNT} detected after ${ELAPSED}s)"
        nvidia-smi -L
        exit 0
    fi
    log "  ⏱ ${COUNT}/${EXPECTED_GPUS} GPUs detected... waiting (${ELAPSED}s / ${TIMEOUT}s)"
    sleep "$INTERVAL"
    ELAPSED=$(( ELAPSED + INTERVAL ))
done

log "❌ Timeout: only ${COUNT}/${EXPECTED_GPUS} GPU(s) available after ${TIMEOUT}s"
nvidia-smi -L 2>/dev/null || true
exit 1