From f3b57dcd6ff9c3ef94f68a8096dcfd53cfa56bc0 Mon Sep 17 00:00:00 2001
From: Tariq Ibrahim <tibrahim@nvidia.com>
Date: Thu, 14 Mar 2024 11:27:08 -0700
Subject: [PATCH] run go build in the cuda base image rather golang

Signed-off-by: Tariq Ibrahim <tibrahim@nvidia.com>
---
 deployments/container/Dockerfile.ubi8 | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/deployments/container/Dockerfile.ubi8 b/deployments/container/Dockerfile.ubi8
index fc5e3ebd..352bbc4e 100644
--- a/deployments/container/Dockerfile.ubi8
+++ b/deployments/container/Dockerfile.ubi8
@@ -14,9 +14,8 @@
 
 ARG CUDA_VERSION
 ARG BASE_DIST
-ARG GOLANG_VERSION=x.x.x
 
-FROM golang:${GOLANG_VERSION} AS build
+FROM nvcr.io/nvidia/cuda:${CUDA_VERSION}-base-${BASE_DIST}
 
 WORKDIR /work
 
@@ -25,9 +24,20 @@ COPY go.sum go.sum
 COPY vendor vendor
 COPY cmd/nvdrain cmd/nvdrain
 
-RUN GOOS=linux go build -o nvdrain ./cmd/nvdrain
+ARG GOLANG_VERSION=0.0.0
+RUN set -eux; \
+    \
+    arch="$(uname -m)"; \
+    case "${arch##*-}" in \
+        x86_64 | amd64) ARCH='amd64' ;; \
+        ppc64el | ppc64le) ARCH='ppc64le' ;; \
+        aarch64 | arm64) 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
 
-FROM nvidia/cuda:${CUDA_VERSION}-base-${BASE_DIST}
+RUN GOOS=linux go build -o nvdrain ./cmd/nvdrain
 
 ARG TARGETARCH
 
@@ -39,7 +49,7 @@ RUN chmod a+x /usr/bin/kubectl
 
 COPY driver-manager /usr/local/bin
 COPY scripts/vfio-manage /usr/local/bin
-COPY --from=build /work/nvdrain /usr/local/bin
+COPY /work/nvdrain /usr/local/bin
 
 LABEL io.k8s.display-name="NVIDIA Driver Upgrade Manager for Kubernetes"
 LABEL name="NVIDIA Driver Upgrade Manager for Kubernetes"