Extend build-commit workflow to also work for ARM builds

.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,34 @@
+import sys
+from pathlib import Path
+
+import boto3
+import constants
+import wheellib
+
+if __name__ == "__main__":
+    commit_hash = sys.argv[1]
+    platform_substring = sys.argv[2]
+    path_to_wheel_dir = Path(sys.argv[3])
+
+    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 Exception(f"No wheels found that match the given platform substring: {platform_substring}; expected 1")
+    elif length > 1:
+        raise Exception(
+            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(wheelpath, constants.BUCKET_NAME, str(destination))
+    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,35 @@ name: build-commit
 
 on:
   workflow_dispatch:
-  workflow_call:
-    secrets:
-      ACTIONS_AWS_ROLE_ARN:
-        description: The ARN of the AWS role to assume
-        required: true
-    outputs:
-      wheel:
-        description: The wheel file that was built
-        value: ${{ jobs.build-commit.outputs.wheel }}
+    inputs:
+      build_for_arm:
+        type: boolean
+        description: "Build for ARM"
+        required: false
+        default: false
+      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.build_for_arm && 'buildjet-8vcpu-ubuntu-2204-arm' || 'buildjet-8vcpu-ubuntu-2004' }}
+    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.build_for_arm }}" == "true" ]; then
+          platform_substring=aarch
+        else
+          platform_substring=x86
+        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 +41,52 @@ 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.build_for_arm && 'arm' || 'x86' }}-cargo-deps-${{ hashFiles('**/Cargo.lock') }}
+        restore-keys: ${{ runner.os }}-${{ inputs.build_for_arm && 'arm' || 'x86' }}-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
+        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: |
+        url="https://us-west-2.console.aws.amazon.com/s3/object/github-actions-artifacts-bucket?prefix=builds/${{ github.sha }}/$wheel_name"
         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
+        echo "$url" >> $GITHUB_STEP_SUMMARY

pyproject.toml

@@ -90,6 +90,9 @@ typeCheckingMode = "off"
 venv = ".venv"
 venvPath = "."
 
+[[tool.pyright.executionEnvironments]]
+root = ".github/ci-scripts"
+
 [tool.pytest.ini_options]
 addopts = "-m 'not (integration or benchmark or hypothesis)'"
 minversion = "6.0"