[FEAT] extend build-commit workflow to support different compile-ar…

…chs (#3459)

# Overview
This PR extends the build commit workflow to also support building for
ARM architectures. The current workflow only built for x86.

## Limitations
- We are building using the `buildjet` service.
- Buildjet provides ARM machines. However, their ARM machines only run
Ubuntu 22.04.
  - Ubuntu 22.04 runs glibc2.34.
  - glibc is forwards-compatible, but *not* backwards compatible.
- This means that whenever you are running the built wheel, you must be
running it on a machine which contains glibc2.34 or greater. Installs of
the wheel will fail on platforms which contain older versions of glibc.
(Note that most AWS AMIs contain older versions of glibc).

## Notes
The storage structure of the wheels in AWS S3 is as below:
```txt
github-actions-artifacts-bucket
|- builds
   |- ${{ GITHUB.SHA }}
      |- ${{ WHEEL_NAME_FOR_X86 }}
      |- ${{ WHEEL_NAME_FOR_ARM }}
```

.github/ci-scripts/constants.py

@@ -0,0 +1,2 @@
+BUCKET_NAME = "github-actions-artifacts-bucket"
+WHEEL_SUFFIX = ".whl"

.github/ci-scripts/get_wheel_name_from_s3.py

@@ -0,0 +1,44 @@
+"""
+Given a commit hash and a "platform substring", prints the wheelname of the wheel (if one exists) to stdout.
+
+# Example
+
+```bash
+COMMIT_HASH="abcdef0123456789"
+PLATFORM_SUBSTRING="x86"
+WHEELNAME=$(python get_wheel_name_from_s3.py $COMMIT_HASH $PLATFORM_SUBSTRING)
+
+echo $WHEELNAME
+# Will echo the wheelname if a wheel exists that matches the platform substring.
+# Otherwise, will echo nothing.
+```
+"""
+
+import sys
+from pathlib import Path
+
+import boto3
+import wheellib
+
+if __name__ == "__main__":
+    commit_hash = sys.argv[1]
+    platform_substring = sys.argv[2]
+
+    s3 = boto3.client("s3")
+    response = s3.list_objects_v2(Bucket="github-actions-artifacts-bucket", Prefix=f"builds/{commit_hash}/")
+    matches = []
+    for content in response.get("Contents", []):
+        wheelname = Path(content["Key"]).name
+        platform_tag = wheellib.get_platform_tag(wheelname)
+        if platform_substring in platform_tag:
+            matches.append(wheelname)
+
+    if len(matches) > 1:
+        raise Exception(
+            f"Multiple wheels found that match the given platform substring: {platform_substring}; expected just 1"
+        )
+
+    try:
+        print(next(iter(matches)))
+    except StopIteration:
+        pass

.github/ci-scripts/upload_wheel_to_s3.py

@@ -0,0 +1,45 @@
+import argparse
+from pathlib import Path
+
+import boto3
+import constants
+import wheellib
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--commit-hash", required=True)
+    parser.add_argument("--platform-substring", required=True, choices=["x86", "aarch", "arm"])
+    parser.add_argument("--path-to-wheel-dir", required=True)
+    args = parser.parse_args()
+
+    commit_hash = args.commit_hash
+    platform_substring = args.platform_substring
+    path_to_wheel_dir = Path(args.path_to_wheel_dir)
+
+    assert path_to_wheel_dir.exists(), f"Path to wheel directory does not exist: {path_to_wheel_dir}"
+    wheelpaths = iter(filepath for filepath in path_to_wheel_dir.iterdir() if filepath.suffix == constants.WHEEL_SUFFIX)
+
+    def f(wheelpath: Path) -> bool:
+        platform_tag = wheellib.get_platform_tag(wheelpath.name)
+        return platform_substring in platform_tag
+
+    filtered_wheelpaths: list[Path] = list(filter(f, wheelpaths))
+
+    length = len(filtered_wheelpaths)
+    if length == 0:
+        raise RuntimeError(f"No wheels found that match the given platform substring: {platform_substring}; expected 1")
+    elif length > 1:
+        raise RuntimeError(
+            f"""Multiple wheels found that match the given platform substring: {platform_substring}; expected just 1
+Wheels available: {wheelpaths}"""
+        )
+    [wheelpath] = filtered_wheelpaths
+    s3 = boto3.client("s3")
+    destination = Path("builds") / commit_hash / wheelpath.name
+    s3.upload_file(
+        Filename=wheelpath,
+        Bucket=constants.BUCKET_NAME,
+        Key=str(destination),
+        ExtraArgs={"ACL": "public-read"},
+    )
+    print(wheelpath.name)

.github/ci-scripts/wheellib.py

@@ -0,0 +1,7 @@
+from packaging.utils import parse_wheel_filename
+
+
+def get_platform_tag(wheelname: str) -> str:
+    distribution, version, build_tag, tags = parse_wheel_filename(wheelname)
+    assert len(tags) == 1, "Multiple tags found"
+    return next(iter(tags)).platform

.github/workflows/build-commit.yaml

@@ -2,26 +2,38 @@ name: build-commit
 
 on:
   workflow_dispatch:
-  workflow_call:
-    secrets:
-      ACTIONS_AWS_ROLE_ARN:
-        description: The ARN of the AWS role to assume
+    inputs:
+      arch:
+        type: choice
+        options:
+        - x86
+        - arm
+        description: The machine architecture to build for
         required: true
-    outputs:
-      wheel:
-        description: The wheel file that was built
-        value: ${{ jobs.build-commit.outputs.wheel }}
+        default: x86
+      python_version:
+        type: string
+        description: The version of python to use
+        required: false
+        default: "3.9"
 
 jobs:
   build-commit:
-    runs-on: buildjet-8vcpu-ubuntu-2004
-    timeout-minutes: 15 # Remove for ssh debugging
+    runs-on: ${{ inputs.arch == 'x86' && 'buildjet-8vcpu-ubuntu-2004' || 'buildjet-8vcpu-ubuntu-2204-arm' }}
+    timeout-minutes: 30 # Remove for ssh debugging
     permissions:
       id-token: write
       contents: read
-    outputs:
-      wheel: ${{ steps.upload.outputs.wheel }}
     steps:
+    - name: Set platform substring
+      run: |
+        if [ "${{ inputs.arch }}" == "x86" ]; then
+          platform_substring=x86
+        else
+          platform_substring=aarch
+        fi
+        echo "platform_substring=$platform_substring" >> $GITHUB_ENV
+        echo "Running on $platform_substring build machines"
     - name: Checkout repo
       uses: actions/checkout@v4
       with:
@@ -32,59 +44,57 @@ jobs:
         aws-region: us-west-2
         role-session-name: build-commit-workflow
         role-to-assume: ${{ secrets.ACTIONS_AWS_ROLE_ARN }}
-    - name: Install rust + uv + python
+    - name: Install uv, rust, python
       uses: ./.github/actions/install
+      with:
+        python_version: ${{ inputs.python_version }}
     - name: Restore cached build artifacts
       uses: buildjet/cache@v4
       with:
-        path: ~/target
-        key: ${{ runner.os }}-cargo-deps-${{ hashFiles('**/Cargo.lock') }}
-        restore-keys: ${{ runner.os }}-cargo-deps-
+        path: ~/target/release
+        key: ${{ runner.os }}-${{ inputs.arch == 'x86' && 'x86' || 'arm' }}-cargo-deps-${{ hashFiles('**/Cargo.lock') }}
+        restore-keys: ${{ runner.os }}-${{ inputs.arch == 'x86' && 'x86' || 'arm' }}-cargo-deps-
+    - name: Setup uv environment
+      run: |
+        uv v
+        source .venv/bin/activate
+        uv pip install boto3 packaging
     - name: Check if build already exists in AWS S3
       run: |
-        RESULT=$(aws s3 ls s3://github-actions-artifacts-bucket/builds/${{ github.sha }}/ | wc -l)
-        if [ "$RESULT" -gt 1 ]; then
-            echo "Error: More than one artifact found. Failing the job."
-            exit 1
-        elif [ "$RESULT" -eq 0 ]; then
-            echo "COMMIT_BUILT=0" >> $GITHUB_ENV
-            echo "No artifacts found; will proceed with building a new wheel"
-        elif [ "$RESULT" -eq 1 ]; then
-            echo "COMMIT_BUILT=1" >> $GITHUB_ENV
-            echo "Commit already built; reusing existing wheel"
+        source .venv/bin/activate
+        wheel_name=$(python .github/ci-scripts/get_wheel_name_from_s3.py ${{ github.sha }} $platform_substring)
+        if [ "$wheel_name" ]; then
+          echo "Python wheel for this commit already built and uploaded"
+        else
+          echo "No python wheel for this commit found; proceeding with build"
         fi
+        echo "wheel_name=$wheel_name" >> $GITHUB_ENV
     - name: Build release wheel
       run: |
-        if [ "$COMMIT_BUILT" -eq 1 ]; then
-          echo "Commit already built"
+        if [ "$wheel_name" ]; then
+          echo "Python wheel for this commit already built and uploaded"
           exit 0
         fi
         export CARGO_TARGET_DIR=~/target
-        uv v
         source .venv/bin/activate
         uv pip install pip maturin boto3
         maturin build --release
     - name: Upload wheel to AWS S3
-      id: upload
       run: |
-        if [ "$COMMIT_BUILT" -eq 1 ]; then
-          echo "Commit already built"
-          wheel=$(aws s3 ls s3://github-actions-artifacts-bucket/builds/${{ github.sha }}/ | awk '{print $4}')
-          echo "wheel=$wheel" >> $GITHUB_OUTPUT
-        else
-          count=$(ls ~/target/wheels/*.whl 2> /dev/null | wc -l)
-          if [ "$count" -gt 1 ]; then
-            echo "Found more than 1 wheel"
-            exit 1
-          elif [ "$count" -eq 0 ]; then
-            echo "Found no wheels"
-            exit 1
-          fi
-          for file in ~/target/wheels/*.whl; do
-            aws s3 cp $file s3://github-actions-artifacts-bucket/builds/${{ github.sha }}/ --acl public-read --no-progress;
-            file_basename=$(basename $file)
-            echo "wheel=$file_basename" >> $GITHUB_OUTPUT
-          done
+        if [ "$wheel_name" ]; then
+          echo "Python wheel for this commit already built and uploaded"
+          exit 0
         fi
-        echo "Output python-release-wheel location:" >> $GITHUB_STEP_SUMMARY
-        echo "https://us-west-2.console.aws.amazon.com/s3/buckets/github-actions-artifacts-bucket?prefix=builds/${{ github.sha }}/" >> $GITHUB_STEP_SUMMARY
+        source .venv/bin/activate
+        wheel_name=$(python .github/ci-scripts/upload_wheel_to_s3.py ${{ github.sha }} $platform_substring ~/target/wheels)
+        echo "wheel_name=$wheel_name" >> $GITHUB_ENV
+    - name: Print url of the built wheel to GitHub Actions Summary Page
+      run: |
+        console_url="https://us-west-2.console.aws.amazon.com/s3/object/github-actions-artifacts-bucket?prefix=builds/${{ github.sha }}/$wheel_name"
+        download_url="https://github-actions-artifacts-bucket.s3.us-west-2.amazonaws.com/builds/${{ github.sha }}/$wheel_name"
+
+        echo "View the location of the built wheel in the AWS console here:" >> $GITHUB_STEP_SUMMARY
+        echo "$console_url" >> $GITHUB_STEP_SUMMARY
+
+        echo "Directly download the wheel here:" >> $GITHUB_STEP_SUMMARY
+        echo "$download_url" >> $GITHUB_STEP_SUMMARY