add support for arm64 for go image

scroll-tech · Apr 26, 2024 · 1880a3e · 1880a3e
1 parent cc88e8e
commit 1880a3e
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 18 deletions.
diff --git a/.github/workflows/intermediate-docker.yml b/.github/workflows/intermediate-docker.yml
@@ -86,17 +86,35 @@ jobs:
           else
             echo "no BASE_IMAGE match"
           fi
-      - name: Build image
-        id: build
+
+      - 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=amd64
+
+      - 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 }}
             GO_VERSION=${{ github.event.inputs.GO_VERSION }}
             RUST_VERSION=${{ github.event.inputs.RUST_VERSION }}
             PYTHON_VERSION=${{ github.event.inputs.PYTHON_VERSION }}
+            PLATFORM=arm64
diff --git a/build/dockerfiles/intermediate/cuda-go-rust-builder.Dockerfile b/build/dockerfiles/intermediate/cuda-go-rust-builder.Dockerfile
@@ -1,9 +1,11 @@
-ARG CUDA_VERSION=11.7.1
-ARG GO_VERSION=1.21
-ARG RUST_VERSION=nightly-2023-12-03
-ARG CARGO_CHEF_TAG=0.1.41
-
 FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
+
+ARG CUDA_VERSION
+ARG GO_VERSION
+ARG RUST_VERSION
+ARG CARGO_CHEF_TAG
+ARG PLATFORM
+
 RUN apt-get update
 # Install basic packages
 RUN apt-get install build-essential curl wget git pkg-config --no-install-recommends -y
@@ -29,7 +31,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
@@ -1,9 +1,10 @@
-ARG GO_VERSION=1.21
-ARG RUST_VERSION=nightly-2023-12-03
-ARG CARGO_CHEF_TAG=0.1.41
-
 FROM ubuntu:20.04
 
+ARG GO_VERSION
+ARG RUST_VERSION
+ARG CARGO_CHEF_TAG
+ARG PLATFORM
+
 RUN apt-get update && ln -fs /usr/share/zoneinfo/America/New_York /etc/localtime
 
 # Install basic packages
@@ -25,7 +26,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}"