From 1497c37bd405ca7204cbb18aca3ca926a135e263 Mon Sep 17 00:00:00 2001
From: Ali Kuwajerwala <ali@kscale.dev>
Date: Fri, 20 Dec 2024 20:13:03 -0800
Subject: [PATCH] make seperate job for QEMU

---
 .github/workflows/publish.yml | 108 +++++++++++++++++++++++++++++++++-
 1 file changed, 107 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index ea462cc..c7f634b 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -14,6 +14,9 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  #
+  # 1) Fast builds for ubuntu + macOS
+  #
   build-wheels:
     strategy:
       matrix:
@@ -35,6 +38,9 @@ jobs:
         with:
           toolchain: stable
 
+      # Notice: we do NOT set up QEMU here,
+      # so macOS won't fail on QEMU steps.
+
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
@@ -60,12 +66,17 @@ jobs:
             CARGO_NET_GIT_FETCH_WITH_CLI=true
         run: |
           cibuildwheel --output-dir dist
+
       - name: Upload wheel artifacts
         uses: actions/upload-artifact@v3
         with:
           name: wheels-${{ matrix.os }}
           path: |
             dist/*.whl
+
+  #
+  # 2) Build source distribution
+  #
   build-source-dist:
     name: Build and publish Python package (source distribution)
     runs-on: ubuntu-latest
@@ -93,12 +104,17 @@ jobs:
       - name: Build source distribution
         run: |
           python -m build --sdist --outdir dist
+
       - name: Upload source distribution
         uses: actions/upload-artifact@v3
         with:
           name: source-dist
           path: |
             dist/*.tar.gz
+
+  #
+  # 3) Publish the main wheels first
+  #
   publish-wheels:
     needs: [build-wheels, build-source-dist]
     name: Publish Python wheels
@@ -115,11 +131,15 @@ jobs:
           mkdir -p final_dist
           find dist -name "*.whl" -exec mv {} final_dist/ \;
           find dist -name "*.tar.gz" -exec mv {} final_dist/ \;
+
       - name: Publish package
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           packages-dir: final_dist/
 
+  #
+  # 4) Publish Rust package
+  #
   publish-rust:
     name: Build and publish Rust package
     runs-on: ubuntu-latest
@@ -137,6 +157,7 @@ jobs:
         run: |
           sudo apt-get update
           sudo apt-get install -y libudev-dev pkg-config
+
       - name: Install protoc
         uses: arduino/setup-protoc@v3
 
@@ -147,6 +168,7 @@ jobs:
           key: ${{ runner.os }}-cargo-registry
           restore-keys: |
             ${{ runner.os }}-cargo-registry
+
       - name: Cache Cargo index
         uses: actions/cache@v2
         with:
@@ -154,8 +176,92 @@ jobs:
           key: ${{ runner.os }}-cargo-index
           restore-keys: |
             ${{ runner.os }}-cargo-index
+
       - name: Publish K-Rec package to crates.io
         env:
           CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}
         run: |
-          cargo publish -p krec
\ No newline at end of file
+          cargo publish -p krec
+
+  #
+  # 5) Build QEMU-based wheels for aarch64 only (no timeouts)
+  #
+  build-qemu-wheels:
+    name: Build QEMU-based wheels (aarch64 only)
+    needs: [publish-wheels]
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.11"
+
+      - name: Set up Rust
+        uses: actions-rs/toolchain@v1
+        with:
+          toolchain: stable
+
+      # We do QEMU setup only here, for aarch64
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v2
+        # "platforms: aarch64" ensures we only enable QEMU for ARM64
+        # so we don't attempt ppc64le or s390x, etc.
+        with:
+          platforms: aarch64
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install cibuildwheel
+        shell: bash
+
+      - name: Build package with QEMU
+        env:
+          CIBW_ARCHS_LINUX: "aarch64"
+          CIBW_QEMU: "true"
+          # Skip PyPy, musllinux, i686, ppc64le, s390x, etc.
+          CIBW_SKIP: "pp* *-musllinux* *-i686* *-ppc64le* *-s390x*"
+          CIBW_BEFORE_ALL_LINUX: |
+            curl -L https://github.com/protocolbuffers/protobuf/releases/download/v28.3/protoc-28.3-linux-x86_64.zip -o protoc.zip
+            unzip protoc.zip -d protoc
+            mv protoc/bin/protoc /usr/local/bin/protoc
+            mv protoc/include/google /usr/local/include/google
+          CIBW_BEFORE_BUILD: |
+            curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+            source $HOME/.cargo/env
+            pip install setuptools-rust
+          CIBW_ENVIRONMENT: |
+            PATH="/usr/local/bin:$HOME/.cargo/bin:$PATH"
+            CARGO_NET_GIT_FETCH_WITH_CLI=true
+        run: |
+          cibuildwheel --output-dir dist
+
+      - name: Upload QEMU wheel artifacts
+        uses: actions/upload-artifact@v3
+        with:
+          name: wheels-qemu
+          path: dist/*.whl
+
+  #
+  # 6) Publish the QEMU-based wheels
+  #
+  publish-qemu-wheels:
+    name: Publish QEMU-based wheels
+    needs: [build-qemu-wheels]
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Download QEMU wheels
+        uses: actions/download-artifact@v3
+        with:
+          name: wheels-qemu
+          path: dist
+
+      - name: Publish QEMU wheels
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: dist