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 1/6] 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)

From aff1d10951834e048980d223c50d12dd5420a3b4 Mon Sep 17 00:00:00 2001
From: Swati Gupta <swatig@nvidia.com>
Date: Thu, 24 Oct 2024 22:40:46 +0000
Subject: [PATCH 2/6] Add R560

---
 .github/workflows/image.yaml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/image.yaml b/.github/workflows/image.yaml
index 5b65bae6..68fe8eea 100644
--- a/.github/workflows/image.yaml
+++ b/.github/workflows/image.yaml
@@ -36,6 +36,7 @@ jobs:
         driver:
           - 535.216.01
           - 550.127.05
+	  - 560.28.03
         dist:
           - ubuntu20.04
           - ubuntu22.04

From 3f3607778c65b00a86389d455730e2466e1966dc Mon Sep 17 00:00:00 2001
From: Swati Gupta <swatig@nvidia.com>
Date: Thu, 24 Oct 2024 22:48:22 +0000
Subject: [PATCH 3/6] Fix syntax and driver version

---
 .github/workflows/image.yaml       |  2 +-
 ubuntu22.04/Dockerfile             | 11 ++++++-----
 ubuntu22.04/precompiled/Dockerfile | 15 +++++++--------
 3 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/.github/workflows/image.yaml b/.github/workflows/image.yaml
index 68fe8eea..64f2deb9 100644
--- a/.github/workflows/image.yaml
+++ b/.github/workflows/image.yaml
@@ -36,7 +36,7 @@ jobs:
         driver:
           - 535.216.01
           - 550.127.05
-	  - 560.28.03
+          - 560.35.03
         dist:
           - ubuntu20.04
           - ubuntu22.04
diff --git a/ubuntu22.04/Dockerfile b/ubuntu22.04/Dockerfile
index 35476573..2c690fb3 100644
--- a/ubuntu22.04/Dockerfile
+++ b/ubuntu22.04/Dockerfile
@@ -88,11 +88,12 @@ RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$TARGETARCH" != "arm64" ]; then \
 
 # 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
+ENV IMEX_DRIVER_VERSION=560
+RUN if [ "$(printf '%s\n' "$IMEX_DRIVER_VERSION" "$DRIVER_VERSION" | sort -V | head -n1)" = "$IMEX_DRIVER_VERSION" ] && [ "$TARGETARCH" = "arm64" ]; then \
+    curl -fsSL -o /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
+    dpkg -i /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
+    nvidia-imex --version && \
+    rm -rf /tmp/nvidia-imex_${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb; fi
 
 WORKDIR  /drivers
 
diff --git a/ubuntu22.04/precompiled/Dockerfile b/ubuntu22.04/precompiled/Dockerfile
index cf3fba9f..64bc2a83 100644
--- a/ubuntu22.04/precompiled/Dockerfile
+++ b/ubuntu22.04/precompiled/Dockerfile
@@ -59,15 +59,14 @@ 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+
+# 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
-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
+ENV IMEX_DRIVER_VERSION=560
+RUN if [ "$(printf '%s\n' "$IMEX_DRIVER_VERSION" "$DRIVER_VERSION" | sort -V | head -n1)" = "$IMEX_DRIVER_VERSION" ] && [ "$TARGETARCH" = "arm64" ]; then \
+    curl -fsSL -o /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
+    dpkg -i /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
+    nvidia-imex --version && \
+    rm -rf /tmp/nvidia-imex_${IMEX_DRIVER_VERSION}_${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.

From 000e6f832a1a2fc1b8d4c99343b6b4a556298c79 Mon Sep 17 00:00:00 2001
From: Swati Gupta <swatig@nvidia.com>
Date: Mon, 28 Oct 2024 23:35:59 +0000
Subject: [PATCH 4/6] run imex without systemd

---
 ubuntu22.04/nvidia-driver             | 11 ++++++-----
 ubuntu22.04/precompiled/nvidia-driver | 13 +++++++------
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/ubuntu22.04/nvidia-driver b/ubuntu22.04/nvidia-driver
index d48c445e..ad490f9f 100755
--- a/ubuntu22.04/nvidia-driver
+++ b/ubuntu22.04/nvidia-driver
@@ -189,7 +189,7 @@ _create_driver_package() (
 
 # Check if imex node config is present
 _assert_imex_node_config() {
-    [ -f /etc/nvidia-imex/node_config.cfg ]] || return 1
+    [ -f /etc/nvidia-imex/nodes_config.cfg ] || return 1
 }
 
 _assert_nvswitch_system() {
@@ -333,7 +333,7 @@ _load_driver() {
 
     if _assert_imex_node_config; then
 	echo "Starting NVIDIA imex daemon..."
-	systemctl start nvidia-imex.service
+	/usr/bin/nvidia-imex -c /etc/nvidia-imex/config.cfg
     fi
 
     return 0
@@ -394,9 +394,10 @@ _unload_driver() {
     fi
 
     # stop nvidia-imex daemon
-    if [ systemctl is-active --quiet nvidia-imex.service ]; then
-	echo "Stopping NVIDIA imex daemon..."
-	systemctl stop nvidia-imex.service
+    local imex_pid=$(pgrep -f /usr/bin/nvidia-imex)
+    if [ -n "${imex_pid}" ]; then
+        echo "Stopping NVIDIA imex daemon..."
+	kill -SIGTERM "${imex_pid}"
     fi
 
     echo "Unloading NVIDIA driver kernel modules..."
diff --git a/ubuntu22.04/precompiled/nvidia-driver b/ubuntu22.04/precompiled/nvidia-driver
index 5253b550..3734d446 100755
--- a/ubuntu22.04/precompiled/nvidia-driver
+++ b/ubuntu22.04/precompiled/nvidia-driver
@@ -23,7 +23,7 @@ _update_package_cache() {
 }
 
 _assert_imex_node_config() {
-    [ -f /etc/nvidia-imex/node_config.cfg ]] || return 1
+    [ -f /etc/nvidia-imex/nodes_config.cfg ] || return 1
 }
 
 _assert_nvswitch_system() {
@@ -135,7 +135,7 @@ _load_driver() {
 
     if _assert_imex_node_config; then
 	echo "Starting NVIDIA imex daemon..."
-	systemctl start nvidia-imex.service
+        /usr/bin/nvidia-imex -c /etc/nvidia-imex/config.cfg
     fi
 
     return 0
@@ -179,10 +179,11 @@ _unload_driver() {
             return 1
         fi
     fi
-
-    if [ systemctl is-active --quiet nvidia-imex.service ]; then
-	echo "Stopping NVIDIA imex daemon..."
-	systemctl stop nvidia-imex.service
+    # stop nvidia-imex daemon
+    local imex_pid=$(pgrep -f /usr/bin/nvidia-imex)
+    if [ -n "${imex_pid}" ]; then
+        echo "Stopping NVIDIA imex daemon..."
+	kill -SIGTERM "${imex_pid}"
     fi
 
     echo "Unloading NVIDIA driver kernel modules..."

From 2cdc30770b28cab93717de7ded4ff0129f9de14e Mon Sep 17 00:00:00 2001
From: Swati Gupta <swatig@nvidia.com>
Date: Thu, 14 Nov 2024 21:24:11 +0000
Subject: [PATCH 5/6] change nodes_config.cfg path as its conflicting
 installation

---
 ubuntu22.04/Dockerfile                | 1 +
 ubuntu22.04/nvidia-driver             | 5 ++++-
 ubuntu22.04/precompiled/Dockerfile    | 1 +
 ubuntu22.04/precompiled/nvidia-driver | 5 ++++-
 4 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/ubuntu22.04/Dockerfile b/ubuntu22.04/Dockerfile
index 2c690fb3..fd0bd3a5 100644
--- a/ubuntu22.04/Dockerfile
+++ b/ubuntu22.04/Dockerfile
@@ -93,6 +93,7 @@ RUN if [ "$(printf '%s\n' "$IMEX_DRIVER_VERSION" "$DRIVER_VERSION" | sort -V | h
     curl -fsSL -o /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
     dpkg -i /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
     nvidia-imex --version && \
+    ls /etc/nvidia-imex && \
     rm -rf /tmp/nvidia-imex_${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb; fi
 
 WORKDIR  /drivers
diff --git a/ubuntu22.04/nvidia-driver b/ubuntu22.04/nvidia-driver
index ad490f9f..d2916f0f 100755
--- a/ubuntu22.04/nvidia-driver
+++ b/ubuntu22.04/nvidia-driver
@@ -189,7 +189,7 @@ _create_driver_package() (
 
 # Check if imex node config is present
 _assert_imex_node_config() {
-    [ -f /etc/nvidia-imex/nodes_config.cfg ] || return 1
+    [ -f /imex/nodes_config.cfg ] || return 1
 }
 
 _assert_nvswitch_system() {
@@ -332,6 +332,9 @@ _load_driver() {
     fi
 
     if _assert_imex_node_config; then
+	echo "Copying imex nodes_config.cfg to /etc/nvidia-imex"
+	# folder is created during imex installation
+	cp /imex/nodes_config.cfg /etc/nvidia-imex/nodes_config.cfg
 	echo "Starting NVIDIA imex daemon..."
 	/usr/bin/nvidia-imex -c /etc/nvidia-imex/config.cfg
     fi
diff --git a/ubuntu22.04/precompiled/Dockerfile b/ubuntu22.04/precompiled/Dockerfile
index 64bc2a83..8f1897bd 100644
--- a/ubuntu22.04/precompiled/Dockerfile
+++ b/ubuntu22.04/precompiled/Dockerfile
@@ -66,6 +66,7 @@ RUN if [ "$(printf '%s\n' "$IMEX_DRIVER_VERSION" "$DRIVER_VERSION" | sort -V | h
     curl -fsSL -o /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
     dpkg -i /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
     nvidia-imex --version && \
+    ls /etc/nvidia-imex && \
     rm -rf /tmp/nvidia-imex_${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb; fi
 
 # update pkg cache and download pkgs for driver module installation during runtime.
diff --git a/ubuntu22.04/precompiled/nvidia-driver b/ubuntu22.04/precompiled/nvidia-driver
index 3734d446..6166e64b 100755
--- a/ubuntu22.04/precompiled/nvidia-driver
+++ b/ubuntu22.04/precompiled/nvidia-driver
@@ -23,7 +23,7 @@ _update_package_cache() {
 }
 
 _assert_imex_node_config() {
-    [ -f /etc/nvidia-imex/nodes_config.cfg ] || return 1
+    [ -f /imex/nodes_config.cfg ] || return 1
 }
 
 _assert_nvswitch_system() {
@@ -134,6 +134,9 @@ _load_driver() {
     fi
 
     if _assert_imex_node_config; then
+	echo "Copying imex nodes_config.cfg to /etc/nvidia-imex"
+	# folder is created during imex installation
+        cp /imex/nodes_config.cfg /etc/nvidia-imex/nodes_config.cfg
 	echo "Starting NVIDIA imex daemon..."
         /usr/bin/nvidia-imex -c /etc/nvidia-imex/config.cfg
     fi

From ebccc8877881cccfad94c97c96c133f1c99aa3ea Mon Sep 17 00:00:00 2001
From: Swati Gupta <swatig@nvidia.com>
Date: Thu, 14 Nov 2024 22:24:58 +0000
Subject: [PATCH 6/6] update driver version

---
 ubuntu22.04/Dockerfile             | 2 +-
 ubuntu22.04/precompiled/Dockerfile | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/ubuntu22.04/Dockerfile b/ubuntu22.04/Dockerfile
index fd0bd3a5..58eccf3b 100644
--- a/ubuntu22.04/Dockerfile
+++ b/ubuntu22.04/Dockerfile
@@ -88,7 +88,7 @@ RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$TARGETARCH" != "arm64" ]; then \
 
 # 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
-ENV IMEX_DRIVER_VERSION=560
+ENV IMEX_DRIVER_VERSION=565
 RUN if [ "$(printf '%s\n' "$IMEX_DRIVER_VERSION" "$DRIVER_VERSION" | sort -V | head -n1)" = "$IMEX_DRIVER_VERSION" ] && [ "$TARGETARCH" = "arm64" ]; then \
     curl -fsSL -o /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
     dpkg -i /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
diff --git a/ubuntu22.04/precompiled/Dockerfile b/ubuntu22.04/precompiled/Dockerfile
index 8f1897bd..2bbabdc8 100644
--- a/ubuntu22.04/precompiled/Dockerfile
+++ b/ubuntu22.04/precompiled/Dockerfile
@@ -61,7 +61,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends nvidia-driver-$
 
 # 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
-ENV IMEX_DRIVER_VERSION=560
+ENV IMEX_DRIVER_VERSION=565
 RUN if [ "$(printf '%s\n' "$IMEX_DRIVER_VERSION" "$DRIVER_VERSION" | sort -V | head -n1)" = "$IMEX_DRIVER_VERSION" ] && [ "$TARGETARCH" = "arm64" ]; then \
     curl -fsSL -o /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \
     dpkg -i /tmp/nvidia-imex-${IMEX_DRIVER_VERSION}_${DRIVER_VERSION}-1_${TARGETARCH}.deb && \