ci: build wheels for linux-aarch64

Author: d4l3k

.github/workflows/build_wheels_linux_aarch64.yml

@@ -0,0 +1,44 @@
+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@d4l3k/conda_no_capture
+    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
+      runner: linux.arm64.r7g.12xlarge.memory

build_ncclx.sh

@@ -134,7 +134,7 @@ function build_third_party {
     build_fb_oss_library "https://github.com/fastfloat/fast_float.git" "v8.0.2" fast_float "-DFASTFLOAT_INSTALL=ON"
     build_fb_oss_library "https://github.com/libevent/libevent.git" "release-2.1.12-stable" event
     build_fb_oss_library "https://github.com/google/double-conversion.git" "v3.3.1" double-conversion
-    build_fb_oss_library "https://github.com/facebook/folly.git" "$third_party_tag" folly "-DUSE_STATIC_DEPS_ON_UNIX=ON"
+    build_fb_oss_library "https://github.com/facebook/folly.git" "$third_party_tag" folly "-DUSE_STATIC_DEPS_ON_UNIX=ON -DOPENSSL_USE_STATIC_LIBS=ON"
   else
     if [[ -z "${NCCL_SKIP_CONDA_INSTALL}" ]]; then
       DEPS=(

scripts/_build_wheel.sh

@@ -10,9 +10,7 @@ set -ex
 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
+conda install -y -c conda-forge libopenssl-static
 
 # 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,14 +8,35 @@
 
 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
+
+# Python 3.10
 #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
+
+# Python 3.13
 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
 
+# Python 3.13t
+#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 python-freethreading -c conda-forge
+
 CONDA_RUN="conda run --no-capture-output -p ${CONDA_ENV}"
 
 ${CONDA_RUN} pip install torch --pre --index-url https://download.pytorch.org/whl/nightly/cu128
 
+# Uncomment to debug the build w/ docker exec
+#${CONDA_RUN} bash -c "sleep 10000000000000"
+
 ${CONDA_RUN} bash scripts/_build_wheel.sh

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

scripts/docker_build_wheel.sh

@@ -8,6 +8,9 @@
 
 set -ex
 
+docker stop torchcomms || true
+docker rm torchcomms || true
+
 docker run --name torchcomms \
     --net=host \
     -i \