From 7fb04878c73b091ce0a16afc93240021a626d942 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Tue, 23 Nov 2021 20:59:52 +0100
Subject: [PATCH] Include all architecture packages in toolkit container

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 .common-ci.yml                    |  5 ++---
 build/container/Dockerfile.centos | 13 +++++++------
 build/container/Dockerfile.ubuntu | 13 +++++++------
 build/container/Makefile          | 13 +++++++++----
 4 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/.common-ci.yml b/.common-ci.yml
index c252332e..bd820c0d 100644
--- a/.common-ci.yml
+++ b/.common-ci.yml
@@ -242,9 +242,8 @@ image-ubuntu18.04:
     - .dist-ubuntu18.04
   needs:
     - package-ubuntu18.04-amd64
-    # TODO: These will be required once we generate multi-arch images
-    # - package-ubuntu18.04-arm64
-    # - package-ubuntu18.04-ppc64le
+    - package-ubuntu18.04-arm64
+    - package-ubuntu18.04-ppc64le
   script:
     - make -f build/container/Makefile build-${DIST}
     - make -f build/container/Makefile push-${DIST}
diff --git a/build/container/Dockerfile.centos b/build/container/Dockerfile.centos
index b47778ec..35e65968 100644
--- a/build/container/Dockerfile.centos
+++ b/build/container/Dockerfile.centos
@@ -46,16 +46,17 @@ ENV NVIDIA_DISABLE_REQUIRE="true"
 ENV NVIDIA_VISIBLE_DEVICES=all
 ENV NVIDIA_DRIVER_CAPABILITIES=utility
 
+ARG ARTIFACTS_DIR
+COPY ${ARTIFACTS_DIR}/ /artifacts/packages/
+
 WORKDIR /artifacts/packages
 
-ARG ARTIFACTS_DIR
-COPY ${ARTIFACTS_DIR}/* /artifacts/packages/
-
 ARG PACKAGE_VERSION
+ARG PACKAGE_ARCH
 RUN yum localinstall -y \
-    libnvidia-container1-${PACKAGE_VERSION}*.rpm \
-    libnvidia-container-tools-${PACKAGE_VERSION}*.rpm \
-    nvidia-container-toolkit-${PACKAGE_VERSION}*.rpm
+    ${PACKAGE_ARCH}/libnvidia-container1-${PACKAGE_VERSION}*.rpm \
+    ${PACKAGE_ARCH}/libnvidia-container-tools-${PACKAGE_VERSION}*.rpm \
+    ${PACKAGE_ARCH}/nvidia-container-toolkit-${PACKAGE_VERSION}*.rpm
 
 WORKDIR /work
 
diff --git a/build/container/Dockerfile.ubuntu b/build/container/Dockerfile.ubuntu
index 6cd3bfa1..a1d610d9 100644
--- a/build/container/Dockerfile.ubuntu
+++ b/build/container/Dockerfile.ubuntu
@@ -52,16 +52,17 @@ ENV NVIDIA_DISABLE_REQUIRE="true"
 ENV NVIDIA_VISIBLE_DEVICES=all
 ENV NVIDIA_DRIVER_CAPABILITIES=utility
 
+ARG ARTIFACTS_DIR
+COPY ${ARTIFACTS_DIR}/ /artifacts/packages/
+
 WORKDIR /artifacts/packages
 
-ARG ARTIFACTS_DIR
-COPY ${ARTIFACTS_DIR}/* /artifacts/packages/
-
 ARG PACKAGE_VERSION
+ARG PACKAGE_ARCH
 RUN dpkg -i \
-    libnvidia-container1_${PACKAGE_VERSION}*.deb \
-    libnvidia-container-tools_${PACKAGE_VERSION}*.deb \
-    nvidia-container-toolkit_${PACKAGE_VERSION}*.deb
+    ${PACKAGE_ARCH}/libnvidia-container1_${PACKAGE_VERSION}*.deb \
+    ${PACKAGE_ARCH}/libnvidia-container-tools_${PACKAGE_VERSION}*.deb \
+    ${PACKAGE_ARCH}/nvidia-container-toolkit_${PACKAGE_VERSION}*.deb
 
 WORKDIR /work
 
diff --git a/build/container/Makefile b/build/container/Makefile
index 4e6fb1ce..38c9064d 100644
--- a/build/container/Makefile
+++ b/build/container/Makefile
@@ -58,22 +58,26 @@ build-%: DOCKERFILE = $(CURDIR)/build/container/Dockerfile.$(DOCKERFILE_SUFFIX)
 
 # Use a generic build target to build the relevant images
 $(BUILD_TARGETS): build-%: $(ARTIFACTS_DIR)
-	$(DOCKER) build --pull \
+	DOCKER_BUILDKIT=1 \
+		$(DOCKER) build --platform=linux/amd64 --pull \
 		--tag $(IMAGE) \
 		--build-arg ARTIFACTS_DIR="$(ARTIFACTS_DIR)" \
 		--build-arg BASE_DIST="$(BASE_DIST)" \
 		--build-arg CUDA_VERSION="$(CUDA_VERSION)" \
 		--build-arg GOLANG_VERSION="$(GOLANG_VERSION)" \
 		--build-arg PACKAGE_VERSION="$(PACKAGE_VERSION)" \
+		--build-arg PACKAGE_ARCH="$(PACKAGE_ARCH)" \
 		--build-arg VERSION="$(VERSION)" \
 		-f $(DOCKERFILE) \
 		$(CURDIR)
 
 
 ARTIFACTS_ROOT ?= $(shell realpath --relative-to=$(CURDIR) $(DIST_DIR))
+ARTIFACTS_DIR = $(ARTIFACTS_ROOT)/$(*)
+
+PACKAGE_ARCH = amd64
 
 build-ubuntu%: DOCKERFILE_SUFFIX := ubuntu
-build-ubuntu%: ARTIFACTS_DIR = $(ARTIFACTS_ROOT)/$(*)/amd64
 build-ubuntu%: PACKAGE_VERSION := $(LIB_VERSION)$(if $(LIB_TAG),~$(LIB_TAG))
 
 build-ubuntu18.04: BASE_DIST := ubuntu18.04
@@ -81,13 +85,14 @@ build-ubuntu20.04: BASE_DIST := ubuntu20.04
 
 build-ubi8: DOCKERFILE_SUFFIX := centos
 # TODO: Update this to use the centos8 packages
-build-ubi8: ARTIFACTS_DIR = $(ARTIFACTS_ROOT)/centos7/x86_64
+build-ubi8: ARTIFACTS_DIR = $(ARTIFACTS_ROOT)/centos7
 build-ubi8: PACKAGE_VERSION := $(LIB_VERSION)-$(if $(LIB_TAG),0.1.$(LIB_TAG),1)
 build-ubi8: BASE_DIST := ubi8
+build-ubi8: PACKAGE_ARCH := x86_64
 
 build-centos%: DOCKERFILE_SUFFIX := centos
-build-centos%: ARTIFACTS_DIR = $(ARTIFACTS_ROOT)/$(*)/x86_64
 build-centos%: PACKAGE_VERSION := $(LIB_VERSION)-$(if $(LIB_TAG),0.1.$(LIB_TAG),1)
+build-centos%: PACKAGE_ARCH := x86_64
 
 build-centos7: BASE_DIST := centos7
 build-centos8: BASE_DIST := centos8