From 7023d7ad150c730d8a95f53372e4a48124e779a4 Mon Sep 17 00:00:00 2001
From: Mergen Nachin <mnachin@meta.com>
Date: Wed, 12 Mar 2025 15:51:58 -0400
Subject: [PATCH] Create Arm64 Graviton instance docker image

Summary:

- Let's see if we can set this up first.

In follow up PRs, here's the plan:

- We will add a few CI jobs to run on this docker image and on Arm64 Graviton instance.
- Will migrate some jobs from macOS runners to this.
---
 .ci/docker/build.sh                 |  4 ++++
 .ci/docker/common/install_conda.sh  |  5 ++++-
 .github/workflows/docker-builds.yml | 22 ++++++++++++++--------
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/.ci/docker/build.sh b/.ci/docker/build.sh
index 16b62b81784..dec80a3da38 100755
--- a/.ci/docker/build.sh
+++ b/.ci/docker/build.sh
@@ -29,6 +29,10 @@ case "${IMAGE_NAME}" in
     LINTRUNNER=""
     CLANG_VERSION=12
     ;;
+  executorch-ubuntu-22.04-gcc11-aarch64)
+    LINTRUNNER=""
+    GCC_VERSION=11
+    ;;
   executorch-ubuntu-22.04-linter)
     LINTRUNNER=yes
     CLANG_VERSION=12
diff --git a/.ci/docker/common/install_conda.sh b/.ci/docker/common/install_conda.sh
index ce51a16452f..ac3e4fc1a0e 100755
--- a/.ci/docker/common/install_conda.sh
+++ b/.ci/docker/common/install_conda.sh
@@ -13,6 +13,9 @@ source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
 install_miniconda() {
   BASE_URL="https://repo.anaconda.com/miniconda"
   CONDA_FILE="Miniconda3-py${PYTHON_VERSION//./}_${MINICONDA_VERSION}-Linux-x86_64.sh"
+  if [[ $(uname -m) == "aarch64" ]]; then 
+    CONDA_FILE="Miniconda3-py${PYTHON_VERSION//./}_${MINICONDA_VERSION}-Linux-aarch64.sh"
+  fi
 
   mkdir -p /opt/conda
   chown ci-user:ci-user /opt/conda
@@ -36,7 +39,7 @@ install_python() {
 
   # From https://github.com/pytorch/pytorch/blob/main/.ci/docker/common/install_conda.sh
   if [[ $(uname -m) == "aarch64" ]]; then
-    conda_install "openblas==0.3.28=*openmp*"
+    conda_install "openblas==0.3.29=*openmp*" -c conda-forge
   else
     conda_install mkl=2022.1.0 mkl-include=2022.1.0
   fi
diff --git a/.github/workflows/docker-builds.yml b/.github/workflows/docker-builds.yml
index 73af9842a20..791a52b96c1 100644
--- a/.github/workflows/docker-builds.yml
+++ b/.github/workflows/docker-builds.yml
@@ -27,19 +27,25 @@ env:
 
 jobs:
   docker-build:
-    runs-on: [self-hosted, linux.2xlarge]
     timeout-minutes: 240
     strategy:
       fail-fast: false
       matrix:
+        runner: [linux.2xlarge]
+        docker-image-name: [
+          executorch-ubuntu-22.04-gcc9,
+          executorch-ubuntu-22.04-clang12,
+          executorch-ubuntu-22.04-linter,
+          executorch-ubuntu-22.04-arm-sdk,
+          executorch-ubuntu-22.04-qnn-sdk,
+          executorch-ubuntu-22.04-mediatek-sdk,
+          executorch-ubuntu-22.04-clang12-android
+          ]
         include:
-          - docker-image-name: executorch-ubuntu-22.04-gcc9
-          - docker-image-name: executorch-ubuntu-22.04-clang12
-          - docker-image-name: executorch-ubuntu-22.04-linter
-          - docker-image-name: executorch-ubuntu-22.04-arm-sdk
-          - docker-image-name: executorch-ubuntu-22.04-qnn-sdk
-          - docker-image-name: executorch-ubuntu-22.04-mediatek-sdk
-          - docker-image-name: executorch-ubuntu-22.04-clang12-android
+          - docker-image-name: executorch-ubuntu-22.04-gcc11-aarch64
+            runner: linux.arm64.2xlarge
+
+    runs-on: [self-hosted, "${{ matrix.runner }}"]
     env:
       DOCKER_IMAGE: 308535385114.dkr.ecr.us-east-1.amazonaws.com/executorch/${{ matrix.docker-image-name }}
     steps: