esh-pfi-infrastructure/scripts/upgrade-irv-ml1-cuda.sh

#!/usr/bin/env bash
# upgrade-irv-ml1-cuda.sh — bump irv-ml1's NVIDIA driver to support
# CUDA 12.9+ (currently capped at 12.8 per driver 570.124.06).
#
# Why two phases: the upgrade involves a reboot, which kills the SSH
# session mid-flight. Splitting at the reboot lets phase 2 wait for
# the host to come back and verify the new driver took.
#
# Usage:
#   scripts/upgrade-irv-ml1-cuda.sh phase1
#       Stop GPU containers, add NVIDIA CUDA APT repo, install
#       cuda-drivers (currently 580 series), reboot.
#       (interactive — sudo password prompted once)
#
#   ... wait ~2 min for irv-ml1 to come back ...
#
#   scripts/upgrade-irv-ml1-cuda.sh phase2
#       Verify new driver + CUDA version, restart containers,
#       and print the next step (flip Kokoro to GPU variant).
#
# Idempotent: rerunning phase1 after the repo is already added is a no-op
# until the apt step. phase2 can be re-run safely; container restarts
# via `docker compose up -d` are idempotent.

set -euo pipefail

HOST=irv-ml1
CONTAINERS=(comfyui cosyvoice qwen3-tts index-tts parakeet kokoro)

say()  { printf '\033[1;36m▸ %s\033[0m\n' "$*"; }
ok()   { printf '\033[1;32m✓ %s\033[0m\n' "$*"; }
warn() { printf '\033[1;33m⚠ %s\033[0m\n' "$*"; }
err()  { printf '\033[1;31m✗ %s\033[0m\n' "$*" >&2; }

phase1() {
    say "Phase 1: pre-flight + install + reboot"
    say "Connecting to $HOST (sudo will prompt once)"

    # -t requests a TTY so sudo can prompt. The whole heredoc runs as
    # a single remote shell — sudo's credential cache covers all
    # subsequent invocations within it.
    ssh -t "$HOST" "bash -s" <<EOF
set -euo pipefail

echo "=== current driver + CUDA ==="
nvidia-smi --query-gpu=driver_version --format=csv,noheader
nvidia-smi | grep "CUDA Version"

echo
echo "=== prime sudo (we only want one password prompt this whole script) ==="
sudo -v

echo
echo "=== adding NVIDIA CUDA APT repo for Debian 12 ==="
sudo install -d -m 0755 /etc/apt/keyrings
if [ ! -f /etc/apt/keyrings/nvidia-cuda.gpg ]; then
    sudo apt-get install -y curl gpg
    curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/debian12/x86_64/3bf863cc.pub \\
        | sudo gpg --dearmor --yes -o /etc/apt/keyrings/nvidia-cuda.gpg
fi
echo "deb [signed-by=/etc/apt/keyrings/nvidia-cuda.gpg] https://developer.download.nvidia.com/compute/cuda/repos/debian12/x86_64/ /" \\
    | sudo tee /etc/apt/sources.list.d/nvidia-cuda.list >/dev/null

sudo apt-get update -qq

echo
echo "=== cuda-drivers candidate ==="
apt-cache policy cuda-drivers | head -10

echo
echo "=== snapshotting current nvidia packages (rollback reference) ==="
dpkg -l | grep -E "^ii\s+(nvidia-|cuda-)" | awk '{print \$2 "=" \$3}' \\
    | sudo tee /tmp/nvidia-pre-upgrade.txt >/dev/null
echo "snapshot at /tmp/nvidia-pre-upgrade.txt"

echo
echo "=== stopping GPU containers ==="
for stack in ${CONTAINERS[*]}; do
    if [ -d "/opt/docker/compose/\$stack" ]; then
        echo "  stopping \$stack"
        (cd "/opt/docker/compose/\$stack" && docker compose stop) || echo "    (no-op)"
    fi
done

echo
echo "=== installing cuda-drivers (DKMS rebuild ~3-5 min) ==="
sudo DEBIAN_FRONTEND=noninteractive apt-get install -y cuda-drivers

echo
echo "=== rebooting in 5 s ==="
sleep 5
sudo systemctl reboot
EOF

    ok "Phase 1 dispatched. Host rebooting."
    say "Wait ~2 min, then run: $0 phase2"
}

phase2() {
    say "Phase 2: wait for host, verify, restart containers"

    say "waiting for $HOST to come back online…"
    for i in $(seq 1 60); do
        if ssh -o ConnectTimeout=3 -o BatchMode=yes "$HOST" true 2>/dev/null; then
            ok "$HOST reachable"
            break
        fi
        printf '.'
        sleep 3
    done
    echo

    ssh "$HOST" 'bash -s' <<'EOF'
set -euo pipefail

echo "=== new driver + CUDA ==="
nvidia-smi --query-gpu=driver_version --format=csv,noheader
nvidia-smi | grep "CUDA Version"

echo
echo "=== restarting GPU containers ==="
for stack in comfyui cosyvoice qwen3-tts index-tts parakeet; do
    if [ -d "/opt/docker/compose/$stack" ]; then
        echo "  starting $stack"
        (cd "/opt/docker/compose/$stack" && docker compose up -d) || echo "    (failed — check manually)"
    fi
done

sleep 8
echo
echo "=== container status ==="
docker ps --format "table {{.Names}}\t{{.Status}}" | head -15
EOF

    ok "Phase 2 complete."
    cat <<'EOF'

────────────────────────────────────────────────────────────────────
NEXT — flip Kokoro from cpu → gpu variant now that CUDA 12.9 is live:

  ssh irv-ml1 '
    cd /opt/docker/compose/kokoro
    sed -i "s/^KOKORO_VARIANT=.*/KOKORO_VARIANT=gpu/" .env
    sed -i "s/^KOKORO_USE_GPU=.*/KOKORO_USE_GPU=true/" .env
    sed -i "s/^KOKORO_GPU_DEVICES=.*/KOKORO_GPU_DEVICES=0/" .env
  '
  scripts/elway irv-ml1 --playbook playbooks/deploy-kokoro.yaml

After that, three more deploys (re-runnable now that the driver is current):
  scripts/elway irv-ml1 --playbook playbooks/deploy-vibevoice.yaml
  scripts/elway irv-ml1 --playbook playbooks/deploy-chatterbox.yaml
────────────────────────────────────────────────────────────────────
EOF
}

# rollback helper — restores the snapshot from /tmp/nvidia-pre-upgrade.txt
rollback() {
    say "Rollback: reinstalling pre-upgrade nvidia packages on $HOST"
    ssh -t "$HOST" 'bash -s' <<'EOF'
set -euo pipefail
[ -f /tmp/nvidia-pre-upgrade.txt ] || { echo "no snapshot found"; exit 1; }
sudo apt-get install -y --allow-downgrades $(cat /tmp/nvidia-pre-upgrade.txt | tr '\n' ' ')
sudo systemctl reboot
EOF
    ok "Rollback dispatched. Host rebooting."
}

case "${1:-}" in
    phase1)   phase1 ;;
    phase2)   phase2 ;;
    rollback) rollback ;;
    *) echo "usage: $0 {phase1|phase2|rollback}" >&2; exit 2 ;;
esac