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] 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.