#!/usr/bin/env bash
# VoltageGPU dedicated bare metal: NVIDIA data center driver and Fabric Manager, version 1.
# Verified on a production 8x H100 SXM node on 2026-10-24 (Ubuntu 24.04 LTS, the image as delivered).
#
# Everything NVIDIA comes from NVIDIA's CUDA repository for Ubuntu 24.04, pinned to one version:
#   - cuda-keyring 1.1-1, the repository signing key (its SHA-256 is checked before it is installed)
#   - nvidia-driver-pinning-595.112.07, then nvidia-open 595.112.07: driver branch R595 (NVIDIA's
#     Production Branch), open kernel modules built by DKMS for the running kernel
#   - nvidia-fabricmanager 595.112.07, the same version, only on NVSwitch boards (HGX 8x H100 or
#     H200), where CUDA does not start without it
# From Ubuntu: the running kernel's headers and its compiler (DKMS needs both), python3-venv and
# pciutils (used by the hwcheck acceptance test). Then it starts Fabric Manager and the persistence
# daemon and waits until the NVLink fabric is ready. Safe to run again: what is done stays done.
# Needs root or sudo. Exit code 0 = ready, 1 = error (the last lines say why), 3 = reboot, then run it again.
#
#   curl -fsSL https://voltagegpu.com/dedicated-setup.sh | bash
set -euo pipefail

DRIVER_VERSION="595.112.07"
DEB_VERSION="595.112.07-1ubuntu1"
REPO="https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64"
KEYRING="cuda-keyring_1.1-1_all.deb"
KEYRING_SHA256="d2a6b11c096396d868758b86dab1823b25e14d70333f1dfa74da5ddaf6a06dba"

log() { printf '[%s] %s\n' "$(date -u +%H:%M:%S)" "$*"; }
die() { log "ERROR: $*" >&2; exit 1; }
as_root() { if [ "$(id -u)" -eq 0 ]; then "$@"; else sudo "$@"; fi; }
apt_get() {
  as_root env DEBIAN_FRONTEND=noninteractive NEEDRESTART_SUSPEND=1 apt-get -y -q -o DPkg::Lock::Timeout=600 \
    -o Dpkg::Options::=--force-confdef -o Dpkg::Options::=--force-confold "$@"
}
installed() { [ "$(dpkg-query -W -f='${Status}' "$1" 2>/dev/null || true)" = "install ok installed" ]; }
# pci_count CLASS: NVIDIA PCI functions of that class (0x0302 = GPU, 0x0680 = NVSwitch)
pci_count() {
  local d n=0
  for d in /sys/bus/pci/devices/*; do
    [ "$(cat "$d/vendor" 2>/dev/null)" = 0x10de ] || continue
    case "$(cat "$d/class" 2>/dev/null)" in "$1"*) n=$((n + 1)) ;; esac
  done
  echo "$n"
}
secure_boot_on() {
  local f
  for f in /sys/firmware/efi/efivars/SecureBoot-*; do
    [ -r "$f" ] && [ "$(od -An -t u1 -j 4 -N 1 "$f" | tr -d ' ')" = 1 ] && return 0
  done
  return 1
}
fabric() { nvidia-smi --query-gpu=fabric.state,fabric.status --format=csv,noheader 2>/dev/null | tr -d ' ' || true; }

# The whole run is one function, called on the last line: piped through bash, nothing runs
# before the script has fully arrived, and no command can read the rest of it from stdin.
main() {
  trap 'log "ERROR: this command failed (exit $?): $BASH_COMMAND" >&2; exit 1' ERR
  # ── Checks before anything is installed ──────────────────────────────────────
  # shellcheck disable=SC1091
  . /etc/os-release
  { [ "${ID:-}" = ubuntu ] && [ "${VERSION_ID:-}" = "24.04" ]; } || die "Ubuntu 24.04 only (this machine: ${PRETTY_NAME:-unknown})"
  [ "$(uname -m)" = x86_64 ] || die "x86_64 only"
  GPUS=$(pci_count 0x0302)
  NVSW=$(pci_count 0x0680)
  log "VoltageGPU dedicated setup: driver $DRIVER_VERSION; $GPUS NVIDIA GPU(s) and $NVSW NVSwitch(es) on the PCI bus; kernel $(uname -r)"
  [ "$GPUS" -gt 0 ] || die "no NVIDIA GPU on the PCI bus"
  if secure_boot_on; then
    die "Secure Boot is on: the kernel would refuse the modules DKMS builds here (a MOK key must be enrolled at the console)"
  fi
  if command -v nvidia-smi >/dev/null 2>&1; then
    CUR=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader -i 0 2>/dev/null || true)
    if [ -n "$CUR" ] && [ "$CUR" != "$DRIVER_VERSION" ]; then
      die "NVIDIA driver $CUR is already loaded; this script installs $DRIVER_VERSION on a machine without a driver"
    fi
  fi

  # ── 1. NVIDIA repository key ─────────────────────────────────────────────────
  if installed cuda-keyring; then
    log "1/5 cuda-keyring already installed"
  else
    log "1/5 cuda-keyring from $REPO"
    TMP=$(mktemp -d)
    trap 'rm -rf "$TMP"' EXIT
    curl -fsSL --retry 3 --max-time 120 -o "$TMP/$KEYRING" "$REPO/$KEYRING"
    echo "$KEYRING_SHA256  $TMP/$KEYRING" | sha256sum -c --quiet - || die "$KEYRING does not match its pinned SHA-256"
    as_root dpkg -i "$TMP/$KEYRING"
  fi

  # ── 2. Version pin first, so that every NVIDIA package resolves to $DRIVER_VERSION ──
  # gcc explicitly: NVIDIA's dkms package only recommends it, and DKMS needs the kernel's own gcc-NN
  KCC=$(grep -o 'gcc-[0-9][0-9]*' /proc/version | head -n 1 || true)
  log "2/5 package lists, version pin, kernel headers, gcc ${KCC:-}, python3-venv, pciutils"
  apt_get update
  apt_get install --no-install-recommends "nvidia-driver-pinning-$DRIVER_VERSION" "linux-headers-$(uname -r)" gcc ${KCC:+"$KCC"} \
    python3-venv pciutils

  # ── 3. Driver, and Fabric Manager of the same version on NVSwitch boards ───────
  PKGS=("nvidia-open=$DEB_VERSION")
  if [ "$NVSW" -gt 0 ]; then PKGS+=("nvidia-fabricmanager=$DEB_VERSION"); fi
  log "3/5 ${PKGS[*]} (the DKMS build of the kernel modules takes a few minutes)"
  apt_get install --no-install-recommends "${PKGS[@]}"

  # ── 4. Kernel modules ────────────────────────────────────────────────────────
  log "4/5 kernel modules"
  if ! grep -q '^nvidia ' /proc/modules; then
    if grep -q '^nouveau ' /proc/modules && ! as_root modprobe -r nouveau; then
      log "the nouveau driver holds the GPUs: reboot, then run this script again"
      exit 3
    fi
    as_root modprobe nvidia
  fi
  as_root modprobe nvidia_uvm
  SEEN=$(nvidia-smi -L | grep -c '^GPU ' || true)
  [ "$SEEN" -ge "$GPUS" ] || die "nvidia-smi sees $SEEN of the $GPUS GPUs on the PCI bus"

  # ── 5. Fabric Manager and persistence ────────────────────────────────────────
  log "5/5 services"
  if [ "$NVSW" -gt 0 ]; then as_root systemctl enable --now nvidia-fabricmanager; fi
  if systemctl cat nvidia-persistenced.service >/dev/null 2>&1; then as_root systemctl enable --now nvidia-persistenced; fi
  case "$(nvidia-smi --query-gpu=persistence_mode --format=csv,noheader)" in
    *Disabled*) as_root nvidia-smi -pm 1 >/dev/null ;;
  esac
  if [ "$NVSW" -gt 0 ]; then
    for _ in $(seq 1 60); do
      FAB=$(fabric)
      if [ -n "$FAB" ] && ! grep -qv '^Completed,Success$' <<< "$FAB"; then break; fi
      sleep 3
    done
    FAB=$(fabric)
    if [ -z "$FAB" ] || grep -qv '^Completed,Success$' <<< "$FAB"; then
      die "the NVLink fabric is not ready ($(paste -sd ' ' <<< "$FAB")): see journalctl -u nvidia-fabricmanager"
    fi
  fi

  log "ready in $((SECONDS / 60)) min $((SECONDS % 60)) s: $SEEN GPUs, driver $(nvidia-smi --query-gpu=driver_version --format=csv,noheader -i 0)$([ "$NVSW" -gt 0 ] && echo ", fabric Completed / Success on every GPU, Fabric Manager $(systemctl is-active nvidia-fabricmanager)"), persistence $(nvidia-smi --query-gpu=persistence_mode --format=csv,noheader | sort -u | paste -sd ' ')"
  log "next: nvidia-smi topo -m (NV18 between every pair on 8x H100 SXM), then the acceptance test:"
  log "  curl -fsSL https://voltagegpu.com/hwcheck.sh | bash -s -- --acceptance --gpus $SEEN --bare-metal"
}

main "$@" < /dev/null