Install and start imex

Author: guptaNswati

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

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)

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.

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)