From 22632c895d5046605492863e20c8c0a3dfc0406d Mon Sep 17 00:00:00 2001
From: Eric Larson <larson.eric.d@gmail.com>
Date: Sat, 1 Feb 2025 06:11:31 -0500
Subject: [PATCH] ENH: Build aarch64 natively (#258)

* ENH: Build aarch64 natively

* FIX: Simplify

* FIX: Paren

* FIX: Numba should be available

* FIX: Why

* FIX: Mystery solved
---
 .github/workflows/cibuildwheel.yml | 37 +++++++-----------------------
 pyproject.toml                     | 10 ++++----
 tests/test_heartbeats.py           |  2 ++
 3 files changed, 15 insertions(+), 34 deletions(-)

diff --git a/.github/workflows/cibuildwheel.yml b/.github/workflows/cibuildwheel.yml
index 0c1fc65a..d3db1a41 100644
--- a/.github/workflows/cibuildwheel.yml
+++ b/.github/workflows/cibuildwheel.yml
@@ -30,42 +30,25 @@ jobs:
     needs: style
     strategy:
       matrix:
-        os: [ubuntu-22.04, windows-2022, macos-15, macos-13]
-        arch: [native]
-        include:
-          - os: ubuntu-22.04
-            arch: aarch64
+        os: [ubuntu-22.04, ubuntu-24.04-arm, windows-2022, macos-15, macos-13]
       fail-fast: false
-    timeout-minutes: 15  # Linux ~1 min, Windows ~4 min, aarch64 emulated ~8 min with tests (~3 min without)
-    name: cibuildwheel (${{ matrix.os }} ${{ matrix.arch }})
+    timeout-minutes: 10  # Linux ~1 min, Windows ~4 min
+    name: cibuildwheel (${{ matrix.os }})
     runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v4
-      - name: Set wheel tests to be skipped on emulated archs on PRs
-        run: |
-          if [[ "${{ github.event_name }}" == "pull_request" ]] && [[ "${{ matrix.arch }}" == "aarch64" ]]; then
-            echo "CIBW_TEST_SKIP=*" | tee -a $GITHUB_ENV
-          fi
-        shell: bash
-      # For aarch64 support https://cibuildwheel.pypa.io/en/stable/faq/#emulation
-      - uses: docker/setup-qemu-action@v3
-        with:
-          platforms: all
-        if: runner.os == 'Linux' && matrix.arch == 'aarch64'
       - name: Build wheels and run tests
         uses: pypa/cibuildwheel@v2.22.0
-        env:
-          CIBW_ARCHS: ${{ matrix.arch }}
       - uses: actions/upload-artifact@v4
         with:
-          name: artifact-${{ matrix.os }}-${{ matrix.arch }}
+          name: artifact-${{ matrix.os }}
           path: ./wheelhouse/*.whl
 
   test:
     needs: build_wheels
     strategy:
       matrix:
-        os: [ubuntu-22.04, windows-2022, macos-15, macos-13]
+        os: [ubuntu-22.04, ubuntu-24.04-arm, windows-2022, macos-15, macos-13]
         python-version: ['3.12']
         include:
           - os: ubuntu-22.04
@@ -84,16 +67,12 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - uses: actions/download-artifact@v4
         with:
-          pattern: artifact-${{ matrix.os }}-native
+          pattern: artifact-${{ matrix.os }}
           merge-multiple: true
           path: dist
       - run: ls -al dist
-      - name: Triage dependencies
-        run: |
-          echo "PIP_EXTRA=wfdb numba pytest-error-for-skips" | tee -a $GITHUB_ENV
-          echo "PYTEST_EXTRA=--error-for-skips" | tee -a $GITHUB_ENV
-      - run: python -m pip install ./dist/*.whl pytest pytest-cov edfio $PIP_EXTRA --only-binary="numpy,numba,edfio"
-      - run: pytest -rfEXs --cov=sleepecg --cov-report=xml --tb=short --cov-branch --color=yes $PYTEST_EXTRA tests/
+      - run: python -m pip install ./dist/*.whl pytest pytest-cov edfio wfdb numba pytest-error-for-skips --only-binary="numpy,numba,llvmlite,edfio"
+      - run: pytest -rfEXs --cov=sleepecg --cov-report=xml --tb=short --cov-branch --color=yes --error-for-skips tests/
       - uses: codecov/codecov-action@v5
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
diff --git a/pyproject.toml b/pyproject.toml
index 36b4eecd..cce90a67 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -63,7 +63,10 @@ doc = [  # RTD uses this when building the documentation
 
 cibw = [  # cibuildwheel uses this for running the test suite
     "edfio >= 0.4.4",
-    "numba >= 0.61.0",
+    # numba 0.61.0 is available on aarch64, but is not compatible with the image
+    # we use to build it. So skip numba tests on aarch64 during cibuildwheel there,
+    # and rely on CIs with Ubuntu to test it with numba.
+    "numba >= 0.61.0; platform_machine != 'aarch64'",
     "wfdb >= 4.2.0",
 ]
 
@@ -79,6 +82,7 @@ test-extras = "cibw"
 test-command = "pytest {package}"
 skip = 'pp* *musllinux*'
 build = "cp310-*"  # abi3
+archs = "native"
 
 [[tool.cibuildwheel.overrides]]
 select = "*"
@@ -89,10 +93,6 @@ repair-wheel-command = [
 
 [tool.cibuildwheel.linux]
 manylinux-x86_64-image = "manylinux2014"
-archs = "x86_64"
-
-[tool.cibuildwheel.macos]
-archs = "x86_64 arm64"
 
 [tool.cibuildwheel.windows]
 archs = "AMD64"
diff --git a/tests/test_heartbeats.py b/tests/test_heartbeats.py
index 0a1a1c34..46199323 100644
--- a/tests/test_heartbeats.py
+++ b/tests/test_heartbeats.py
@@ -33,6 +33,8 @@ def mitdb_234_MLII():
 @pytest.mark.parametrize("backend", ["c", "numba", "python"])
 def test_detect_heartbeats(mitdb_234_MLII, backend):
     """Test heartbeat detection on mitdb:234:MLII."""
+    if backend == "numba":
+        pytest.importorskip("numba")
     record = mitdb_234_MLII
     detection = detect_heartbeats(record.ecg, record.fs, backend=backend)
     TP, FP, FN = compare_heartbeats(detection, record.annotation, int(record.fs / 10))