From eb5270f00a1cd82e070803c74f4a27751d871daa 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     | 36 ++++++++++++++++---
 .../cuda-go-rust-builder.Dockerfile           |  6 ++--
 .../intermediate/go-rust-builder.Dockerfile   |  6 ++--
 3 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/intermediate-docker.yml b/.github/workflows/intermediate-docker.yml
index f158eb4207..82801acbd5 100644
--- a/.github/workflows/intermediate-docker.yml
+++ b/.github/workflows/intermediate-docker.yml
@@ -53,7 +53,7 @@ defaults:
     working-directory: "build/dockerfiles/intermediate"
 
 jobs:
-  build-and-publish-intermediate:
+  setup:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code
@@ -86,14 +86,42 @@ jobs:
           else
             echo "no BASE_IMAGE match"
           fi
-      - name: Build image
-        id: build
+  build-and-publish-intermediate-amd64:
+    runs-on: ubuntu-latest
+    needs: setup
+    steps:
+      - 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 }}
+          env:
+            PLATFORM: arm64
+          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 }}
+
+  build-and-publish-intermediate-arm64:
+    runs-on: ubuntu-latest
+    needs: setup
+    steps:
+      - name: Build arm64 image
+        id: build-arm64
         uses: docker/build-push-action@v5
+        needs: setup
+        env:
+          PLATFORM: arm64
         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..1136fda34f 100644
--- a/build/dockerfiles/intermediate/go-rust-builder.Dockerfile
+++ b/build/dockerfiles/intermediate/go-rust-builder.Dockerfile
@@ -25,7 +25,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}"