Use nvidia/cuda base images for build

Signed-off-by: Evan Lezar <[email protected]>

deployments/container/Dockerfile.ubi8

@@ -16,7 +16,28 @@ ARG GOLANG_VERSION=1.20.4
 ARG CUDA_IMAGE=cuda
 ARG CUDA_VERSION=11.8.0
 ARG BASE_DIST=ubi8
-FROM golang:${GOLANG_VERSION} as build
+FROM nvidia/cuda:${CUDA_VERSION}-base-${BASE_DIST} as build
+
+RUN yum install -y \
+    wget make git gcc \
+     && \
+    rm -rf /var/cache/yum/*
+
+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 | 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
+
+ENV GOPATH /go
+ENV PATH $GOPATH/bin:/usr/local/go/bin:$PATH
 
 WORKDIR /build
 COPY . .

deployments/container/Dockerfile.ubuntu

@@ -16,7 +16,28 @@ ARG GOLANG_VERSION=1.20.4
 ARG CUDA_IMAGE=cuda
 ARG CUDA_VERSION=11.8.0
 ARG BASE_DIST=ubuntu20.04
-FROM golang:${GOLANG_VERSION} as build
+FROM nvidia/cuda:${CUDA_VERSION}-base-${BASE_DIST} as build
+
+RUN apt-get update && \
+    apt-get install -y wget make git gcc \
+    && \
+    rm -rf /var/lib/apt/lists/*
+
+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 | 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
+
+ENV GOPATH /go
+ENV PATH $GOPATH/bin:/usr/local/go/bin:$PATH
 
 WORKDIR /build
 COPY . .