#!/usr/bin/env bash
# =============================================================================
# verify-install.sh — post-install verification for get-cuda-mirror.sh
#
# Run it with the SAME flags you installed with, so it knows what to expect:
#
#   sudo ./verify-install.sh --nvidia-driver=580 --cuda-toolkit=13-0 \
#        --dcgm=true --nvitop=true --no-nv=true
#
#   sudo ./verify-install.sh --doca=true --nvidia-driver=580 --cuda-toolkit=13-0 \
#        --dcgm=true --nvitop=true
#
# With no flags it auto-detects what is on the box and checks internal
# consistency only. Flags turn "is it consistent" into "is it what you asked
# for", which is the whole point of a pin-exact mirror.
#
# Exit code: 0 = everything the flags asked for is there, 1 = at least one FAIL.
# WARNs never fail the run — they are things that are expected to be unfinished
# before a reboot, or that depend on hardware this node may not have.
#
# Read-only. It installs nothing, changes nothing, starts nothing.
# =============================================================================
set -uo pipefail

if [[ -t 1 ]]; then
    C_RED=$'\033[1;31m'; C_GRN=$'\033[0;32m'; C_YLW=$'\033[0;33m'
    C_CYN=$'\033[0;36m'; C_WHT=$'\033[1;37m'; C_DIM=$'\033[2m'; C_NC=$'\033[0m'
else
    C_RED=""; C_GRN=""; C_YLW=""; C_CYN=""; C_WHT=""; C_DIM=""; C_NC=""
fi

N_PASS=0; N_FAIL=0; N_WARN=0
declare -a FAILED=() WARNED=()

hdr()  { printf '\n%s\n' "${C_CYN}====${C_NC} ${C_WHT}$*${C_NC}"; }
pass() { printf '%s\n' "${C_GRN}[PASS]${C_NC} $*"; N_PASS=$((N_PASS+1)); }
fail() { printf '%s\n' "${C_RED}[FAIL]${C_NC} $*"; N_FAIL=$((N_FAIL+1)); FAILED+=("$*"); }
warn() { printf '%s\n' "${C_YLW}[WARN]${C_NC} $*"; N_WARN=$((N_WARN+1)); WARNED+=("$*"); }
info() { printf '%s\n' "${C_DIM}       $*${C_NC}"; }

# ---------------------------------------------------------------- expectations
EXP_DRIVER=""; EXP_CUDA=""; EXP_DCGM=""; EXP_NVITOP=""; EXP_DOCA=""; EXP_NONV=""
DEEP=false

for arg in "$@"; do
    case "$arg" in
        --nvidia-driver=*) EXP_DRIVER="${arg#*=}" ;;
        --cuda-toolkit=*)  EXP_CUDA="${arg#*=}" ;;
        --dcgm=true)       EXP_DCGM=true ;;   --dcgm=false)   EXP_DCGM=false ;;
        --nvitop=true)     EXP_NVITOP=true ;; --nvitop=false) EXP_NVITOP=false ;;
        --doca=true)       EXP_DOCA=true ;;   --doca=false)   EXP_DOCA=false ;;
        --no-nv=true)      EXP_NONV=true ;;   --no-nv=false)  EXP_NONV=false ;;
        --deep)            DEEP=true ;;
        # Accepted and ignored so an install command pastes across verbatim:
        # none of these change what a finished node should look like.
        --open=*|--doca-version=*|--mirror=*|--local-install=*|--force-install=*) ;;
        --install-ip*|--yes|--aws=*|--force=*) ;;
        -h|--help)
            sed -n '2,22p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
        *) echo "unknown option: $arg  (try --help)"; exit 2 ;;
    esac
done

# ------------------------------------------------------------------- utilities

# Installed version of a package, empty if not installed.
pver() {
    local v
    v="$(dpkg-query -W -f='${db:Status-Status} ${Version}' "$1" 2>/dev/null)" || return 1
    [[ "$v" == installed\ * ]] || return 1
    printf '%s' "${v#installed }"
}
have() { pver "$1" >/dev/null 2>&1; }

# Strip epoch and Debian revision: "1:580.178.04-1ubuntu1" -> "580.178.04"
upstream() { local v="${1#*:}"; printf '%s' "${v%%-*}"; }

# An Ubuntu-archive build carries the distro triplet (-0ubuntu0.24.04.1);
# NVIDIA's own repo uses -1ubuntu1 / -0ubuntu1. This is the lineage test the
# mirror itself uses, and telling the two apart is the reason this project exists.
is_ubuntu_lineage() { [[ "$1" =~ -[0-9]+ubuntu[0-9]+\.[0-9]+\.[0-9]+ ]]; }

# Every installed package in the NVIDIA driver namespace: "name<TAB>version".
nvidia_pkgs() {
    dpkg-query -W -f='${db:Status-Status}\t${Package}\t${Version}\n' \
        'nvidia-*' 'libnvidia-*' 'xserver-xorg-video-nvidia*' \
        'nvlink5*' 'libnvsdm*' 'libxnvctrl*' 2>/dev/null \
        | awk -F'\t' '$1=="installed"{print $2 "\t" $3}'
}

# A driver-versioned package has an upstream like 580.178.04 — three parts, the
# first three digits. Enablers (dkms 1:3.4.1, container-toolkit 1.20.0,
# egl-gbm 1.1.3, DCGM 1:4.6.1) deliberately do NOT match: they track their own
# release trains and must not be compared against the driver.
is_driver_versioned() { [[ "$(upstream "$1")" =~ ^[0-9]{3}\.[0-9]+\.[0-9]+$ ]]; }

hdr "Node"
printf '%s\n' "${C_DIM}       $(. /etc/os-release 2>/dev/null; echo "${PRETTY_NAME:-unknown}")  ·  kernel $(uname -r)  ·  $(hostname)${C_NC}"
[[ $EUID -ne 0 ]] && info "(not root — a few checks are skipped; re-run with sudo for the full set)"

# =============================================================================
hdr "1. Kernel modules"
# =============================================================================
loaded_ver=""
if [[ -r /proc/driver/nvidia/version ]]; then
    loaded_ver="$(awk '{for(i=1;i<=NF;i++) if ($i ~ /^[0-9]{3}\.[0-9]+\.[0-9]+$/) {print $i; exit}}' /proc/driver/nvidia/version)"
fi
# /sys/module is the kernel's own answer and it is the primary source; lsmod is a
# second opinion. Parsing lsmod's columns reported 'nvidia' NOT loaded on a node
# where nvidia_uvm, nvidia_modeset and nvidia_drm all PASSed — which cannot happen,
# since each of those carries 'nvidia' in its modinfo depends and the loader will
# not insert them without it. In the same run /proc/driver/nvidia/version was
# readable, nvidia-smi answered 575.57.08 and gpu_burn ran clean. The identical miss
# hit 'nvidia_peermem' in §9 while the boot journal showed systemd-modules-load had
# inserted it minutes earlier. Two modules, one text match, both wrong the same way.
# Reporting "the driver is not running" on a healthy node is worse than no check at
# all — it sends someone to reboot for nothing. Ask the kernel. Either source
# proving the module present is enough; both must miss it before this fails.
mod_loaded() { [[ -d "/sys/module/${1}" ]] || lsmod 2>/dev/null | grep -q "^${1}[[:space:]]"; }
for m in nvidia nvidia_uvm nvidia_modeset nvidia_drm; do
    if mod_loaded "$m"; then
        pass "module loaded: $m"
    else
        fail "module NOT loaded: $m"
        info "the driver is installed but not running — reboot, or 'sudo modprobe $m'"
    fi
done
[[ -n "$loaded_ver" ]] && info "running module version: $loaded_ver"

# One pending reboot produces a page of unrelated-looking failures: NVML refuses to
# initialise, dcgmi finds no GPU, nvitop cannot read the device. All three are the
# same single fact — a new driver is on disk while the old module is still resident —
# and printing them as three independent FAILs buries the one line that explains them.
# That is exactly the state the intended bare-metal loop passes through on purpose
# (uninstall, install, reboot ONCE), so it is worth reporting as what it is.
# Establish the fact here, ahead of everything that consumes it, by comparing what is
# running against what depmod indexed on disk. Both readings come straight from the
# kernel and neither needs NVML, so this still works precisely when nvidia-smi does
# not. Downstream checks soften to warnings that point back here — but ONLY when this
# is set, so a genuinely dead GPU on a settled node still fails loudly.
disk_ver="$(modinfo -F version nvidia 2>/dev/null || true)"
reboot_pending=0
if [[ -n "$loaded_ver" && -n "$disk_ver" && "$loaded_ver" != "$disk_ver" ]]; then
    reboot_pending=1
fi

# =============================================================================
hdr "2. Driver version lock"
# =============================================================================
# The core promise of the mirror: pick X, get exactly X — on every package, on
# every machine, forever. Three separate ways that can break, checked separately.
smi_ver=""; smi_cuda=""; pkg_ver=""; smi_out=""; smi_err=""
if command -v nvidia-smi >/dev/null 2>&1; then
    # When NVML refuses to load, nvidia-smi writes the reason to stdout and the
    # --query-gpu form hands that sentence back where a version belongs. So the
    # value is only accepted once it actually looks like a driver version — taking
    # it on trust turns a broken node into a PASS with a garbled version string.
    smi_out="$(nvidia-smi 2>&1)"
    smi_ver="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -n1 | tr -d ' ')"
    [[ "$smi_ver" =~ ^[0-9]{3}\.[0-9]+(\.[0-9]+)?$ ]] || smi_ver=""
    smi_cuda="$(grep -oE 'CUDA (UMD )?Version: *[0-9]+\.[0-9]+' <<< "$smi_out" | grep -oE '[0-9]+\.[0-9]+' | head -n1)"
    # Some builds keep the max CUDA out of the header table; -q still carries it.
    [[ -z "$smi_cuda" ]] && smi_cuda="$(nvidia-smi -q 2>/dev/null | awk -F': *' '/CUDA Version/{print $2; exit}' | tr -d ' ')"
    [[ "$smi_cuda" =~ ^[0-9]+\.[0-9]+$ ]] || smi_cuda=""
    [[ -z "$smi_ver" ]] && smi_err="$(grep -m1 -iE 'mismatch|failed|error|no devices' <<< "$smi_out")"
fi
if [[ -n "$smi_ver" ]]; then
    pass "nvidia-smi reports driver $smi_ver  (max CUDA ${smi_cuda:-unknown})"
elif [[ "$reboot_pending" == 1 ]]; then
    # Already explained by the module-vs-package FAIL below, which names both
    # versions. Repeating it as a second FAIL would double-count one cause.
    warn "nvidia-smi cannot report a driver version yet — reboot pending (see below)"
    [[ -n "$smi_err" ]] && info "nvidia-smi said: $smi_err"
else
    fail "nvidia-smi did not report a usable driver version"
    [[ -n "$smi_err" ]] && info "nvidia-smi said: $smi_err"
    info "most often a pending reboot — a new driver is unpacked while the old kernel"
    info "module is still loaded. The module-vs-package check below settles it."
fi

# 2a. every driver-versioned package must carry the SAME version
mapfile -t drvrows < <(nvidia_pkgs | while IFS=$'\t' read -r p v; do
    is_driver_versioned "$v" && printf '%s\t%s\n' "$p" "$v"
done)
if ((${#drvrows[@]} == 0)); then
    fail "no driver-versioned NVIDIA packages installed at all"
else
    # The thing that must be identical is the DRIVER BUILD — the upstream version.
    # The Debian revision after the dash is NVIDIA's own packaging counter and it
    # legitimately differs between sub-packages of one build: on 575.57.08 the
    # NVIDIA repo ships the driver at -0ubuntu1, nvidia-fabricmanager-575 at -1 and
    # nvidia-persistenced at -1ubuntu1, all resolved by apt in a single pinned
    # transaction. Comparing full version strings flagged that as "the exact bug the
    # mirror exists to prevent", which is the opposite of the truth — it is one
    # coherent build. 580 happened to use one revision throughout, so the false
    # positive only surfaced on a Gen-1 branch. Group by upstream, report the
    # revision spread as detail.
    mapfile -t uniqv  < <(printf '%s\n' "${drvrows[@]}" | cut -f2 | sort -u)
    # upstream() prints without a trailing newline, so the echo is load-bearing:
    # without it sort -u sees one concatenated line and ALWAYS reports a single
    # build — the SPLIT failure below would never fire again.
    mapfile -t uniqup < <(printf '%s\n' "${uniqv[@]}" | while read -r v; do echo "$(upstream "$v")"; done | sort -u)
    if ((${#uniqup[@]} == 1)); then
        pass "all ${#drvrows[@]} driver packages at one build: ${uniqup[0]}"
        if ((${#uniqv[@]} > 1)); then
            info "${#uniqv[@]} packaging revisions of that build, which is normal for NVIDIA's repo:"
            for v in "${uniqv[@]}"; do
                printf '%s\n' "${C_DIM}         ${v#*-} :  $(printf '%s\n' "${drvrows[@]}" | awk -F'\t' -v w="$v" '$2==w{printf "%s ", $1}')${C_NC}"
            done
        fi
    else
        fail "driver packages are SPLIT across ${#uniqup[@]} builds — this is the exact bug the mirror exists to prevent:"
        for u in "${uniqup[@]}"; do
            printf '%s\n' "${C_DIM}         $u :  $(printf '%s\n' "${drvrows[@]}" | awk -F'\t' -v w="$u" '{u=$2; sub(/^[^:]*:/,"",u); sub(/-.*$/,"",u); if (u==w) printf "%s ", $1}')${C_NC}"
        done
    fi
    pkg_ver="${uniqv[0]}"
    # 2b. the running module must be the installed one (else a reboot is pending)
    if [[ -n "$loaded_ver" && -n "$pkg_ver" ]]; then
        if [[ "$loaded_ver" == "$(upstream "$pkg_ver")" ]]; then
            pass "running module matches the installed package ($loaded_ver)"
        else
            fail "running module $loaded_ver != installed package $(upstream "$pkg_ver") — REBOOT PENDING"
        fi
    fi
fi

# 2c. lineage — nothing from the Ubuntu archive may have crept in
leak="$(nvidia_pkgs | while IFS=$'\t' read -r p v; do
            is_ubuntu_lineage "$v" && printf '  %s %s\n' "$p" "$v"
        done)"
if [[ -z "$leak" ]]; then
    pass "no Ubuntu-archive builds in the NVIDIA stack (NVIDIA lineage only)"
else
    fail "Ubuntu-archive packages found in the NVIDIA stack:"
    printf '%s\n' "${C_DIM}$leak${C_NC}"
    info "Ubuntu publishes one build per branch and rewrites it in place — these"
    info "will silently move to a different build on the next machine."
fi

# 2d. did we get the build that was asked for?
if [[ -n "$EXP_DRIVER" && -n "$smi_ver" ]]; then
    if [[ "$EXP_DRIVER" == *.* ]]; then
        [[ "$smi_ver" == "$EXP_DRIVER" ]] \
            && pass "driver is the pin-exact build requested: $EXP_DRIVER" \
            || fail "asked for driver $EXP_DRIVER but got $smi_ver"
    else
        [[ "${smi_ver%%.*}" == "$EXP_DRIVER" ]] \
            && pass "driver is on the requested branch $EXP_DRIVER (build $smi_ver)" \
            || fail "asked for branch $EXP_DRIVER but got $smi_ver"
    fi
fi
# Everything downstream (FM name, module flavor, DCGM) needs the branch. When
# nvidia-smi is down the installed package still knows it, so fall back to that
# rather than carrying an empty branch through five more sections.
drv_branch="${smi_ver%%.*}"
[[ -z "$drv_branch" && -n "$pkg_ver" ]] && drv_branch="$(upstream "$pkg_ver")" && drv_branch="${drv_branch%%.*}"

# 2e. kernel module flavor, for the record
if have nvidia-dkms-open || have "nvidia-dkms-${drv_branch}-open" \
   || have nvidia-driver-open || have "nvidia-driver-${drv_branch}-open"; then
    info "kernel module flavor: open"
else
    info "kernel module flavor: proprietary"
fi

# =============================================================================
hdr "3. Secure Boot"
# =============================================================================
if command -v mokutil >/dev/null 2>&1 && [[ "$(mokutil --sb-state 2>/dev/null)" == *"SecureBoot enabled"* ]]; then
    signer="$(modinfo -F signer nvidia 2>/dev/null | head -n1)"
    if [[ -n "$signer" ]]; then
        pass "Secure Boot ON, nvidia module signed by: $signer"
    else
        fail "Secure Boot is ON but the nvidia module carries NO signature"
        info "it will build fine and then refuse to load after reboot ('no devices found')"
    fi
else
    info "Secure Boot is off or unavailable — module signing not required"
fi

# =============================================================================
hdr "4. CUDA toolkit"
# =============================================================================
nvcc_bin=""
for c in /usr/local/cuda/bin/nvcc "$(command -v nvcc 2>/dev/null)"; do
    [[ -n "$c" && -x "$c" ]] && { nvcc_bin="$c"; break; }
done
nvcc_rel=""
if [[ -n "$nvcc_bin" ]]; then
    nvcc_rel="$("$nvcc_bin" --version 2>/dev/null | grep -oE 'release [0-9]+\.[0-9]+' | awk '{print $2}')"
    pass "nvcc $nvcc_rel  ($nvcc_bin)"
else
    if [[ -n "$EXP_CUDA" ]]; then
        fail "nvcc not found"
        info "the installer puts it on PATH via /etc/profile.d/cuda.sh (new logins only)"
    else
        info "nvcc not found — no CUDA toolkit expected"
    fi
fi
if [[ -e /usr/local/cuda ]]; then
    tgt="$(readlink -f /usr/local/cuda)"
    pass "/usr/local/cuda -> $tgt"
    others="$(ls -d /usr/local/cuda-1*.* 2>/dev/null | grep -v "^${tgt}$" | tr '\n' ' ')"
    [[ -n "$others" ]] && info "other toolkits also present: $others"
fi
if [[ -z "$EXP_CUDA" ]]; then
    # No --cuda-toolkit given: nothing to assert, so at least name what is here.
    tk="$(dpkg-query -W -f='${db:Status-Status} ${Package} ${Version}\n' 'cuda-toolkit-*' 2>/dev/null \
          | awk '$1=="installed" && $2 ~ /^cuda-toolkit-[0-9]+-[0-9]+$/ {printf "%s %s   ", $2, $3}')"
    [[ -n "$tk" ]] && info "installed toolkits: $tk"
    info "pass --cuda-toolkit=<ver> to check this against what you asked for"
else
    want_dot="${EXP_CUDA//-/.}"
    tv="$(pver "cuda-toolkit-$EXP_CUDA")" \
        && pass "cuda-toolkit-$EXP_CUDA $tv" \
        || fail "cuda-toolkit-$EXP_CUDA is NOT installed"
    if [[ -n "$nvcc_rel" ]]; then
        if [[ "$nvcc_rel" == "$want_dot" ]]; then
            pass "nvcc is the requested toolkit ($want_dot)"
        else
            fail "asked for CUDA $want_dot but nvcc reports $nvcc_rel"
            if [[ -d "/usr/local/cuda-$want_dot" ]]; then
                info "both toolkits are installed; the /usr/local/cuda symlink just points at $nvcc_rel."
                info "fix: sudo ln -sfn /usr/local/cuda-$want_dot /usr/local/cuda"
            else
                info "toolkit $want_dot is NOT under /usr/local — install it first:"
                info "    sudo apt-get install cuda-toolkit-${EXP_CUDA}"
            fi
        fi
    fi

    # The 'cuda-<ver>' meta = cuda-toolkit + cuda-runtime. Up to CUDA 13.0 the
    # runtime hard-Depends on ONE driver branch (cuda-runtime-13-0 ->
    # libnvidia-compute-580), so on a mismatched branch the installer skips the
    # meta on purpose and nothing is lost. From 13.1 NVIDIA dropped that dep, so
    # the meta must install on any branch. Work out which case this is instead of
    # guessing, so a deliberate skip never reads as a failure.
    need_branch="$(LC_ALL=C apt-cache show "cuda-runtime-$EXP_CUDA" 2>/dev/null \
                   | awk -F'libnvidia-compute-' '/^Depends:/ && /libnvidia-compute-/ {print $2; exit}' \
                   | grep -oE '^[0-9]+')"
    if have "cuda-$EXP_CUDA"; then
        pass "meta cuda-$EXP_CUDA $(pver "cuda-$EXP_CUDA") installed"
    elif [[ -n "$need_branch" && -n "$drv_branch" && "$need_branch" != "$drv_branch" ]]; then
        pass "meta cuda-$EXP_CUDA correctly skipped — cuda-runtime-$EXP_CUDA needs driver branch $need_branch, this node runs $drv_branch"
        info "cosmetic only: the toolkit holds every library; the runtime would only add a driver"
    else
        warn "meta cuda-$EXP_CUDA is not installed although branch $drv_branch should satisfy it"
        info "check the install log for a 'Skipping meta' line and its reason"
    fi
fi

# NVLink state is read here, before Fabric Manager, because FM's expectations
# depend on it: an --no-nv install has no FM by design. Run without flags, the
# marker file is the only evidence of which kind of install this node had.
nvl="$(awk -F': *' '/NvLinkDisable/{print $2}' /proc/driver/nvidia/params 2>/dev/null | tr -d ' ')"
conf_exists=false; [[ -e /etc/modprobe.d/disable-nvlink.conf ]] && conf_exists=true
nonv_eff="$EXP_NONV"
[[ -z "$nonv_eff" && "$conf_exists" == true ]] && nonv_eff=inferred

# =============================================================================
hdr "5. Fabric Manager"
# =============================================================================
# FM is installed only when --no-nv is absent, and must match the driver
# EXACTLY (not merely the same branch) — that coupling is what dragged pinned
# drivers forward when FM came from the Ubuntu archive.
fm_pkg=""
for f in "nvidia-fabricmanager-${drv_branch}" nvidia-fabricmanager; do
    have "$f" && { fm_pkg="$f"; break; }
done
if [[ "$nonv_eff" == true || "$nonv_eff" == inferred ]]; then
    src_note="as expected with --no-nv=true"
    [[ "$nonv_eff" == inferred ]] && src_note="expected: disable-nvlink.conf shows this node was installed with --no-nv=true"
    if [[ -z "$fm_pkg" ]]; then
        pass "no Fabric Manager, $src_note"
    else
        warn "--no-nv is in effect but Fabric Manager is installed ($fm_pkg $(pver "$fm_pkg"))"
        info "probably left over from an earlier run on this node"
    fi
elif [[ -n "$fm_pkg" ]]; then
    fm_ver="$(pver "$fm_pkg")"
    if [[ -n "$pkg_ver" && "$(upstream "$fm_ver")" == "$(upstream "$pkg_ver")" ]]; then
        pass "$fm_pkg $fm_ver — exact version match with the driver"
    else
        fail "$fm_pkg is $fm_ver but the driver is $pkg_ver — FM must match the driver exactly"
    fi
    st="$(systemctl is-active nvidia-fabricmanager 2>/dev/null)"
    en="$(systemctl is-enabled nvidia-fabricmanager 2>/dev/null)"
    if [[ "$st" == active ]]; then
        pass "nvidia-fabricmanager.service is active (enabled: ${en:-unknown})"
    else
        # Whether a stopped FM is correct or broken depends on ONE thing: does this
        # box actually have NVSwitches? Deciding that from the journal text alone was
        # wrong in both directions. NVIDIA's real message on a plain GPU node is
        #   "request to query NVSwitch device information ... failed with
        #    error:WARNING Nothing to do [NV_WARN_NOTHING_TO_DO]"
        # which matched none of the old patterns, so every non-NVSwitch node WARNed
        # forever. Widening the grep would have been worse: that same line on an HGX
        # box means the NVSwitches are not visible to the driver — a genuine fault
        # that would then have been reported as a PASS. So count the hardware instead
        # and let the journal line be commentary. Two independent sources, because the
        # nvswitch procfs tree only exists once the nvidia-nvswitch driver has bound.
        nvsw_n=0
        if [[ -d /proc/driver/nvidia-nvswitch/devices ]]; then
            nvsw_n="$(find /proc/driver/nvidia-nvswitch/devices -mindepth 1 -maxdepth 1 2>/dev/null | wc -l)"
        fi
        if [[ "$nvsw_n" -eq 0 ]] && command -v lspci >/dev/null 2>&1; then
            nvsw_n="$(lspci -d 10de: 2>/dev/null | grep -ci nvswitch)"
        fi
        why="$(journalctl -u nvidia-fabricmanager -n 20 --no-pager 2>/dev/null \
               | grep -oiE 'no nvswitch[^.]*|not supported[^.]*|nvswitch.*not found|nothing to do[^.]*|NV_WARN_NOTHING_TO_DO' | tail -n1)"
        if [[ "$nvsw_n" -eq 0 ]]; then
            pass "nvidia-fabricmanager not running — no NVSwitch on this node, so FM has nothing to manage"
            info "FM stays installed and enabled on purpose: it is version-locked to the driver,"
            info "and it starts by itself if this image is ever deployed on NVSwitch hardware."
            [[ -n "$why" ]] && info "journal: $why"
        else
            # FM down on a box that HAS NVSwitches means the NVLink fabric was never
            # initialised. The GPUs are up and nvidia-smi looks normal, but they cannot
            # reach each other over NVLink — that is a failure, not a warning.
            fail "nvidia-fabricmanager.service is '$st' but this node has $nvsw_n NVSwitch device(s)"
            info "the NVLink fabric is NOT initialised — GPUs cannot talk over NVLink. Inspect with:"
            info "  sudo systemctl status nvidia-fabricmanager --no-pager -l"
            [[ -n "$why" ]] && info "journal: $why"
        fi
    fi
else
    warn "no Fabric Manager installed"
    info "expected on an NVSwitch/HGX node. If the mirror had no FM at this exact"
    info "driver build the installer says so and installs the driver alone."
fi

# =============================================================================
hdr "6. NVLink"
# =============================================================================
# Always state the facts first — with no flags there is nothing to judge against,
# and a section that printed nothing at all looked like the check had crashed.
info "NvLinkDisable: ${nvl:-unreadable}   ·   disable-nvlink.conf: $([[ "$conf_exists" == true ]] && echo present || echo absent)"
if [[ -e /etc/modprobe.d/nvlink-denylist.conf ]]; then
    fail "/etc/modprobe.d/nvlink-denylist.conf exists — left by an OLD --no-nv run"
    info "it blacklists nvidia_uvm, which kills the CUDA runtime. Re-run the installer to clear it."
fi
if [[ "$EXP_NONV" == true ]]; then
    [[ "$conf_exists" == true ]] \
        && pass "/etc/modprobe.d/disable-nvlink.conf present" \
        || fail "--no-nv=true but /etc/modprobe.d/disable-nvlink.conf is missing"
    if [[ "$nvl" == 1 ]]; then
        pass "NvLinkDisable: 1 — NVLink is off, as requested"
    else
        warn "NvLinkDisable: ${nvl:-unknown} — the setting only takes effect after a REBOOT"
    fi
elif [[ -n "$EXP_NONV" || -n "$EXP_DOCA" || -n "$EXP_DRIVER" ]]; then
    # No --no-nv on this run: NVLink should be ON. The installer does NOT delete
    # a disable-nvlink.conf written by an earlier --no-nv run, so a node used for
    # both tests keeps NVLink off and nothing on the console says so.
    if [[ "$conf_exists" == true ]]; then
        fail "--no-nv was NOT used, but /etc/modprobe.d/disable-nvlink.conf still exists"
        info "left behind by an earlier --no-nv run on this node; the installer does not remove it."
        info "  sudo rm /etc/modprobe.d/disable-nvlink.conf && sudo update-initramfs -u && sudo reboot"
    else
        pass "no disable-nvlink.conf — NVLink not suppressed"
    fi
    if [[ "$nvl" == 0 ]]; then
        pass "NvLinkDisable: 0 — NVLink enabled"
    else
        fail "NvLinkDisable: ${nvl:-unknown} — NVLink is OFF but this run did not ask for that"
    fi
elif [[ "$conf_exists" == true ]]; then
    info "this node carries the --no-nv marker, so NVLink is deliberately off"
    info "pass --no-nv=true (or omit it) to have that checked instead of just reported"
else
    info "NVLink not suppressed"
fi

# =============================================================================
hdr "7. DCGM"
# =============================================================================
if [[ "$EXP_DCGM" == false ]]; then
    info "not requested — skipped"
elif have datacenter-gpu-manager-4-core || [[ "$EXP_DCGM" == true ]]; then
    core="$(pver datacenter-gpu-manager-4-core)" \
        && pass "datacenter-gpu-manager-4-core $core" \
        || fail "datacenter-gpu-manager-4-core is NOT installed"
    # The dcgmi diag plugin follows the DRIVER's max CUDA, not the toolkit:
    # 580 -> 13.0 -> cuda13, 575 -> 12.9 -> cuda12.
    want_major="${smi_cuda%%.*}"
    if [[ -n "$want_major" ]]; then
        if have "datacenter-gpu-manager-4-cuda${want_major}"; then
            pass "plugin datacenter-gpu-manager-4-cuda${want_major} matches the driver's max CUDA ($smi_cuda)"
        else
            got="$(dpkg-query -W -f='${db:Status-Status} ${Package}\n' 'datacenter-gpu-manager-4-cuda*' 2>/dev/null \
                   | awk '$1=="installed"{printf "%s ", $2}')"
            fail "expected plugin datacenter-gpu-manager-4-cuda${want_major}, installed: ${got:-none}"
            info "dcgmi diag loads the plugin for the driver's CUDA, so a mismatch makes diag fail to run"
        fi
    else
        # Deliberately not guessed from the branch: the branch -> max-CUDA map
        # moves with each release, and a wrong guess here fails a healthy node.
        info "plugin match not checked — nvidia-smi could not report the driver's max CUDA"
    fi
    st="$(systemctl is-active nvidia-dcgm 2>/dev/null)"
    [[ "$st" == active ]] && pass "nvidia-dcgm.service active" || warn "nvidia-dcgm.service is '$st'"
    if command -v dcgmi >/dev/null 2>&1; then
        if out="$(timeout 60 dcgmi discovery -l 2>&1)"; then
            n="$(grep -oE '[0-9]+ GPUs? found' <<< "$out" | grep -oE '^[0-9]+' | head -n1)"
            extra="$(grep -oE '[0-9]+ (NvSwitches|ConnectX) found' <<< "$out" | tr '\n' '|' | sed 's/|$//; s/|/, /g')"
            if [[ "${n:-0}" -gt 0 ]]; then
                pass "dcgmi discovery: ${n} GPU(s)${extra:+, $extra}"
            elif [[ "$reboot_pending" == 1 ]]; then
                # DCGM talks to the GPU through NVML, so a resident old module hides
                # every device from it. Nothing is wrong with DCGM here — the same
                # mismatch §2 already failed on is the whole explanation.
                warn "dcgmi finds no GPU — reboot pending, NVML cannot reach the device"
            else
                # dcgmi exits 0 while reporting nothing when the driver is half-up.
                fail "dcgmi ran but found no GPU"
                printf '%s\n' "${C_DIM}$(sed 's/^/         /' <<< "$out" | head -n 6)${C_NC}"
            fi
        else
            fail "dcgmi discovery -l failed"
            printf '%s\n' "${C_DIM}$(sed 's/^/         /' <<< "$out" | head -n 8)${C_NC}"
        fi
        if [[ "$DEEP" == true ]]; then
            info "running dcgmi diag -r 2 (a few minutes) ..."
            if timeout 900 dcgmi diag -r 2 >/tmp/dcgmi-diag.$$ 2>&1; then
                grep -qiE '\bfail\b' /tmp/dcgmi-diag.$$ \
                    && { fail "dcgmi diag -r 2 reported a failure — see /tmp/dcgmi-diag.$$"; } \
                    || pass "dcgmi diag -r 2 all Pass"
            else
                fail "dcgmi diag -r 2 did not complete — see /tmp/dcgmi-diag.$$"
            fi
        fi
    else
        fail "dcgmi not on PATH"
    fi
else
    info "DCGM not installed and not requested"
fi

# =============================================================================
hdr "8. nvitop"
# =============================================================================
if [[ "$EXP_NVITOP" == false ]]; then
    info "not requested — skipped"
elif command -v nvitop >/dev/null 2>&1; then
    pass "nvitop $(nvitop --version 2>/dev/null | grep -oE '[0-9]+\.[0-9]+\.[0-9]+' | head -n1)"
    timeout 60 nvitop -1 >/dev/null 2>&1 \
        && pass "nvitop -1 runs and reads the GPU" \
        || warn "nvitop is installed but 'nvitop -1' did not succeed"
elif [[ "$EXP_NVITOP" == true ]]; then
    fail "nvitop not on PATH"
else
    info "nvitop not installed and not requested"
fi

# =============================================================================
hdr "9. DOCA-OFED"
# =============================================================================
if [[ "$EXP_DOCA" == true ]] || have doca-ofed; then
    v="$(pver doca-ofed)" && pass "doca-ofed $v" || fail "doca-ofed is NOT installed"
    v="$(pver mft)"       && pass "mft $v"       || fail "mft is NOT installed"
    if command -v ofed_info >/dev/null 2>&1; then
        pass "ofed_info -s: $(ofed_info -s 2>/dev/null | head -n1)"
    else
        warn "ofed_info not on PATH"
    fi
    # ib_core must be DOCA's, not the inbox one — that only happens after reboot.
    if mod_loaded ib_core; then
        pass "ib_core loaded"
    else
        warn "ib_core not loaded — reboot to replace the inbox IB modules"
    fi
    # nvidia_peermem is the whole point of pairing DOCA with the driver:
    # GPUDirect RDMA does not work without it.
    if mod_loaded nvidia_peermem; then
        pass "nvidia_peermem loaded — GPUDirect RDMA path is up"
    else
        warn "nvidia_peermem NOT loaded — GPUDirect RDMA is not active"
        info "after the reboot:  sudo modprobe nvidia_peermem  &&  lsmod | grep nvidia_peermem"
    fi
    sym="/usr/src/ofa_kernel-dkms/x86_64/$(uname -r)/Module.symvers"
    if [[ -f "$sym" ]] && grep -q ib_register_peer_memory_client "$sym"; then
        pass "ib_peer_mem symbols present for $(uname -r)"
    else
        warn "ib_peer_mem symbols not found for the running kernel"
        info "nvidia_peermem cannot link without them — expected before the first reboot"
    fi
    if command -v ibstat >/dev/null 2>&1; then
        ports="$(ibstat 2>/dev/null | grep -c 'State: Active')"
        if [[ "${ports:-0}" -gt 0 ]]; then
            pass "ibstat: $ports active IB port(s)"
        else
            warn "ibstat reports no active port (no cable, or IB fabric down)"
        fi
    else
        warn "ibstat not on PATH"
    fi
elif [[ "$EXP_DOCA" == false ]]; then
    have doca-ofed && warn "--doca was not requested but doca-ofed is installed" \
                   || pass "no DOCA, as expected"
else
    info "DOCA not installed and not requested"
fi

# =============================================================================
hdr "10. Container toolkit"
# =============================================================================
if v="$(pver nvidia-container-toolkit)"; then
    pass "nvidia-container-toolkit $v"
else
    info "nvidia-container-toolkit not installed (optional)"
fi

# =============================================================================
hdr "11. Hygiene"
# =============================================================================
# The installer writes a hard apt pin while it works and removes it on exit via
# a trap. A leftover pin makes every later apt run behave strangely for reasons
# nothing on screen explains, so check for it explicitly.
pins="$(ls /etc/apt/preferences.d/zz-nvidia-driver-lock-* 2>/dev/null)"
if [[ -z "$pins" ]]; then
    pass "no leftover apt driver pin"
else
    fail "leftover apt pin(s) — apt will keep refusing packages until removed:"
    printf '%s\n' "${C_DIM}$(sed 's/^/         /' <<< "$pins")${C_NC}"
    info "  sudo rm $pins"
fi
holds="$(apt-mark showhold 2>/dev/null | grep -E 'nvidia|cuda' | tr '\n' ' ')"
[[ -z "$holds" ]] && pass "no nvidia/cuda packages on hold" || warn "packages on hold: $holds"
if [[ $EUID -eq 0 ]]; then
    broken="$(LC_ALL=C apt-get -s install -f 2>/dev/null | grep -cE '^(Inst|Remv) ')"
    [[ "${broken:-0}" -eq 0 ]] \
        && pass "apt dependency tree is clean" \
        || warn "'apt-get install -f' wants to change ${broken} package(s) — the tree is not clean"
fi
[[ -e /var/run/reboot-required ]] && warn "a reboot is pending (/var/run/reboot-required)"

printf '%s\n' "${C_DIM}       disk:${C_NC}"
df -Pm / /usr/local /var/cache/apt/archives 2>/dev/null | awk 'NR>1{printf "         %-42s %7s MB free   %s\n", $1, $4, $6}' | sort -u
cache_mb="$(df -Pm /var/cache/apt/archives 2>/dev/null | awk 'NR==2{print $4}')"
# Free space alone is the wrong number: apt stages .debs here and leaves them
# there, so a node that just installed CUDA is "short" purely because it is
# holding the .debs it no longer needs. Warning on that is noise — it fires on a
# perfectly healthy node and never clears, which trains people to ignore it.
# What matters is free + reclaimable, i.e. what 'apt-get clean' can actually give
# back. Only warn when even that total cannot stage the next toolkit.
cache_used_mb="$(du -sm /var/cache/apt/archives 2>/dev/null | awk '{print $1}')"
if [[ -n "$cache_mb" ]]; then
    reclaim_mb=$(( cache_mb + ${cache_used_mb:-0} ))
    if [[ "$reclaim_mb" -lt 4500 ]]; then
        warn "${cache_mb} MB free + ${cache_used_mb:-0} MB reclaimable where apt stages .debs"
        info "  a CUDA toolkit needs ~4500 MB to download; this filesystem cannot reach that"
        info "  even after 'sudo apt-get clean' — grow $(df -P /var/cache/apt/archives 2>/dev/null | awk 'NR==2{print $6}') before the next toolkit install"
    elif [[ "$cache_mb" -lt 4500 ]]; then
        info "       ${cache_mb} MB free where apt stages .debs, + ${cache_used_mb} MB reclaimable"
        info "       with 'sudo apt-get clean' (safe: pure cache) — enough for the next toolkit"
    fi
fi

# =============================================================================
printf '\n%s\n' "${C_CYN}========================================${C_NC}"
if ((N_FAIL == 0)); then
    printf '%s\n' "${C_GRN}RESULT: PASS${C_NC}  —  ${N_PASS} passed, ${N_WARN} warning(s)"
else
    printf '%s\n' "${C_RED}RESULT: FAIL${C_NC}  —  ${N_PASS} passed, ${N_FAIL} failed, ${N_WARN} warning(s)"
    printf '\n%s\n' "${C_WHT}Failures:${C_NC}"
    for f in "${FAILED[@]}"; do printf '%s\n' "  ${C_RED}·${C_NC} $f"; done
fi
if ((N_WARN > 0)); then
    printf '\n%s\n' "${C_WHT}Warnings:${C_NC}"
    for w in "${WARNED[@]}"; do printf '%s\n' "  ${C_YLW}·${C_NC} $w"; done
fi
printf '%s\n' "${C_CYN}========================================${C_NC}"
exit $(( N_FAIL > 0 ? 1 : 0 ))
