From 76faea1bc82943fa5932aee713f86e7a5b59d409 Mon Sep 17 00:00:00 2001
From: Swati Gupta <swatig@nvidia.com>
Date: Thu, 24 Oct 2024 22:18:36 +0000
Subject: [PATCH] Install and start imex

---
 ubuntu22.04/Dockerfile                |  8 ++++++++
 ubuntu22.04/nvidia-driver             | 16 ++++++++++++++++
 ubuntu22.04/precompiled/Dockerfile    | 10 ++++++++++
 ubuntu22.04/precompiled/nvidia-driver | 14 ++++++++++++++
 4 files changed, 48 insertions(+)

diff --git a/ubuntu22.04/Dockerfile b/ubuntu22.04/Dockerfile
index 5ef5041a..35476573 100644
--- a/ubuntu22.04/Dockerfile
+++ b/ubuntu22.04/Dockerfile
@@ -86,6 +86,14 @@ RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$TARGETARCH" != "arm64" ]; then \
     apt-get install -y --no-install-recommends nvidia-fabricmanager-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 \
     libnvidia-nscq-${DRIVER_BRANCH}=${DRIVER_VERSION}-1; fi
 
+# For now install imex only on arm64 and R560+
+# https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/nvidia-imex_560.28.03-1_arm64.deb
+RUN DRIVER_MAJOR_VERSION=$(echo "$DRIVER_VERSION" | cut -d '.' -f 1)
+RUN if [ "$DRIVER_MAJOR_VERSION" -ge 560 ] && [ "$TARGETARCH" = "arm64" ]; then \
+    curl -fsSL -o /tmp/nvidia-imex_$DRIVER_VERSION-1_$TARGETARCH.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/nvidia-imex_560.28.03-1_arm64.deb && \
+    dpkg -i /tmp/nvidia-imex_$DRIVER_VERSION-1_$TARGETARCH.deb && \
+    rm -rf /tmp/nvidia-imex_$DRIVER_VERSION-1_$TARGETARCH.deb; fi
+
 WORKDIR  /drivers
 
 ARG PUBLIC_KEY=empty
diff --git a/ubuntu22.04/nvidia-driver b/ubuntu22.04/nvidia-driver
index aedeeea2..d48c445e 100755
--- a/ubuntu22.04/nvidia-driver
+++ b/ubuntu22.04/nvidia-driver
@@ -187,6 +187,11 @@ _create_driver_package() (
     mv ${pkg_name} precompiled
 )
 
+# Check if imex node config is present
+_assert_imex_node_config() {
+    [ -f /etc/nvidia-imex/node_config.cfg ]] || return 1
+}
+
 _assert_nvswitch_system() {
     [ -d /proc/driver/nvidia-nvswitch/devices ] || return 1
     if [ -z "$(ls -A /proc/driver/nvidia-nvswitch/devices)" ]; then
@@ -326,6 +331,11 @@ _load_driver() {
         nv-fabricmanager -c /usr/share/nvidia/nvswitch/fabricmanager.cfg
     fi
 
+    if _assert_imex_node_config; then
+	echo "Starting NVIDIA imex daemon..."
+	systemctl start nvidia-imex.service
+    fi
+
     return 0
 }
 
@@ -383,6 +393,12 @@ _unload_driver() {
         fi
     fi
 
+    # stop nvidia-imex daemon
+    if [ systemctl is-active --quiet nvidia-imex.service ]; then
+	echo "Stopping NVIDIA imex daemon..."
+	systemctl stop nvidia-imex.service
+    fi
+
     echo "Unloading NVIDIA driver kernel modules..."
     if [ -f /sys/module/nvidia_modeset/refcnt ]; then
         nvidia_modeset_refs=$(< /sys/module/nvidia_modeset/refcnt)
diff --git a/ubuntu22.04/precompiled/Dockerfile b/ubuntu22.04/precompiled/Dockerfile
index f582b73d..cf3fba9f 100644
--- a/ubuntu22.04/precompiled/Dockerfile
+++ b/ubuntu22.04/precompiled/Dockerfile
@@ -59,6 +59,16 @@ RUN apt-get update && apt-get install -y --no-install-recommends nvidia-driver-$
     xserver-xorg-video-nvidia-${DRIVER_BRANCH}-server && \
     rm -rf /var/lib/apt/lists/*;
 
+# For now install imex only with on arm64 and R560+
+# https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/nvidia-imex_560.28.03-1_arm64.deb
+ARG TARGETARCH
+ENV TARGETARCH=$TARGETARCH
+RUN DRIVER_MAJOR_VERSION=$(echo "$DRIVER_VERSION" | cut -d '.' -f 1)
+RUN if [ "$DRIVER_MAJOR_VERSION" -ge 560 ] && [ "$TARGETARCH" = "arm64" ]; then \
+    curl -fsSL -o /tmp/nvidia-imex_$DRIVER_VERSION-1_$TARGETARCH.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/nvidia-imex_$DRIVER_VERSION-1_$TARGETARCH.deb && \
+    dpkg -i /tmp/nvidia-imex_$DRIVER_VERSION-1_$TARGETARCH.deb && \
+    rm -rf /tmp/nvidia-imex_$DRIVER_VERSION-1_$TARGETARCH.deb; fi
+
 # update pkg cache and download pkgs for driver module installation during runtime.
 # this is done to avoid shipping .ko files.
 # avoid cleaning the cache after this to retain these packages during runtime.
diff --git a/ubuntu22.04/precompiled/nvidia-driver b/ubuntu22.04/precompiled/nvidia-driver
index 97369be2..5253b550 100755
--- a/ubuntu22.04/precompiled/nvidia-driver
+++ b/ubuntu22.04/precompiled/nvidia-driver
@@ -22,6 +22,10 @@ _update_package_cache() {
     fi
 }
 
+_assert_imex_node_config() {
+    [ -f /etc/nvidia-imex/node_config.cfg ]] || return 1
+}
+
 _assert_nvswitch_system() {
     [ -d /proc/driver/nvidia-nvswitch/devices ] || return 1
     if [ -z "$(ls -A /proc/driver/nvidia-nvswitch/devices)" ]; then
@@ -129,6 +133,11 @@ _load_driver() {
         nv-fabricmanager -c /usr/share/nvidia/nvswitch/fabricmanager.cfg
     fi
 
+    if _assert_imex_node_config; then
+	echo "Starting NVIDIA imex daemon..."
+	systemctl start nvidia-imex.service
+    fi
+
     return 0
 }
 
@@ -171,6 +180,11 @@ _unload_driver() {
         fi
     fi
 
+    if [ systemctl is-active --quiet nvidia-imex.service ]; then
+	echo "Stopping NVIDIA imex daemon..."
+	systemctl stop nvidia-imex.service
+    fi
+
     echo "Unloading NVIDIA driver kernel modules..."
     if [ -f /sys/module/nvidia_drm/refcnt ]; then
         nvidia_drm_refs=$(< /sys/module/nvidia_drm/refcnt)