From cf84c0ca1599ef2a5e1d1fb4caf5d1268c6bd621 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Fri, 12 Jul 2024 17:11:29 +0200
Subject: [PATCH] TOFIX: packages in image

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 deployments/container/Dockerfile.ubi8 | 92 +++++++++++++++++++++------
 deployments/container/Makefile        | 12 +++-
 2 files changed, 82 insertions(+), 22 deletions(-)

diff --git a/deployments/container/Dockerfile.ubi8 b/deployments/container/Dockerfile.ubi8
index c6b3be14..e0638a46 100644
--- a/deployments/container/Dockerfile.ubi8
+++ b/deployments/container/Dockerfile.ubi8
@@ -47,33 +47,85 @@ ARG VERSION="N/A"
 ARG GIT_COMMIT="unknown"
 RUN make PREFIX=/artifacts cmd-nvidia-ctk-installer
 
-FROM nvcr.io/nvidia/cuda:12.8.0-base-ubi8
+FROM nvcr.io/nvidia/cuda:12.8.0-base-ubi8 AS packaging
+
+ARG ARTIFACTS_ROOT
+COPY ${ARTIFACTS_ROOT} /artifacts/packages/
+
+WORKDIR /artifacts/packages
+
+# build-args are added to the manifest.txt file below.
+ARG PACKAGE_DIST
+ARG PACKAGE_VERSION
+ARG GIT_BRANCH
+ARG GIT_COMMIT
+ARG GIT_COMMIT_SHORT
+ARG SOURCE_DATE_EPOCH
+ARG VERSION
+
+# Create a manifest.txt file with the absolute paths of all deb and rpm packages in the container
+RUN echo "#IMAGE_EPOCH=$(date '+%s')" > /artifacts/manifest.txt && \
+    env | sed 's/^/#/g' >> /artifacts/manifest.txt && \
+    find /artifacts/packages -iname '*.deb' -o -iname '*.rpm' >> /artifacts/manifest.txt
+
+RUN mkdir /licenses && mv /NGC-DL-CONTAINER-LICENSE /licenses/NGC-DL-CONTAINER-LICENSE
+
+# The rpmpackages stage is used to extract the contents of the rpm packages.
+FROM nvcr.io/nvidia/cuda:12.8.0-base-ubi8 AS rpmpackages
+RUN dnf install -y cpio
+
+ARG TARGETARCH
+ARG PACKAGE_DIST_RPM=centos7
+
+COPY --from=packaging /artifacts/packages/${PACKAGE_DIST_RPM} /rpm-packages
+
+RUN mkdir -p /artifacts/rpm
+RUN set -eux; \
+    \
+    case "${TARGETARCH}" in \
+        x86_64 | amd64) ARCH='x86_64' ;; \
+        ppc64el | ppc64le) ARCH='ppc64le' ;; \
+        aarch64 | arm64) ARCH='aarch64' ;; \
+        *) echo "unsupported architecture" ; exit 1 ;; \
+    esac; \
+    for p in $(ls /rpm-packages/${ARCH}/*.rpm); do rpm2cpio $p | cpio -idmv -D /artifacts/rpm; done
+
+# The debpackages stage is used to extract the contents of deb packages.
+FROM nvcr.io/nvidia/cuda:12.8.0-base-ubuntu20.04 AS debpackages
+
+ARG TARGETARCH
+ARG PACKAGE_DIST_DEB=ubuntu18.04
+
+COPY --from=packaging /artifacts/packages/${PACKAGE_DIST_DEB} /deb-packages
+
+RUN mkdir -p /artifacts/deb
+RUN set -eux; \
+    \
+    case "${TARGETARCH}" in \
+        x86_64 | amd64) ARCH='amd64' ;; \
+        ppc64el | ppc64le) ARCH='ppc64le' ;; \
+        aarch64 | arm64) ARCH='arm64' ;; \
+        *) echo "unsupported architecture" ; exit 1 ;; \
+    esac; \
+    for p in $(ls /deb-packages/${ARCH}/*.deb); do dpkg-deb -xv $p /artifacts/deb/; done
+
+FROM nvcr.io/nvidia/cuda:12.8.0-base-ubi8 AS artifacts
+
+COPY --from=rpmpackages /artifacts/rpm /artifacts/rpm
+COPY --from=debpackages /artifacts/deb /artifacts/deb
+COPY --from=build /artifacts/bin /artifacts/build
+
+FROM nvcr.io/nvidia/cuda:12.6.2-base-ubi8
 
 ENV NVIDIA_DISABLE_REQUIRE="true"
 ENV NVIDIA_VISIBLE_DEVICES=void
 ENV NVIDIA_DRIVER_CAPABILITIES=utility
 
-ARG ARTIFACTS_ROOT
-ARG PACKAGE_DIST
-COPY ${ARTIFACTS_ROOT}/${PACKAGE_DIST} /artifacts/packages/${PACKAGE_DIST}
-
-WORKDIR /artifacts/packages
-
-ARG PACKAGE_VERSION
-ARG TARGETARCH
-ENV PACKAGE_ARCH=${TARGETARCH}
-
-RUN PACKAGE_ARCH=${PACKAGE_ARCH/amd64/x86_64} && PACKAGE_ARCH=${PACKAGE_ARCH/arm64/aarch64} && \
-    yum localinstall -y \
-    ${PACKAGE_DIST}/${PACKAGE_ARCH}/libnvidia-container1-1.*.rpm \
-    ${PACKAGE_DIST}/${PACKAGE_ARCH}/libnvidia-container-tools-1.*.rpm \
-    ${PACKAGE_DIST}/${PACKAGE_ARCH}/nvidia-container-toolkit*-${PACKAGE_VERSION}*.rpm
+COPY --from=artifacts /artifacts/rpm /artifacts/rpm
+COPY --from=artifacts /artifacts/deb /artifacts/deb
+COPY --from=artifacts /artifacts/build /work
 
 WORKDIR /work
-
-COPY --from=build /artifacts/nvidia-ctk-installer /work/nvidia-ctk-installer
-RUN ln -s nvidia-ctk-installer nvidia-toolkit
-
 ENV PATH=/work:$PATH
 
 ARG VERSION
diff --git a/deployments/container/Makefile b/deployments/container/Makefile
index a2a1f23b..8c450709 100644
--- a/deployments/container/Makefile
+++ b/deployments/container/Makefile
@@ -90,10 +90,13 @@ $(IMAGE_TARGETS): image-%: $(ARTIFACTS_ROOT)
 		--provenance=false --sbom=false \
 		$(DOCKER_BUILD_OPTIONS) \
 		$(DOCKER_BUILD_PLATFORM_OPTIONS) \
+		$(INTERMEDIATE_TARGET) \
 		--tag $(IMAGE) \
 		--build-arg ARTIFACTS_ROOT="$(ARTIFACTS_ROOT)" \
 		--build-arg GOLANG_VERSION="$(GOLANG_VERSION)" \
 		--build-arg PACKAGE_DIST="$(PACKAGE_DIST)" \
+		--build-arg PACKAGE_DIST_DEB="$(PACKAGE_DIST_DEB)" \
+		--build-arg PACKAGE_DIST_RPM="$(PACKAGE_DIST_RPM)" \
 		--build-arg PACKAGE_VERSION="$(PACKAGE_VERSION)" \
 		--build-arg VERSION="$(VERSION)" \
 		--build-arg GIT_COMMIT="$(GIT_COMMIT)" \
@@ -103,14 +106,19 @@ $(IMAGE_TARGETS): image-%: $(ARTIFACTS_ROOT)
 		-f $(DOCKERFILE) \
 		$(CURDIR)
 
+
+PACKAGE_DIST_DEB = ubuntu18.04
+# TODO: This needs to be set to centos8 for ppc64le builds
+PACKAGE_DIST_RPM = centos7
+
 build-ubuntu%: DOCKERFILE_SUFFIX := ubuntu
 build-ubuntu%: PACKAGE_DIST = ubuntu18.04
 
 build-ubi8: DOCKERFILE_SUFFIX := ubi8
 build-ubi8: PACKAGE_DIST = centos7
 
-build-packaging: DOCKERFILE_SUFFIX := packaging
-build-packaging: PACKAGE_ARCH := amd64
+build-packaging: DOCKERFILE_SUFFIX := ubi8
+build-packaging: INTERMEDIATE_TARGET := --target=packaging
 build-packaging: PACKAGE_DIST = all
 
 # Test targets