From ef2d9d1c54a453d93e32fc7a79e33ca438d41607 Mon Sep 17 00:00:00 2001
From: Christopher Desiniotis <cdesiniotis@nvidia.com>
Date: Mon, 27 Jan 2025 10:29:28 -0800
Subject: [PATCH 1/2] Add nvidia-imex package

Signed-off-by: Christopher Desiniotis <cdesiniotis@nvidia.com>
---
 rhel8/Dockerfile                   | 3 +++
 rhel9/Dockerfile                   | 3 +++
 ubuntu20.04/Dockerfile             | 8 +++++++-
 ubuntu22.04/Dockerfile             | 8 +++++++-
 ubuntu22.04/precompiled/Dockerfile | 5 +++++
 ubuntu24.04/Dockerfile             | 8 +++++++-
 ubuntu24.04/precompiled/Dockerfile | 5 +++++
 7 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/rhel8/Dockerfile b/rhel8/Dockerfile
index a4625d30..24b95e4b 100644
--- a/rhel8/Dockerfile
+++ b/rhel8/Dockerfile
@@ -81,6 +81,9 @@ RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$TARGETARCH" != "arm64" ]; then \
     dnf module enable -y nvidia-driver:${DRIVER_BRANCH}-dkms && \
     dnf install -y ${fmPackage} ${nscqPackage}; fi
 
+RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$DRIVER_BRANCH" -ge "550" ]; then \
+    dnf install -y nvidia-imex-${DRIVER_BRANCH}-${DRIVER_VERSION}-1; fi
+
 COPY nvidia-driver /usr/local/bin
 COPY ocp_dtk_entrypoint /usr/local/bin
 COPY common.sh /usr/local/bin
diff --git a/rhel9/Dockerfile b/rhel9/Dockerfile
index 5f27bc80..0cc92794 100644
--- a/rhel9/Dockerfile
+++ b/rhel9/Dockerfile
@@ -75,6 +75,9 @@ RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$TARGETARCH" != "arm64" ]; then \
     dnf module enable -y nvidia-driver:${DRIVER_BRANCH}-dkms && \
     dnf install -y nvidia-fabric-manager-${DRIVER_VERSION}-1 libnvidia-nscq-${DRIVER_BRANCH}-${DRIVER_VERSION}-1; fi
 
+RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$DRIVER_BRANCH" -ge "550" ]; then \
+    dnf install -y nvidia-imex-${DRIVER_BRANCH}-${DRIVER_VERSION}-1; fi
+
 COPY nvidia-driver /usr/local/bin
 COPY ocp_dtk_entrypoint /usr/local/bin
 COPY common.sh /usr/local/bin
diff --git a/ubuntu20.04/Dockerfile b/ubuntu20.04/Dockerfile
index e65ed401..dcfcb8e2 100644
--- a/ubuntu20.04/Dockerfile
+++ b/ubuntu20.04/Dockerfile
@@ -84,7 +84,13 @@ RUN if [ "$DRIVER_TYPE" != "vgpu" ]; then \
 RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$TARGETARCH" != "arm64" ]; then \
     apt-get update && \
     apt-get install -y --no-install-recommends nvidia-fabricmanager-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 \
-    libnvidia-nscq-${DRIVER_BRANCH}=${DRIVER_VERSION}-1; fi
+    libnvidia-nscq-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 && \
+    rm -rf /var/lib/apt/lists/*; fi
+
+RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$DRIVER_BRANCH" -ge "550" ]; then \
+    apt-get update && \
+    apt-get install -y --no-install-recommends nvidia-imex-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 && \
+    rm -rf /var/lib/apt/lists/*; fi
 
 WORKDIR  /drivers
 
diff --git a/ubuntu22.04/Dockerfile b/ubuntu22.04/Dockerfile
index 979aceef..16052497 100644
--- a/ubuntu22.04/Dockerfile
+++ b/ubuntu22.04/Dockerfile
@@ -84,7 +84,13 @@ RUN if [ "$DRIVER_TYPE" != "vgpu" ]; then \
 RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$TARGETARCH" != "arm64" ]; then \
     apt-get update && \
     apt-get install -y --no-install-recommends nvidia-fabricmanager-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 \
-    libnvidia-nscq-${DRIVER_BRANCH}=${DRIVER_VERSION}-1; fi
+    libnvidia-nscq-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 && \
+    rm -rf /var/lib/apt/lists/*; fi
+
+RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$DRIVER_BRANCH" -ge "550" ]; then \
+    apt-get update && \
+    apt-get install -y --no-install-recommends nvidia-imex-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 && \
+    rm -rf /var/lib/apt/lists/*; fi
 
 WORKDIR  /drivers
 
diff --git a/ubuntu22.04/precompiled/Dockerfile b/ubuntu22.04/precompiled/Dockerfile
index 22b3e881..1d494635 100644
--- a/ubuntu22.04/precompiled/Dockerfile
+++ b/ubuntu22.04/precompiled/Dockerfile
@@ -54,6 +54,11 @@ RUN apt-get update && apt-get install -y --download-only --no-install-recommends
     libnvidia-nscq-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 && \
     rm -rf /var/lib/apt/lists/*;
 
+RUN if [ "$DRIVER_BRANCH" -ge "550" ]; then \
+    apt-get update && \
+    apt-get install -y --no-install-recommends nvidia-imex-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 && \
+    rm -rf /var/lib/apt/lists/*; 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/ubuntu24.04/Dockerfile b/ubuntu24.04/Dockerfile
index f486bf4a..fd183458 100644
--- a/ubuntu24.04/Dockerfile
+++ b/ubuntu24.04/Dockerfile
@@ -81,7 +81,13 @@ RUN if [ "$DRIVER_TYPE" != "vgpu" ]; then \
 RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$TARGETARCH" != "arm64" ]; then \
     apt-get update && \
     apt-get install -y --no-install-recommends nvidia-fabricmanager-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 \
-    libnvidia-nscq-${DRIVER_BRANCH}=${DRIVER_VERSION}-1; fi
+    libnvidia-nscq-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 && \
+    rm -rf /var/lib/apt/lists/*; fi
+
+RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$DRIVER_BRANCH" -ge "550" ]; then \
+    apt-get update && \
+    apt-get install -y --no-install-recommends nvidia-imex-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 && \
+    rm -rf /var/lib/apt/lists/*; fi
 
 WORKDIR  /drivers
 
diff --git a/ubuntu24.04/precompiled/Dockerfile b/ubuntu24.04/precompiled/Dockerfile
index a8b77ba3..df1841bc 100644
--- a/ubuntu24.04/precompiled/Dockerfile
+++ b/ubuntu24.04/precompiled/Dockerfile
@@ -49,6 +49,11 @@ RUN apt-get update && apt-get install -y --download-only --no-install-recommends
     libnvidia-nscq-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 && \
     rm -rf /var/lib/apt/lists/*;
 
+RUN if [ "$DRIVER_BRANCH" -ge "550" ]; then \
+    apt-get update && \
+    apt-get install -y --no-install-recommends nvidia-imex-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 && \
+    rm -rf /var/lib/apt/lists/*; 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.

From acf3325504e72132498b149f1013f583e7ba5994 Mon Sep 17 00:00:00 2001
From: Christopher Desiniotis <cdesiniotis@nvidia.com>
Date: Mon, 27 Jan 2025 23:44:40 -0800
Subject: [PATCH 2/2] [ubuntu20.04] force install nvidia-imex pkg to workaround
 unnecessary dependency issues

The debian packages for Ubuntu 20.04 incorrectly mark openssl >= 3 as a dependency
of the nvidia-imex-$DRIVER_BRANCH package, even though openssl is statically linked.
This commit acts as workaround until the upstream imex packages are fixed.

Signed-off-by: Christopher Desiniotis <cdesiniotis@nvidia.com>
---
 ubuntu20.04/Dockerfile | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/ubuntu20.04/Dockerfile b/ubuntu20.04/Dockerfile
index dcfcb8e2..aba032d1 100644
--- a/ubuntu20.04/Dockerfile
+++ b/ubuntu20.04/Dockerfile
@@ -87,10 +87,16 @@ RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$TARGETARCH" != "arm64" ]; then \
     libnvidia-nscq-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 && \
     rm -rf /var/lib/apt/lists/*; fi
 
+# Force install the imex package to workaround a dependency issue.
+# The debian packages for Ubuntu 20.04 incorrectly mark openssl >= 3
+# as a dependency of the nvidia-imex-$DRIVER_BRANCH package, even
+# though openssl is statically linked.
 RUN if [ "$DRIVER_TYPE" != "vgpu" ] && [ "$DRIVER_BRANCH" -ge "550" ]; then \
-    apt-get update && \
-    apt-get install -y --no-install-recommends nvidia-imex-${DRIVER_BRANCH}=${DRIVER_VERSION}-1 && \
-    rm -rf /var/lib/apt/lists/*; fi
+    OS_ARCH=${TARGETARCH/amd64/x86_64} && OS_ARCH=${OS_ARCH/arm64/sbsa} && \
+    curl -fSsL https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/${OS_ARCH}/nvidia-imex-${DRIVER_BRANCH}_${DRIVER_VERSION}-1_${TARGETARCH}.deb \
+        -o nvidia-imex.deb && \
+    dpkg --force-all -i nvidia-imex.deb && \
+    rm nvidia-imex.deb; fi
 
 WORKDIR  /drivers