From 7f3c3073ec6d5e52c6d0728d2ab94f0aa05a5ff1 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Sun, 28 Jan 2024 12:22:58 +0100
Subject: [PATCH] Switch to CUDA base images

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 deployments/gpu-operator/Dockerfile.ubi8   | 34 ++++++++++++++++------
 deployments/gpu-operator/Dockerfile.ubuntu | 34 ++++++++++++++++------
 2 files changed, 50 insertions(+), 18 deletions(-)

diff --git a/deployments/gpu-operator/Dockerfile.ubi8 b/deployments/gpu-operator/Dockerfile.ubi8
index 6342edab..1c90d124 100644
--- a/deployments/gpu-operator/Dockerfile.ubi8
+++ b/deployments/gpu-operator/Dockerfile.ubi8
@@ -12,18 +12,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+ARG GOLANG_VERSION=1.20.5
+ARG CUDA_IMAGE=cuda
+ARG CUDA_VERSION=12.2.0
 ARG BASE_DIST=ubi8
-ARG CUDA_VERSION
-ARG GOLANG_VERSION=x.x.x
+FROM nvidia/cuda:${CUDA_VERSION}-base-${BASE_DIST} as build
 
-# NOTE: In cases where the libc version is a concern, we would have to use an
-# image based on the target OS to build the golang executables here -- especially
-# if cgo code is included.
-FROM golang:${GOLANG_VERSION} AS build
+RUN yum install -y \
+    wget make git gcc \
+     && \
+    rm -rf /var/cache/yum/*
 
-SHELL ["/bin/bash", "-c"]
+ARG GOLANG_VERSION=x.x.x
+RUN set -eux; \
+    \
+    arch="$(uname -m)"; \
+    case "${arch##*-}" in \
+        x86_64 | amd64) ARCH='amd64' ;; \
+        ppc64el | ppc64le) ARCH='ppc64le' ;; \
+        aarch64) ARCH='arm64' ;; \
+        *) echo "unsupported architecture" ; exit 1 ;; \
+    esac; \
+    wget -nv -O - https://storage.googleapis.com/golang/go${GOLANG_VERSION}.linux-${ARCH}.tar.gz \
+    | tar -C /usr/local -xz
+
+ENV GOPATH /go
+ENV PATH $GOPATH/bin:/usr/local/go/bin:$PATH
 
-ARG TARGETARCH
 WORKDIR /build
 COPY . .
 
@@ -34,7 +49,8 @@ RUN make PREFIX=/artifacts cmds
 
 RUN cp ./deployments/gpu-operator/reconfigure-mig.sh /artifacts/reconfigure-mig.sh
 
-RUN OS_ARCH=${TARGETARCH/x86_64/amd64} && OS_ARCH=${OS_ARCH/aarch64/arm64} && curl -o /artifacts/kubectl -L "https://dl.k8s.io/release/$(curl -L -s https://dl.k8s.io/release/stable.txt)/bin/linux/${OS_ARCH}/kubectl";
+ARG TARGETARCH
+RUN ["/bin/bash", "-c", "OS_ARCH=${TARGETARCH/x86_64/amd64} && OS_ARCH=${OS_ARCH/aarch64/arm64} && curl -o /artifacts/kubectl -L \"https://dl.k8s.io/release/$(curl -L -s https://dl.k8s.io/release/stable.txt)/bin/linux/${OS_ARCH}/kubectl\""]
 RUN chmod a+x /artifacts/kubectl
 
 # Install the nvidia-ctk binary as a go executable
diff --git a/deployments/gpu-operator/Dockerfile.ubuntu b/deployments/gpu-operator/Dockerfile.ubuntu
index 07c11d22..13fd9778 100644
--- a/deployments/gpu-operator/Dockerfile.ubuntu
+++ b/deployments/gpu-operator/Dockerfile.ubuntu
@@ -12,18 +12,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+ARG GOLANG_VERSION=1.20.5
+ARG CUDA_IMAGE=cuda
+ARG CUDA_VERSION=12.2.0
 ARG BASE_DIST=ubuntu20.04
-ARG CUDA_VERSION
-ARG GOLANG_VERSION=x.x.x
+FROM nvidia/cuda:${CUDA_VERSION}-base-${BASE_DIST} as build
 
-# NOTE: In cases where the libc version is a concern, we would have to use an
-# image based on the target OS to build the golang executables here -- especially
-# if cgo code is included.
-FROM golang:${GOLANG_VERSION} AS build
+RUN apt-get update && \
+    apt-get install -y wget make git gcc curl \
+    && \
+    rm -rf /var/lib/apt/lists/*
 
-SHELL ["/bin/bash", "-c"]
+ARG GOLANG_VERSION=x.x.x
+RUN set -eux; \
+    \
+    arch="$(uname -m)"; \
+    case "${arch##*-}" in \
+        x86_64 | amd64) ARCH='amd64' ;; \
+        ppc64el | ppc64le) ARCH='ppc64le' ;; \
+        aarch64) ARCH='arm64' ;; \
+        *) echo "unsupported architecture" ; exit 1 ;; \
+    esac; \
+    wget -nv -O - https://storage.googleapis.com/golang/go${GOLANG_VERSION}.linux-${ARCH}.tar.gz \
+    | tar -C /usr/local -xz
+
+ENV GOPATH /go
+ENV PATH $GOPATH/bin:/usr/local/go/bin:$PATH
 
-ARG TARGETARCH
 WORKDIR /build
 COPY . .
 
@@ -34,7 +49,8 @@ RUN make PREFIX=/artifacts cmds
 
 RUN cp ./deployments/gpu-operator/reconfigure-mig.sh /artifacts/reconfigure-mig.sh
 
-RUN OS_ARCH=${TARGETARCH/x86_64/amd64} && OS_ARCH=${OS_ARCH/aarch64/arm64} && curl -o /artifacts/kubectl -L "https://dl.k8s.io/release/$(curl -L -s https://dl.k8s.io/release/stable.txt)/bin/linux/${OS_ARCH}/kubectl";
+ARG TARGETARCH
+RUN ["/bin/bash", "-c", "OS_ARCH=${TARGETARCH/x86_64/amd64} && OS_ARCH=${OS_ARCH/aarch64/arm64} && curl -o /artifacts/kubectl -L \"https://dl.k8s.io/release/$(curl -L -s https://dl.k8s.io/release/stable.txt)/bin/linux/${OS_ARCH}/kubectl\""]
 RUN chmod a+x /artifacts/kubectl
 
 # Install the nvidia-ctk binary as a go executable