From 0dec547a656617ab0f7f847d04e6c19cfae2076e Mon Sep 17 00:00:00 2001
From: Benjy Weinberger <benjyw@gmail.com>
Date: Tue, 10 Jan 2023 05:28:05 -0500
Subject: [PATCH] Enable Linux aarch64 CI. (#17918)

---
 .github/workflows/test-cron.yaml              | 152 ++++++++++++
 .github/workflows/test.yaml                   | 230 ++++++++++++++++++
 .../bin/generate_github_workflows.py          |   6 +-
 pants.ci.aarch64.toml                         |   8 +
 4 files changed, 394 insertions(+), 2 deletions(-)
 create mode 100644 pants.ci.aarch64.toml

diff --git a/.github/workflows/test-cron.yaml b/.github/workflows/test-cron.yaml
index 9d94867ab46..b971459db89 100644
--- a/.github/workflows/test-cron.yaml
+++ b/.github/workflows/test-cron.yaml
@@ -7,6 +7,105 @@ env:
   PANTS_CONFIG_FILES: +['pants.ci.toml']
   RUST_BACKTRACE: all
 jobs:
+  bootstrap_pants_linux_arm64:
+    env:
+      PANTS_REMOTE_CACHE_READ: 'false'
+      PANTS_REMOTE_CACHE_WRITE: 'false'
+    if: github.repository_owner == 'pantsbuild'
+    name: Bootstrap Pants, test Rust (Linux-ARM64)
+    runs-on:
+    - self-hosted
+    - Linux
+    - ARM64
+    steps:
+    - name: Check out code
+      uses: actions/checkout@v3
+      with:
+        fetch-depth: 10
+    - name: Tell Pants to use Python ${{ matrix.python-version }}
+      run: 'echo "PY=python${{ matrix.python-version }}" >> $GITHUB_ENV
+
+        echo "PANTS_PYTHON_INTERPRETER_CONSTRAINTS=[''==${{ matrix.python-version
+        }}.*'']" >> $GITHUB_ENV
+
+        '
+    - name: Cache Rust toolchain
+      uses: actions/cache@v3
+      with:
+        key: Linux-ARM64-rustup-${{ hashFiles('rust-toolchain') }}-v2
+        path: '~/.rustup/toolchains/1.66.0-*
+
+          ~/.rustup/update-hashes
+
+          ~/.rustup/settings.toml
+
+          '
+    - name: Cache Cargo
+      uses: benjyw/rust-cache@461b9f8eee66b575bce78977bf649b8b7a8d53f1
+      with:
+        cache-bin: 'false'
+        shared-key: engine
+        workspaces: src/rust/engine
+    - id: get-engine-hash
+      name: Get native engine hash
+      run: echo "hash=$(./build-support/bin/rust/print_engine_hash.sh)" >> $GITHUB_OUTPUT
+      shell: bash
+    - name: Cache native engine
+      uses: actions/cache@v3
+      with:
+        key: Linux-ARM64-engine-${{ steps.get-engine-hash.outputs.hash }}-v1
+        path: '.pants
+
+          src/python/pants/engine/internals/native_engine.so
+
+          src/python/pants/engine/internals/native_engine.so.metadata'
+    - if: github.event_name != 'pull_request'
+      name: Setup toolchain auth
+      run: 'echo TOOLCHAIN_AUTH_TOKEN="${{ secrets.TOOLCHAIN_AUTH_TOKEN }}" >> $GITHUB_ENV
+
+        '
+    - name: Bootstrap Pants
+      run: ./pants version > ${{ runner.temp }}/_pants_version.stdout && [[ -s ${{
+        runner.temp }}/_pants_version.stdout ]]
+    - name: Run smoke tests
+      run: './pants list ::
+
+        ./pants roots
+
+        ./pants help goals
+
+        ./pants help targets
+
+        ./pants help subsystems
+
+        '
+    - continue-on-error: true
+      if: always()
+      name: Upload pants.log
+      uses: actions/upload-artifact@v3
+      with:
+        name: pants-log-bootstrap-Linux-ARM64
+        path: .pants.d/pants.log
+    - name: Upload native binaries
+      uses: actions/upload-artifact@v3
+      with:
+        name: native_binaries.${{ matrix.python-version }}.Linux-ARM64
+        path: '.pants
+
+          src/python/pants/engine/internals/native_engine.so
+
+          src/python/pants/engine/internals/native_engine.so.metadata'
+    - env:
+        TMPDIR: ${{ runner.temp }}
+      if: needs.classify_changes.outputs.rust == 'true'
+      name: Test Rust
+      run: ./cargo test --tests -- --nocapture
+    strategy:
+      matrix:
+        python-version:
+        - '3.8'
+        - '3.9'
+    timeout-minutes: 60
   bootstrap_pants_linux_x86_64:
     env:
       PANTS_REMOTE_CACHE_READ: 'false'
@@ -287,6 +386,59 @@ jobs:
         - '3.8'
         - '3.9'
     timeout-minutes: 30
+  test_python_linux_arm64:
+    env:
+      PANTS_CONFIG_FILES: +['pants.ci.toml','pants.ci.aarch64.toml']
+    if: github.repository_owner == 'pantsbuild'
+    name: Test Python (Linux-ARM64)
+    needs: bootstrap_pants_linux_arm64
+    runs-on:
+    - self-hosted
+    - Linux
+    - ARM64
+    steps:
+    - name: Check out code
+      uses: actions/checkout@v3
+      with:
+        fetch-depth: 10
+    - name: Install AdoptJDK
+      uses: actions/setup-java@v3
+      with:
+        distribution: adopt
+        java-version: '11'
+    - name: Tell Pants to use Python ${{ matrix.python-version }}
+      run: 'echo "PY=python${{ matrix.python-version }}" >> $GITHUB_ENV
+
+        echo "PANTS_PYTHON_INTERPRETER_CONSTRAINTS=[''==${{ matrix.python-version
+        }}.*'']" >> $GITHUB_ENV
+
+        '
+    - name: Download native binaries
+      uses: actions/download-artifact@v3
+      with:
+        name: native_binaries.${{ matrix.python-version }}.Linux-ARM64
+    - if: github.event_name != 'pull_request'
+      name: Setup toolchain auth
+      run: 'echo TOOLCHAIN_AUTH_TOKEN="${{ secrets.TOOLCHAIN_AUTH_TOKEN }}" >> $GITHUB_ENV
+
+        '
+    - name: Run Python tests
+      run: './pants --tag=+platform_specific_behavior test :: -- -m platform_specific_behavior
+
+        '
+    - continue-on-error: true
+      if: always()
+      name: Upload pants.log
+      uses: actions/upload-artifact@v3
+      with:
+        name: pants-log-python-test-Linux-ARM64
+        path: .pants.d/pants.log
+    strategy:
+      matrix:
+        python-version:
+        - '3.8'
+        - '3.9'
+    timeout-minutes: 90
   test_python_linux_x86_64_0:
     env: {}
     if: github.repository_owner == 'pantsbuild'
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index e748e44510a..746356d4310 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -7,6 +7,107 @@ env:
   PANTS_CONFIG_FILES: +['pants.ci.toml']
   RUST_BACKTRACE: all
 jobs:
+  bootstrap_pants_linux_arm64:
+    env:
+      PANTS_REMOTE_CACHE_READ: 'false'
+      PANTS_REMOTE_CACHE_WRITE: 'false'
+    if: (github.repository_owner == 'pantsbuild') && (needs.classify_changes.outputs.docs_only
+      != 'true')
+    name: Bootstrap Pants, test Rust (Linux-ARM64)
+    needs:
+    - classify_changes
+    runs-on:
+    - self-hosted
+    - Linux
+    - ARM64
+    steps:
+    - name: Check out code
+      uses: actions/checkout@v3
+      with:
+        fetch-depth: 10
+    - name: Tell Pants to use Python ${{ matrix.python-version }}
+      run: 'echo "PY=python${{ matrix.python-version }}" >> $GITHUB_ENV
+
+        echo "PANTS_PYTHON_INTERPRETER_CONSTRAINTS=[''==${{ matrix.python-version
+        }}.*'']" >> $GITHUB_ENV
+
+        '
+    - name: Cache Rust toolchain
+      uses: actions/cache@v3
+      with:
+        key: Linux-ARM64-rustup-${{ hashFiles('rust-toolchain') }}-v2
+        path: '~/.rustup/toolchains/1.66.0-*
+
+          ~/.rustup/update-hashes
+
+          ~/.rustup/settings.toml
+
+          '
+    - name: Cache Cargo
+      uses: benjyw/rust-cache@461b9f8eee66b575bce78977bf649b8b7a8d53f1
+      with:
+        cache-bin: 'false'
+        shared-key: engine
+        workspaces: src/rust/engine
+    - id: get-engine-hash
+      name: Get native engine hash
+      run: echo "hash=$(./build-support/bin/rust/print_engine_hash.sh)" >> $GITHUB_OUTPUT
+      shell: bash
+    - name: Cache native engine
+      uses: actions/cache@v3
+      with:
+        key: Linux-ARM64-engine-${{ steps.get-engine-hash.outputs.hash }}-v1
+        path: '.pants
+
+          src/python/pants/engine/internals/native_engine.so
+
+          src/python/pants/engine/internals/native_engine.so.metadata'
+    - if: github.event_name != 'pull_request'
+      name: Setup toolchain auth
+      run: 'echo TOOLCHAIN_AUTH_TOKEN="${{ secrets.TOOLCHAIN_AUTH_TOKEN }}" >> $GITHUB_ENV
+
+        '
+    - name: Bootstrap Pants
+      run: ./pants version > ${{ runner.temp }}/_pants_version.stdout && [[ -s ${{
+        runner.temp }}/_pants_version.stdout ]]
+    - name: Run smoke tests
+      run: './pants list ::
+
+        ./pants roots
+
+        ./pants help goals
+
+        ./pants help targets
+
+        ./pants help subsystems
+
+        '
+    - continue-on-error: true
+      if: always()
+      name: Upload pants.log
+      uses: actions/upload-artifact@v3
+      with:
+        name: pants-log-bootstrap-Linux-ARM64
+        path: .pants.d/pants.log
+    - name: Upload native binaries
+      uses: actions/upload-artifact@v3
+      with:
+        name: native_binaries.${{ matrix.python-version }}.Linux-ARM64
+        path: '.pants
+
+          src/python/pants/engine/internals/native_engine.so
+
+          src/python/pants/engine/internals/native_engine.so.metadata'
+    - env:
+        TMPDIR: ${{ runner.temp }}
+      if: needs.classify_changes.outputs.rust == 'true'
+      name: Test Rust
+      run: ./cargo test --tests -- --nocapture
+    strategy:
+      matrix:
+        python-version:
+        - '3.7'
+    timeout-minutes: 60
   bootstrap_pants_linux_x86_64:
     env:
       PANTS_REMOTE_CACHE_READ: 'false'
@@ -227,6 +328,77 @@ jobs:
         python-version:
         - '3.7'
     timeout-minutes: 60
+  build_wheels_linux_arm64:
+    container:
+      image: quay.io/pypa/manylinux2014_aarch64:latest
+      options: --user 1000:1000
+    env:
+      PANTS_REMOTE_CACHE_READ: 'false'
+      PANTS_REMOTE_CACHE_WRITE: 'false'
+    if: ((github.repository_owner == 'pantsbuild') && (github.event_name == 'push'
+      || needs.classify_changes.outputs.release == 'true')) && (needs.classify_changes.outputs.docs_only
+      != 'true')
+    name: Build wheels (Linux-ARM64)
+    needs:
+    - classify_changes
+    runs-on:
+    - self-hosted
+    - Linux
+    - ARM64
+    steps:
+    - name: Check out code
+      uses: actions/checkout@v3
+      with:
+        fetch-depth: 10
+    - name: Configure Git
+      run: git config --global safe.directory "$GITHUB_WORKSPACE"
+    - name: Install rustup
+      run: 'curl --proto ''=https'' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s --
+        -v -y --default-toolchain none
+
+        echo "${HOME}/.cargo/bin" >> $GITHUB_PATH
+
+        '
+    - name: Expose Pythons
+      run: 'echo "/opt/python/cp37-cp37m/bin" >> $GITHUB_PATH
+
+        echo "/opt/python/cp38-cp38/bin" >> $GITHUB_PATH
+
+        echo "/opt/python/cp39-cp39/bin" >> $GITHUB_PATH
+
+        '
+    - if: github.event_name != 'pull_request'
+      name: Setup toolchain auth
+      run: 'echo TOOLCHAIN_AUTH_TOKEN="${{ secrets.TOOLCHAIN_AUTH_TOKEN }}" >> $GITHUB_ENV
+
+        '
+    - env:
+        PANTS_CONFIG_FILES: +['pants.ci.toml','pants.ci.aarch64.toml']
+      name: Build wheels
+      run: '[[ "${GITHUB_EVENT_NAME}" == "pull_request" ]] && export MODE=debug
+
+        USE_PY39=true ./build-support/bin/release.sh build-local-pex
+
+
+        ./build-support/bin/release.sh build-wheels
+
+        USE_PY38=true ./build-support/bin/release.sh build-wheels
+
+        USE_PY39=true ./build-support/bin/release.sh build-wheels'
+    - continue-on-error: true
+      if: always()
+      name: Upload pants.log
+      uses: actions/upload-artifact@v3
+      with:
+        name: pants-log-wheels-Linux-ARM64
+        path: .pants.d/pants.log
+    - env:
+        AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+        AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+      if: github.event_name == 'push'
+      name: Deploy to S3
+      run: ./build-support/bin/deploy_to_s3.py
+    timeout-minutes: 90
   build_wheels_linux_x86_64:
     container:
       image: quay.io/pypa/manylinux2014_x86_64:latest
@@ -605,8 +777,10 @@ jobs:
     needs:
     - classify_changes
     - check_labels
+    - bootstrap_pants_linux_arm64
     - bootstrap_pants_linux_x86_64
     - bootstrap_pants_macos11_x86_64
+    - build_wheels_linux_arm64
     - build_wheels_linux_x86_64
     - build_wheels_macos10_15_x86_64
     - build_wheels_macos11_arm64
@@ -614,6 +788,7 @@ jobs:
     - check_labels
     - classify_changes
     - lint_python
+    - test_python_linux_arm64
     - test_python_linux_x86_64_0
     - test_python_linux_x86_64_1
     - test_python_linux_x86_64_2
@@ -625,6 +800,61 @@ jobs:
     steps:
     - id: set_merge_ok
       run: echo 'merge_ok=true' >> ${GITHUB_OUTPUT}
+  test_python_linux_arm64:
+    env:
+      PANTS_CONFIG_FILES: +['pants.ci.toml','pants.ci.aarch64.toml']
+    if: (github.repository_owner == 'pantsbuild') && (needs.classify_changes.outputs.docs_only
+      != 'true')
+    name: Test Python (Linux-ARM64)
+    needs:
+    - bootstrap_pants_linux_arm64
+    - classify_changes
+    runs-on:
+    - self-hosted
+    - Linux
+    - ARM64
+    steps:
+    - name: Check out code
+      uses: actions/checkout@v3
+      with:
+        fetch-depth: 10
+    - name: Install AdoptJDK
+      uses: actions/setup-java@v3
+      with:
+        distribution: adopt
+        java-version: '11'
+    - name: Tell Pants to use Python ${{ matrix.python-version }}
+      run: 'echo "PY=python${{ matrix.python-version }}" >> $GITHUB_ENV
+
+        echo "PANTS_PYTHON_INTERPRETER_CONSTRAINTS=[''==${{ matrix.python-version
+        }}.*'']" >> $GITHUB_ENV
+
+        '
+    - name: Download native binaries
+      uses: actions/download-artifact@v3
+      with:
+        name: native_binaries.${{ matrix.python-version }}.Linux-ARM64
+    - if: github.event_name != 'pull_request'
+      name: Setup toolchain auth
+      run: 'echo TOOLCHAIN_AUTH_TOKEN="${{ secrets.TOOLCHAIN_AUTH_TOKEN }}" >> $GITHUB_ENV
+
+        '
+    - name: Run Python tests
+      run: './pants --tag=+platform_specific_behavior test :: -- -m platform_specific_behavior
+
+        '
+    - continue-on-error: true
+      if: always()
+      name: Upload pants.log
+      uses: actions/upload-artifact@v3
+      with:
+        name: pants-log-python-test-Linux-ARM64
+        path: .pants.d/pants.log
+    strategy:
+      matrix:
+        python-version:
+        - '3.7'
+    timeout-minutes: 90
   test_python_linux_x86_64_0:
     env: {}
     if: (github.repository_owner == 'pantsbuild') && (needs.classify_changes.outputs.docs_only
diff --git a/build-support/bin/generate_github_workflows.py b/build-support/bin/generate_github_workflows.py
index 30201f38595..34b4fe0bff7 100644
--- a/build-support/bin/generate_github_workflows.py
+++ b/build-support/bin/generate_github_workflows.py
@@ -312,6 +312,8 @@ def platform_env(self):
             ret["ARCHFLAGS"] = "-arch x86_64"
         if self.platform == Platform.MACOS11_ARM64:
             ret["ARCHFLAGS"] = "-arch arm64"
+        if self.platform == Platform.LINUX_ARM64:
+            ret["PANTS_CONFIG_FILES"] = "+['pants.ci.toml','pants.ci.aarch64.toml']"
         return ret
 
     def wrap_cmd(self, cmd: str) -> str:
@@ -763,7 +765,7 @@ def build_wheels_job(platform: Platform, python_versions: list[str]) -> Jobs:
 def build_wheels_jobs() -> Jobs:
     return {
         **build_wheels_job(Platform.LINUX_X86_64, ALL_PYTHON_VERSIONS),
-        # **build_wheels_job(Platform.LINUX_ARM64, ALL_PYTHON_VERSIONS),
+        **build_wheels_job(Platform.LINUX_ARM64, ALL_PYTHON_VERSIONS),
         **build_wheels_job(Platform.MACOS11_X86_64, ALL_PYTHON_VERSIONS),
         **build_wheels_job(Platform.MACOS10_15_X86_64, ALL_PYTHON_VERSIONS),
         **build_wheels_job(Platform.MACOS11_ARM64, [PYTHON39_VERSION]),
@@ -781,7 +783,7 @@ def test_workflow_jobs(python_versions: list[str], *, cron: bool) -> Jobs:
         },
     }
     jobs.update(**linux_x86_64_test_jobs(python_versions))
-    # jobs.update(**linux_arm64_test_jobs(python_versions))
+    jobs.update(**linux_arm64_test_jobs(python_versions))
     jobs.update(**macos11_x86_64_test_jobs(python_versions))
     if not cron:
         jobs.update(**build_wheels_jobs())
diff --git a/pants.ci.aarch64.toml b/pants.ci.aarch64.toml
new file mode 100644
index 00000000000..4d01f9dda51
--- /dev/null
+++ b/pants.ci.aarch64.toml
@@ -0,0 +1,8 @@
+[GLOBAL]
+# Our physical aarch64 CI instance is a very burly machine with 80 cores.
+# A single GHA self-hosted runner can only run one job at a time, so to use
+# more of the machine's resources, we run multiple self-hosted runners on it
+# and constrain each one here so they don't all try and use all the cores.
+# TODO: Control this externally, on the machine itself?
+rule_threads_core = 8
+process_execution_local_parallelism = 16