build-torch-wheels.yml

name: Build Torch Wheels

on:
  push:
    branches: [main]
    tags: ["*"]
  pull_request:
    # Check all PR

concurrency:
  group: torch-wheels-${{ github.ref }}
  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}

jobs:
  build-torch-sdist:
    runs-on: ubuntu-22.04
    name: sdist
    env:
      SPHERICART_NO_LOCAL_DEPS: 1
    steps:
    - name: Checkout code
      uses: actions/checkout@v4

    - name: Set up Python
      uses: actions/setup-python@v5
      with:
        python-version: "3.12"

    - name: Install dependencies
      run: python -m pip install build

    - name: Build torch sdist
      run: python -m build --sdist ./sphericart-torch --outdir ./dist

    - uses: actions/upload-artifact@v4
      with:
        name: "sphericart-torch-sdist"
        path: ./dist/*.tar.gz

    - name: upload to GitHub release
      if: startsWith(github.ref, 'refs/tags/')
      uses: softprops/action-gh-release@v2
      with:
        files: |
          ./dist/*.tar.gz
        prerelease: ${{ contains(github.ref, '-rc') }}
      env:
        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

  build-torch-wheels:
    runs-on: ${{ matrix.os }}
    strategy:
      matrix:
        pytorch-version: ["2.4", "2.5"]
        arch: ["arm64", "x86_64"]
        os: ["ubuntu-22.04", "macos-14"]
        exclude:
          # remove mismatched arch/os pairs
          - {os: macos-14, arch: x86_64}
          - {os: ubuntu-22.04, arch: arm64}

    name: "Torch ${{matrix.pytorch-version}} ${{matrix.os}} ${{matrix.arch}}"

    steps:
    - name: Checkout code
      uses: actions/checkout@v4

    - name: Build custom manylinux Docker Image with CUDA
      if: matrix.os == 'ubuntu-22.04'
      run: |
        docker build --no-cache \
            -t sphericart_manylinux2014_"${{ matrix.arch }}" \
            scripts/manylinux2014_"${{ matrix.arch }}"

    - name: Set up Python
      uses: actions/setup-python@v5
      with:
        python-version: "3.12"

    - name: Install cibuildwheel
      run: python -m pip install cibuildwheel

    - name: Build torch wheels
      run: python -m cibuildwheel ./sphericart-torch
      env:
        CIBW_BUILD_VERBOSITY: 1
        CIBW_BUILD: "cp312-*"
        CIBW_SKIP: "*-musllinux*"
        CIBW_ARCHS: "${{ matrix.arch }}"
        CIBW_MANYLINUX_X86_64_IMAGE: sphericart_manylinux2014_x86_64
        CIBW_REPAIR_WHEEL_COMMAND_LINUX: |
          auditwheel repair \
              --exclude libc10.so \
              --exclude libtorch.so \
              --exclude libtorch_cpu.so \
              --exclude libtorch_cuda.so \
              --exclude libc10_cuda.so  \
              -w {dest_dir} {wheel}
        CIBW_REPAIR_WHEEL_COMMAND_MACOS: |
          delocate-wheel --ignore-missing-dependencies \
                         --require-archs {delocate_archs} \
                         -w {dest_dir} -v {wheel}
        CIBW_ENVIRONMENT: >
          SPHERICART_ARCH_NATIVE=OFF
          SPHERICART_NO_LOCAL_DEPS=1
          TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0"
          SPHERICART_TORCH_BUILD_WITH_TORCH_VERSION=${{ matrix.pytorch-version }}.*

    - uses: actions/upload-artifact@v4
      with:
        name: torch-single-version-wheel-${{ matrix.os }}-${{ matrix.arch }}-torch-${{ matrix.pytorch-version }}
        path: ./wheelhouse/*.whl

  merge-torch-wheels:
    needs: build-torch-wheels
    runs-on: ubuntu-22.04
    name: merge wheels for ${{ matrix.name }}
    strategy:
      matrix:
        include:
          - name: x86_64 Linux
            os: ubuntu-22.04
            arch: x86_64
          - name: arm64 macOS
            os: macos-14
            arch: arm64
    steps:
      - uses: actions/checkout@v4

      - name: Download wheels
        uses: actions/download-artifact@v4
        with:
          pattern: torch-single-version-wheel-${{ matrix.os }}-${{ matrix.arch }}-*
          merge-multiple: false
          path: dist

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: "3.12"

      - name: install dependencies
        run: python -m pip install twine wheel

      - name: merge wheels
        run: |
          # collect all torch versions used for the build
          REQUIRES_TORCH=$(find dist -name "*.whl" -exec unzip -p {} "sphericart_torch-*.dist-info/METADATA" \; | grep "Requires-Dist: torch")
          MERGED_TORCH_REQUIRE=$(python scripts/create-torch-versions-range.py "$REQUIRES_TORCH")

          echo MERGED_TORCH_REQUIRE=$MERGED_TORCH_REQUIRE

          # unpack all single torch versions wheels in the same directory
          mkdir dist/unpacked
          find dist -name "*.whl" -print -exec python -m wheel unpack --dest dist/unpacked/ {} ';'

          sed -i "s/Requires-Dist: torch.*/$MERGED_TORCH_REQUIRE/" dist/unpacked/sphericart_torch-*/sphericart_torch-*.dist-info/METADATA

          echo "\n\n METADATA = \n\n"
          cat dist/unpacked/sphericart_torch-*/sphericart_torch-*.dist-info/METADATA

          # check the right metadata was added to the file. grep will exit with
          # code `1` if the line is not found, which will stop CI
          grep "$MERGED_TORCH_REQUIRE" dist/unpacked/sphericart_torch-*/sphericart_torch-*.dist-info/METADATA

          # repack the directory as a new wheel
          mkdir wheelhouse
          python -m wheel pack --dest wheelhouse/ dist/unpacked/*

      - name: check wheels with twine
        run: twine check wheelhouse/*

      - uses: actions/upload-artifact@v4
        with:
          name: sphericart-torch-wheel-${{ matrix.os }}-${{ matrix.arch }}
          path: ./wheelhouse/*.whl

      - name: upload to GitHub release
        if: startsWith(github.ref, 'refs/tags/')
        uses: softprops/action-gh-release@v2
        with:
          files: |
            ./wheelhouse/*.whl
          prerelease: ${{ contains(github.ref, '-rc') }}
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}