ci: build wheels for linux-aarch64

Author: d4l3k

.github/workflows/build_wheels_linux_aarch64.yml

@@ -0,0 +1,43 @@
+name: Build Linux Aarch64 Wheels
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+      - nightly
+  workflow_dispatch:
+
+jobs:
+  generate-matrix:
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    with:
+      package-type: wheel
+      os: linux-aarch64
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      with-cuda: enable
+      with-rocm: disable
+      with-cpu: disable
+      # TODO: Python free threading builds are broken due to newer openssl
+      # requirements.
+      python-versions: '["3.10","3.11","3.12","3.13","3.14"]'
+  build:
+    needs: generate-matrix
+    name: ${{ matrix.repository }}
+    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
+    with:
+      repository: meta-pytorch/torchcomms
+      ref: ""
+      pre-script: ""
+      post-script: ""
+      build-platform: "python-build-package"
+      build-command: "scripts/_build_wheel.sh"
+      smoke-test-script: "scripts/smoke_test.py"
+      package-name: torchcomms
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      trigger-event: ${{ github.event_name }}
+      architecture: aarch64
+      #setup-miniconda: false

scripts/_build_wheel.sh

@@ -11,8 +11,10 @@ dnf config-manager --set-enabled powertools
 dnf install -y almalinux-release-devel
 dnf install -y openssl-static ninja-build cmake
 
-# Install older openssl which is compatible with manylinux2.28 system openssl.
-conda install -y openssl==3.0.18
+if [ "$(uname -m)" != "aarch64" ]; then
+    # Install older openssl which is compatible with manylinux2.28 system openssl.
+    conda install -y openssl==3.0.18
+fi
 
 # Nuke conda cmake, ninja and libstdc++ we want to install to use system libraries.
 rm -f "$CONDA_PREFIX/lib/libstdc"* || true

scripts/_emulate_build_wheel.sh

@@ -8,8 +8,19 @@
 
 set -ex
 
+if [ "$(uname -m)" = "aarch64" ]; then
+    echo "Building for aarch64"
+    curl -L -o /mambaforge.sh https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-aarch64.sh
+    chmod +x /mambaforge.sh
+    /mambaforge.sh -b -p /opt/conda
+    rm /mambaforge.sh
+    source /opt/conda/etc/profile.d/conda.sh
+fi
+
 cd /torchcomms
 
+
+
 CONDA_ENV=/tmp/conda_env
 #conda create --yes --quiet --prefix "$CONDA_ENV" python=3.10 cmake=3.31.2 ninja=1.12.1 pkg-config=0.29 wheel=0.37
 conda create --yes --quiet --prefix "$CONDA_ENV" python=3.13 cmake=3.31.2 ninja=1.12.1 pkg-config=0.29 wheel=0.37

scripts/docker_build_linux_aarch64_wheel.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# DO NOT DELETE
+# This script is used to build manylinux releases of torchcomms. The output from
+# this script should be able to be used on most modern OSes as long as the
+# Python version and CUDA version match.
+
+set -ex
+
+docker stop torchcomms || true
+docker rm torchcomms || true
+
+docker run --name torchcomms \
+    --arch=aarch64 \
+    --net=host \
+    -i \
+    -t \
+    -v ".:/torchcomms" \
+    pytorch/manylinuxaarch64-builder:cuda12.9 \
+    bash /torchcomms/scripts/_emulate_build_wheel.sh