Try to build CUDA wheels

.github/workflows/build-jax-wheels.yml

@@ -0,0 +1,63 @@
+name: Build Jax Wheels
+
+on:
+  push:
+    branches: [main]
+    tags: ["*"]
+  pull_request:
+    # Check all PR
+
+concurrency:
+  group: jax-wheels-${{ github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+
+jobs:
+  build-jax-wheels:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        python-version: ["39", "310", "311", "312", "313"]
+        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: "Py${{matrix.python-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 jax wheels
+      run: python -m cibuildwheel ./sphericart-jax
+      env:
+        CIBW_BUILD_VERBOSITY: 1
+        CIBW_BUILD: "cp${{ matrix.python-version }}-*"
+        CIBW_SKIP: "*-musllinux*"
+        CIBW_ARCHS: "${{ matrix.arch }}"
+        CIBW_MANYLINUX_X86_64_IMAGE: sphericart_manylinux2014_x86_64
+        CIBW_ENVIRONMENT: >
+          SPHERICART_ARCH_NATIVE=OFF
+          SPHERICART_NO_LOCAL_DEPS=1
+
+    - uses: actions/upload-artifact@v4
+      with:
+        name: "sphericart-jax-py-${{ matrix.python-version }}-${{ matrix.os }}-${{ matrix.arch }}"
+        path: ./wheelhouse/*.whl

.github/workflows/build-torch-wheels.yml

@@ -0,0 +1,141 @@
+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-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

.github/workflows/build-wheels.yml

@@ -1,4 +1,4 @@
-name: Build Python wheels
+name: Build Wheels
 
 on:
   push:
@@ -7,124 +7,60 @@ on:
   pull_request:
     # Check all PR
 
-env:
-  SPHERICART_NO_LOCAL_DEPS: "1"
-  MACOSX_DEPLOYMENT_TARGET: "11.0"
+concurrency:
+  group: wheels-${{ github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
 
 jobs:
   build-wheels:
     runs-on: ${{ matrix.os }}
-    name: ${{ matrix.os }}
+    name: ${{matrix.name}}
     strategy:
       matrix:
-        # TODO: add windows builder
-        os: [ubuntu-20.04, macos-13]
+        include:
+          - name: x86_64 Linux
+            os: ubuntu-22.04
+            arch: x86_64
+          - name: x86_64 macOS
+            os: macos-13
+            arch: x86_64
+          - name: arm64 macOS
+            os: macos-14
+            arch: arm64
 
     steps:
-      - uses: actions/checkout@v3
-        with:
-          fetch-depth: 0
+    - name: Checkout code
+      uses: actions/checkout@v4
 
-      - name: Set up Python
-        uses: actions/setup-python@v4
-        with:
-          python-version: "3.12"
+    - 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: Install cibuildwheel
-        run: python -m pip install cibuildwheel build
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.12"
 
-      - 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.12
-          CIBW_BUILD: cp312-*
-          # skip musl and 32-bit builds
-          CIBW_SKIP: "*-musllinux* *-win32 *-manylinux_i686"
-          # on macOS, build both Intel & Apple Silicon versions
-          CIBW_ARCHS_MACOS: x86_64 arm64
-          CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014
-          # do not build wheels with -march=native
-          CIBW_ENVIRONMENT: SPHERICART_ARCH_NATIVE=OFF
+    - name: Install cibuildwheel
+      run: python -m pip install cibuildwheel
 
-      - 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_BUILD_VERBOSITY: 3
-          CIBW_BUILD: cp310-*
-          CIBW_SKIP: "*-musllinux* *-win32 *-manylinux_i686"
-          # we can not build wheels for macos-arm64, since the host is always
-          # x86_64, and we assume we can link against the host version of libtorch
-          CIBW_ARCHS_MACOS: x86_64
-          CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014
-          # Use the CPU only version of torch when building/running the code
-          CIBW_ENVIRONMENT: SPHERICART_ARCH_NATIVE=OFF PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cpu
-          # do not complain for missing libtorch.so in sphericart-torch wheel
-          CIBW_REPAIR_WHEEL_COMMAND_MACOS: |
-            delocate-wheel --ignore-missing-dependencies --require-archs {delocate_archs} -w {dest_dir} -v {wheel}
-          CIBW_REPAIR_WHEEL_COMMAND_LINUX: |
-            auditwheel repair --exclude libtorch.so --exclude libtorch_cpu.so --exclude libc10.so -w {dest_dir} {wheel}
+    - name: Build wheels
+      run: python -m cibuildwheel .
+      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_ENVIRONMENT: >
+          SPHERICART_ARCH_NATIVE=OFF
+          SPHERICART_NO_LOCAL_DEPS=1
+          MACOSX_DEPLOYMENT_TARGET=11.0
 
-      - 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
-
-      - uses: actions/upload-artifact@v3
-        with:
-          name: wheels
-          # only upload sphericart wheels, not sphericart-torch or sphericart-jax
-          path: dist/*.whl
-
-      - name: upload wheel to GitHub release
-        if: startsWith(github.ref, 'refs/tags/')
-        uses: softprops/action-gh-release@v1
-        with:
-          # only upload sphericart wheels, not sphericart-torch or sphericart-jax
-          files: dist/*.whl
-        env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-
-  build-sdist:
-    runs-on: ubuntu-20.04
-    name: sdist
-    steps:
-      - uses: actions/checkout@v3
-        with:
-          fetch-depth: 0
-      - name: Set up Python
-        uses: actions/setup-python@v4
-        with:
-          python-version: "3.12"
-      - name: build sdist
-        run: |
-          python -m pip install build
-
-          python -m build --sdist . --outdir ./dist
-          python -m build --sdist sphericart-torch --outdir ./dist
-          python -m build --sdist sphericart-jax --outdir ./dist
-
-      - uses: actions/upload-artifact@v3
-        with:
-          name: wheels
-          path: dist/*.tar.gz
-      - name: upload sdist to GitHub release
-        if: startsWith(github.ref, 'refs/tags/')
-        uses: softprops/action-gh-release@v1
-        with:
-          files: dist/*.tar.gz
-        env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+    - uses: actions/upload-artifact@v4
+      with:
+        name: sphericart-wheel-${{ matrix.os }}-${{ matrix.arch }}
+        path: ./wheelhouse/*.whl

.github/workflows/docs.yml

@@ -7,6 +7,10 @@ on:
   pull_request:
     # Check all PR
 
+concurrency:
+  group: docs-${{ github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+
 jobs:
   build-and-publish:
     runs-on: ubuntu-20.04