build-cuda-wheels.yml

# inspired by https://github.com/AutoGPTQ/AutoGPTQ/blob/main/.github/workflows/build_wheels_cuda.yml

name: Build Python wheels with CUDA

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

env:
  SPHERICART_NO_LOCAL_DEPS: "1"

jobs:
  build_wheels:
    name: ${{ matrix.os }}, torch ${{ matrix.torch }}, cuda ${{ matrix.cuda }}
    runs-on: ${{ matrix.os }}
    strategy:
      matrix:
        os: [ubuntu-20.04]
        torch: ["2.0.0", "2.1.0", "2.2.0"]
        cuda: ["11.8", "12.1"]
        exclude:
          # torch 1.13.0 does not support cuda 12.1
          - os: ubuntu-20.04
            torch: "1.13.1"
            cuda: "12.1"

    steps:
      - name: Prepare cuda_no_point variable
        id: prepare_cuda_no_point
        run: |
          export cuda_no_point=$(echo ${{ matrix.cuda }} | sed 's/\.//')
          echo ::set-output name=cuda_no_point::$cuda_no_point

      - name: free disk space
        run: sudo rm -rf /usr/share/dotnet /usr/local/lib/android || true

      - uses: actions/checkout@v3

      - name: Set up Python
        uses: actions/setup-python@v4
        with:
          python-version: "3.10"

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

      - name: Build sphericart wheels
        run: |
          # ensure we build the wheel from the sdist, not the checkout
          python -m build --sdist . --outdir dist
          python -m cibuildwheel dist/*.tar.gz --output-dir dist
        env:
          CIBW_BUILD_VERBOSITY: 3
          # build wheels on CPython 3.10
          CIBW_BUILD: cp39-* cp310-* cp311-* cp312-*
          # skip musl and 32-bit builds
          CIBW_SKIP: "*-musllinux* *-win32 *-manylinux_i686"
          # on macOS, build both Intel & Apple Silicon versions
          CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014
          # do not build wheels with -march=native
          CIBW_ENVIRONMENT: SPHERICART_ARCH_NATIVE=OFF

      - name: Build sphericart-torch wheels
        run: |
          # ensure we build the wheel from the sdist, not the checkout
          python -m build --sdist sphericart-torch --outdir sphericart-torch/dist
          python -m cibuildwheel sphericart-torch/dist/*.tar.gz --output-dir sphericart-torch/dist
        env:
          CIBW_BEFORE_ALL: bash /host/home/runner/work/sphericart/sphericart/scripts/cibw-cuda-setup.sh ${{ matrix.cuda }}
          CIBW_BUILD_VERBOSITY: 3
          CIBW_BUILD: >
            bash -c 'if [[ "${{ matrix.torch }}" == "2.2.0" ]]; then
              echo "cp39-* cp310-* cp311-* cp312-*"
            else
              echo "cp39-* cp310-* cp311-*"
            fi'
          CIBW_SKIP: "*-musllinux* *-win32 *-manylinux_i686"
          CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014
          # set environment variables for sphericart-torch build
          CIBW_ENVIRONMENT: SPHERICART_ARCH_NATIVE=OFF CUDACXX=/usr/local/cuda/bin/nvcc TORCH_CUDA_ARCH_LIST=All CUDAARCHS=all SPHERICART_TORCH_TORCH_VERSION=${{ matrix.torch }} PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cu${{ steps.prepare_cuda_no_point.outputs.cuda_no_point }}
          # do not complain for missing libtorch.so in sphericart-torch wheel
          CIBW_REPAIR_WHEEL_COMMAND_LINUX: |
            auditwheel repair --exclude libtorch.so --exclude libtorch_cpu.so --exclude libtorch_cuda.so --exclude libc10.so --exclude libc10_cuda.so -w {dest_dir} {wheel}

      # - name: Build sphericart-jax wheels
      #   run: |
      #     # ensure we build the wheel from the sdist, not the checkout
      #     python -m build --sdist sphericart-jax --outdir sphericart-jax/dist
      #     python -m cibuildwheel sphericart-jax/dist/*.tar.gz --output-dir sphericart-jax/dist
      #   env:
      #     CIBW_BUILD_VERBOSITY: 3
      #     CIBW_BUILD: cp310-*
      #     CIBW_SKIP: "*-musllinux* *-win32 *-manylinux_i686"
      #     CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014
      #     CIBW_ENVIRONMENT: SPHERICART_ARCH_NATIVE=OFF