diff --git a/download.sh b/download.sh index b8e1bb0..61c80ee 100755 --- a/download.sh +++ b/download.sh @@ -48,11 +48,40 @@ install_fabric_manager () { mv /opt/gpu/fm_run_package_installer.sh /opt/gpu/fabricmanager-linux-${NVIDIA_FM_ARCH}-${DRIVER_VERSION}/sbin/fm_run_package_installer.sh } -if [[ "${DRIVER_KIND}" == "cuda" ]]; then - # download fabricmanager for nvlink based systems, e.g. multi instance gpu vms. +install_imex () { + # nvidia-imex is the cross-node NVLink (MNNVL / ComputeDomains) coordinator for + # Grace-Blackwell. It is NOT in the driver .run or the fabric-manager redist -- it + # ships only as a separate deb in the CUDA repo. Bundle the version-matched deb into + # the image; it is installed at node boot (see install.sh device_init). The exact deb + # revision suffix (e.g. -1ubuntu1) varies by version, so resolve the filename from the + # repo Packages index rather than hard-coding it. + local repo="https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${VERSION_ID//./}/sbsa" + local deb + # NB: awk must read to EOF (no early `exit`) -- exiting mid-stream closes the pipe and + # SIGPIPEs curl, which `set -o pipefail` would turn into a build failure. Gate on a flag + # to keep only the first match instead. + deb="$(curl -fsSL "${repo}/Packages" \ + | awk -v p="nvidia-imex_${DRIVER_VERSION}-" '$1=="Filename:" && index($2,p) && !f{print $2; f=1}')" + if [[ -z "${deb}" ]]; then + echo "nvidia-imex ${DRIVER_VERSION} not found in ${repo}" + exit 1 + fi + curl -fsSLO "${repo}/${deb#./}" + mv "$(basename "${deb}")" /opt/gpu/ +} + +# download fabricmanager for nvlink based systems, but skip it on arm64: +# arm64 = Grace-Blackwell (GB200/GB300), which uses IMEX, not a node-local FM. +if [[ "${DRIVER_KIND}" == "cuda" && "${TARGETARCH}" != "arm64" ]]; then install_fabric_manager fi +# download nvidia-imex for nvlink based arm64 systems (Grace-Blackwell GB200/GB300): +# GB uses IMEX (not a node-local fabric manager) to coordinate cross-node NVLink. +if [[ "${DRIVER_KIND}" == "cuda" && "${TARGETARCH}" == "arm64" ]]; then + install_imex +fi + # configure nvidia apt repo to cache packages curl -fsSLO https://nvidia.github.io/libnvidia-container/gpgkey diff --git a/install.sh b/install.sh index adc8fcb..a616c5a 100644 --- a/install.sh +++ b/install.sh @@ -166,14 +166,11 @@ initialize_nvidia_driver() { ldconfig nvidia-smi - # install fabricmanager for nvlink based systems - if [[ "${DRIVER_KIND}" == "cuda" ]]; then - NVIDIA_FM_ARCH=$ARCH - if [ "$NVIDIA_FM_ARCH" = "arm64" ]; then - # NVIDIA uses the name "SBSA" for ARM64 platforms for the fabric manager. See https://en.wikipedia.org/wiki/Server_Base_System_Architecture - NVIDIA_FM_ARCH="sbsa" - fi - bash /opt/gpu/fabricmanager-linux-${NVIDIA_FM_ARCH}-${DRIVER_VERSION}/sbin/fm_run_package_installer.sh + # install fabricmanager for nvlink based systems, but skip it on arm64: + # arm64 = Grace-Blackwell (GB200/GB300), which uses IMEX, not a node-local FM. + # (uname -m reports "aarch64" for arm64.) + if [[ "${DRIVER_KIND}" == "cuda" && "${ARCH}" != "aarch64" ]]; then + bash /opt/gpu/fabricmanager-linux-${ARCH}-${DRIVER_VERSION}/sbin/fm_run_package_installer.sh fi } @@ -182,6 +179,18 @@ configure_nvidia_container_runtime() { install_nvidia_container_toolkit + # install nvidia-imex on arm64 (Grace-Blackwell): the cross-node NVLink (MNNVL) + # coordinator that GB uses in place of a node-local fabric manager. The binary must be + # present on the host so the NVIDIA DRA driver (ComputeDomains) can inject it into its + # per-workload IMEX daemon pods. --force-depends: the deb declares nvidia-modprobe, + # which is provided by the runfile driver (present on disk) rather than as a deb. + # The node-wide nvidia-imex.service stays OFF -- IMEX is orchestrated per ComputeDomain + # by DRA, not run cluster-wide from the host. + if [[ "${DRIVER_KIND}" == "cuda" && "${ARCH}" == "aarch64" ]]; then + dpkg -i --force-depends /opt/gpu/nvidia-imex_*_arm64.deb + systemctl disable --now nvidia-imex.service 2>/dev/null || true + fi + mkdir -p /etc/containerd/config.d cp /opt/gpu/10-nvidia-runtime.toml /etc/containerd/config.d/10-nvidia-runtime.toml diff --git a/test/install.bats b/test/install.bats index 870f3a7..145637a 100644 --- a/test/install.bats +++ b/test/install.bats @@ -275,3 +275,82 @@ EOF [ "$status" -ne 0 ] } + +# --- arm64 (Grace-Blackwell GB200/GB300): fabric-manager skip + nvidia-imex --- +# GB uses cross-node IMEX in place of a node-local fabric manager. initialize_nvidia_driver +# must skip FM on aarch64 (installing the nonexistent aarch64 FM tarball previously bricked +# node join), and configure_nvidia_container_runtime must instead stage the nvidia-imex deb +# while leaving the node-wide service OFF (DRA orchestrates IMEX per ComputeDomain). + +@test "initialize_nvidia_driver installs the fabric manager on x86_64 cuda" { + DRIVER_KIND="cuda"; ARCH="x86_64"; DRIVER_VERSION="580.0.0" + nvidia-modprobe() { :; } + cp() { :; } + ldconfig() { :; } + nvidia-smi() { :; } + # the FM installer is invoked as `bash /opt/gpu/fabricmanager-.../fm_run_package_installer.sh` + bash() { echo "bash $*" >> "${TEST_TMP}/fm-calls"; } + + run initialize_nvidia_driver + + [ "$status" -eq 0 ] + run cat "${TEST_TMP}/fm-calls" + [ "$status" -eq 0 ] + [[ "$output" == *"/opt/gpu/fabricmanager-linux-x86_64-580.0.0/sbin/fm_run_package_installer.sh"* ]] +} + +@test "initialize_nvidia_driver skips the fabric manager on arm64 (aarch64) cuda" { + DRIVER_KIND="cuda"; ARCH="aarch64"; DRIVER_VERSION="580.0.0" + nvidia-modprobe() { :; } + cp() { :; } + ldconfig() { :; } + nvidia-smi() { :; } + bash() { echo "bash $*" >> "${TEST_TMP}/fm-calls"; } + + run initialize_nvidia_driver + + [ "$status" -eq 0 ] + # FM installer must never run on aarch64 (GB has no node-local fabric manager) + [ ! -f "${TEST_TMP}/fm-calls" ] +} + +@test "configure_nvidia_container_runtime installs nvidia-imex on arm64 (aarch64) and leaves the service off" { + DRIVER_KIND="cuda"; ARCH="aarch64" + install_nvidia_container_toolkit() { :; } + mkdir() { :; } + cp() { :; } + dirname() { command dirname "$@"; } + dpkg() { echo "dpkg $*" >> "${TEST_TMP}/imex-calls"; } + systemctl() { echo "systemctl $*" >> "${TEST_TMP}/imex-calls"; } + export AKSGPU_NVIDIA_CTK_BIN="${TEST_TMP}/bin/nvidia-ctk" + _stub_bin nvidia-ctk 0 + + run configure_nvidia_container_runtime + + [ "$status" -eq 0 ] + run cat "${TEST_TMP}/imex-calls" + [ "$status" -eq 0 ] + [[ "$output" == *"dpkg -i --force-depends /opt/gpu/nvidia-imex_"*"_arm64.deb"* ]] + # node-wide IMEX service is disabled+stopped; DRA runs it per ComputeDomain + [[ "$output" == *"systemctl disable --now nvidia-imex.service"* ]] +} + +@test "configure_nvidia_container_runtime does not install nvidia-imex on x86_64" { + DRIVER_KIND="cuda"; ARCH="x86_64" + install_nvidia_container_toolkit() { :; } + mkdir() { :; } + cp() { :; } + dirname() { command dirname "$@"; } + dpkg() { echo "dpkg $*" >> "${TEST_TMP}/imex-calls"; } + systemctl() { echo "systemctl $*" >> "${TEST_TMP}/imex-calls"; } + export AKSGPU_NVIDIA_CTK_BIN="${TEST_TMP}/bin/nvidia-ctk" + _stub_bin nvidia-ctk 0 + + run configure_nvidia_container_runtime + + [ "$status" -eq 0 ] + run cat "${TEST_TMP}/imex-calls" + # imex must never be installed on x86_64 (fabric manager handles NVLink there) + [[ "$output" != *"dpkg"* ]] + [[ "$output" != *"nvidia-imex.service"* ]] +}