From 63bad5ba6ebada9f7727418d875bd098dc75f88e Mon Sep 17 00:00:00 2001
From: "sebastien.baizet" <sebastien.baizet@gmail.com>
Date: Fri, 26 Apr 2024 14:35:44 +0200
Subject: [PATCH] add support for arm64 for go image

---
 .github/workflows/intermediate-docker.yml     | 32 +++++++++++++++++--
 .../cuda-go-rust-builder.Dockerfile           |  6 ++--
 .../intermediate/go-rust-builder.Dockerfile   | 11 ++++---
 3 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/intermediate-docker.yml b/.github/workflows/intermediate-docker.yml
index f158eb4207..49608591c4 100644
--- a/.github/workflows/intermediate-docker.yml
+++ b/.github/workflows/intermediate-docker.yml
@@ -86,14 +86,40 @@ jobs:
           else
             echo "no BASE_IMAGE match"
           fi
-      - name: Build image
-        id: build
+
+      - name: Set env PLATFORM to amd64
+        run: echo PLATFORM="amd64" >> $GITHUB_ENV
+
+      - name: test
+        run: echo ${{ env.PLATFORM }}
+
+      - name: Build amd64 image
+        id: build-amd64
+        uses: docker/build-push-action@v5
+        with:
+          context: .
+          push: true
+          file: build/dockerfiles/intermediate/${{ github.event.inputs.BASE_IMAGE }}.Dockerfile
+          platforms: linux/amd64
+          tags: scrolltech/${{ github.event.inputs.BASE_IMAGE }}:${{ env.TAG }}
+          build-args: |
+            CUDA_VERSION=${{ github.event.inputs.CUDA_VERSION }}
+            GO_VERSION=${{ github.event.inputs.GO_VERSION }}
+            RUST_VERSION=${{ github.event.inputs.RUST_VERSION }}
+            PYTHON_VERSION=${{ github.event.inputs.PYTHON_VERSION }}
+            PLATFORM=${{ env.PLATFORM }}
+
+      - name: Set env PLATFORM to arm64
+        run: PLATFORM="arm64"
+
+      - name: Build arm64 image
+        id: build-arm64
         uses: docker/build-push-action@v5
         with:
           context: .
           push: true
           file: build/dockerfiles/intermediate/${{ github.event.inputs.BASE_IMAGE }}.Dockerfile
-          platforms: linux/amd64,linux/arm64
+          platforms: linux/arm64
           tags: scrolltech/${{ github.event.inputs.BASE_IMAGE }}:${{ env.TAG }}
           build-args: |
             CUDA_VERSION=${{ github.event.inputs.CUDA_VERSION }}
diff --git a/build/dockerfiles/intermediate/cuda-go-rust-builder.Dockerfile b/build/dockerfiles/intermediate/cuda-go-rust-builder.Dockerfile
index 757c7c06a9..e018ae5c55 100644
--- a/build/dockerfiles/intermediate/cuda-go-rust-builder.Dockerfile
+++ b/build/dockerfiles/intermediate/cuda-go-rust-builder.Dockerfile
@@ -29,7 +29,7 @@ RUN cargo install cargo-chef --locked --version ${CARGO_CHEF_TAG} \
 # Install Go
 ARG GO_VERSION
 RUN rm -rf /usr/local/go
-RUN wget https://go.dev/dl/go${GO_VERSION}.1.linux-amd64.tar.gz
-RUN tar -C /usr/local -xzf go${GO_VERSION}.1.linux-amd64.tar.gz
-RUN rm go${GO_VERSION}.1.linux-amd64.tar.gz
+RUN wget https://go.dev/dl/go${GO_VERSION}.1.linux-${PLATFORM}.tar.gz
+RUN tar -C /usr/local -xzf go${GO_VERSION}.1.linux-${PLATFORM}.tar.gz
+RUN rm go${GO_VERSION}.1.linux-${PLATFORM}.tar.gz
 ENV PATH="/usr/local/go/bin:${PATH}"
diff --git a/build/dockerfiles/intermediate/go-rust-builder.Dockerfile b/build/dockerfiles/intermediate/go-rust-builder.Dockerfile
index 1881a57029..9b3dd2b00c 100644
--- a/build/dockerfiles/intermediate/go-rust-builder.Dockerfile
+++ b/build/dockerfiles/intermediate/go-rust-builder.Dockerfile
@@ -1,8 +1,11 @@
+FROM ubuntu:20.04
+
 ARG GO_VERSION=1.21
 ARG RUST_VERSION=nightly-2023-12-03
 ARG CARGO_CHEF_TAG=0.1.41
+ARG PLATFORM
 
-FROM ubuntu:20.04
+RUN echo "platform = ${PLATFORM}"
 
 RUN apt-get update && ln -fs /usr/share/zoneinfo/America/New_York /etc/localtime
 
@@ -25,7 +28,7 @@ RUN cargo install cargo-chef --locked --version ${CARGO_CHEF_TAG} \
 # Install Go
 ARG GO_VERSION
 RUN rm -rf /usr/local/go
-RUN wget https://go.dev/dl/go${GO_VERSION}.1.linux-amd64.tar.gz
-RUN tar -C /usr/local -xzf go${GO_VERSION}.1.linux-amd64.tar.gz
-RUN rm go${GO_VERSION}.1.linux-amd64.tar.gz
+RUN wget https://go.dev/dl/go${GO_VERSION}.1.linux-${PLATFORM}.tar.gz
+RUN tar -C /usr/local -xzf go${GO_VERSION}.1.linux-${PLATFORM}.tar.gz
+RUN rm go${GO_VERSION}.1.linux-${PLATFORM}.tar.gz
 ENV PATH="/usr/local/go/bin:${PATH}"