Manage held nvidia APT packages

nvidia-apt.sh

#!/usr/bin/env bash
set -euo pipefail

MODE="${1:-}"

# Headless CUDA/job host package matcher.
# Intentionally includes:
# - nvidia driver/kernel/dkms/open modules
# - libnvidia userspace/CUDA-facing libs
# - libcuda
# - nvidia-container-toolkit stack
# - nvidia-persistenced
# - nvidia-firmware / kernel-common/source
#
# Intentionally does NOT need to include GUI/Xorg/settings packages unless installed.
PATTERN='^(nvidia|libnvidia|libcuda|cuda|nvidia-container|libnvidia-container)'

usage() {
  cat <<EOF
Usage:
  nvidia-apt peek
  nvidia-apt update
  nvidia-apt hold
  nvidia-apt rehold

Modes:
  peek     Show NVIDIA/CUDA package state, holds, upgradables, and simulated upgrade
  update   Unhold NVIDIA/CUDA stack, apt update, simulate, ask, then apt upgrade
  hold     Hold installed NVIDIA/CUDA stack packages
  rehold   Same as hold

After update:
  sudo reboot
  nvidia-smi
EOF
}

nvidia_pkgs_installed() {
  dpkg-query -W -f='${Package}\n' 2>/dev/null \
    | grep -Ei "$PATTERN" \
    | sort -u || true
}

nvidia_pkgs_held() {
  apt-mark showhold 2>/dev/null \
    | grep -Ei "$PATTERN" \
    | sort -u || true
}

nvidia_pkgs_upgradable() {
  apt list --upgradable 2>/dev/null \
    | awk -F/ 'NR > 1 {print $1}' \
    | grep -Ei "$PATTERN" \
    | sort -u || true
}

print_status() {
  echo "== nvidia-smi =="
  nvidia-smi || true
  echo

  echo "== DKMS NVIDIA =="
  dkms status 2>/dev/null | grep -i nvidia || true
  echo

  echo "== Installed NVIDIA/CUDA-ish packages =="
  nvidia_pkgs_installed
  echo

  echo "== Held NVIDIA/CUDA-ish packages =="
  nvidia_pkgs_held
  echo

  echo "== Upgradable NVIDIA/CUDA-ish packages =="
  nvidia_pkgs_upgradable
  echo
}

hold_installed() {
  mapfile -t pkgs < <(nvidia_pkgs_installed)

  if (( ${#pkgs[@]} == 0 )); then
    echo "No installed NVIDIA/CUDA-ish packages found to hold."
    return 0
  fi

  echo "Holding installed NVIDIA/CUDA-ish packages:"
  printf '  %s\n' "${pkgs[@]}"
  echo

  sudo apt-mark hold "${pkgs[@]}"
}

unhold_held() {
  mapfile -t pkgs < <(nvidia_pkgs_held)

  if (( ${#pkgs[@]} == 0 )); then
    echo "No held NVIDIA/CUDA-ish packages found."
    return 0
  fi

  echo "Unholding NVIDIA/CUDA-ish packages:"
  printf '  %s\n' "${pkgs[@]}"
  echo

  sudo apt-mark unhold "${pkgs[@]}"
}

simulate_upgrade() {
  echo "== apt upgrade simulation =="
  sudo apt upgrade --simulate | tee /tmp/nvidia-apt-upgrade-sim.txt
  echo
  echo "Simulation saved to:"
  echo "  /tmp/nvidia-apt-upgrade-sim.txt"
  echo
  echo "NVIDIA/CUDA-ish lines from simulation:"
  grep -Ei 'nvidia|libnvidia|libcuda|cuda|container-toolkit|dkms' /tmp/nvidia-apt-upgrade-sim.txt || true
  echo
}

case "$MODE" in
  peek)
    sudo apt update
    print_status
    simulate_upgrade
    ;;

  update)
    print_status

    echo "This will unhold NVIDIA/CUDA packages and upgrade packages via apt."
    echo "You should expect to reboot after this."
    echo
    read -r -p "Continue? [y/N] " ans
    case "$ans" in
      y|Y|yes|YES) ;;
      *) echo "Aborted."; exit 1 ;;
    esac

    unhold_held

    echo "== apt update =="
    sudo apt update
    echo

    simulate_upgrade

    read -r -p "Run sudo apt upgrade now? [y/N] " ans
    case "$ans" in
      y|Y|yes|YES) ;;
      *) echo "Aborted before upgrade. You may want to run: nvidia-apt hold"; exit 1 ;;
    esac

    sudo apt upgrade

    echo
    echo "Upgrade done."
    echo
    echo "Recommended next:"
    echo "  sudo reboot"
    echo
    echo "After reboot:"
    echo "  nvidia-smi"
    echo "  nvidia-apt hold"
    ;;

  hold|rehold)
    hold_installed
    echo
    echo "Current holds:"
    nvidia_pkgs_held
    ;;

  *)
    usage
    exit 2
    ;;
esac