feat: install nvidia-imex on arm64 (Grace-Blackwell GB200/GB300)

GB uses IMEX -- not a node-local fabric manager -- to coordinate cross-node
NVLink (MNNVL / ComputeDomains). The nvidia-imex daemon ships only as a separate
deb in the CUDA repo (not in the driver .run or the fabric-manager redist), so:

- download.sh bundles the version-matched nvidia-imex arm64 deb into the image,
  resolving the exact filename from the repo Packages index (the revision suffix
  varies by version). Skipped on x86 (which uses fabric manager instead).
- install.sh installs it at node boot via dpkg -i --force-depends (the deb
  declares nvidia-modprobe, provided by the runfile driver rather than a deb),
  and leaves the node-wide nvidia-imex.service OFF -- IMEX is orchestrated per
  ComputeDomain by the NVIDIA DRA driver, not run cluster-wide from the host.

The host binary must be present so the DRA driver can CDI-inject it into its
per-workload IMEX daemon pods. Pairs with the fabric-manager arm64 skip: on GB,
skip FM and install IMEX in the same code path where x86 installs FM.

Author: xuexu6666

download.sh

@@ -48,12 +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
 }
 
+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

install.sh

@@ -166,6 +166,18 @@ device_init() {
         bash /opt/gpu/fabricmanager-linux-${ARCH}-${DRIVER_VERSION}/sbin/fm_run_package_installer.sh
     fi
 
+    # 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