Build wheels

.github/workflows/build-wheels.yml

@@ -0,0 +1,236 @@
+name: Build wheels
+
+on:
+  push:
+    branches: [main]
+    tags: ["*"]
+  pull_request:
+    # Check all PR
+
+
+concurrency:
+  group: python-wheels-${{ github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+
+
+jobs:
+  build-wheels:
+    runs-on: ${{ matrix.os }}
+    name: ${{ matrix.name }} (torch v${{ matrix.torch-version }})
+    strategy:
+      matrix:
+        torch-version: ["2.3", "2.4", "2.5", "2.6"]
+        arch: ["arm64", "x86_64"]
+        os: ["ubuntu-22.04", "ubuntu-22.04-arm", "macos-14"]
+        exclude:
+          # remove mismatched arch/os pairs
+          - {os: macos-14, arch: x86_64}
+          - {os: ubuntu-22.04, arch: arm64}
+          - {os: ubuntu-22.04-arm, arch: x86_64}
+        include:
+          # add `cibw-arch` to the different configurations
+          - name: x86_64 Linux
+            os: ubuntu-22.04
+            arch: x86_64
+            cibw-arch: x86_64
+          - name: arm64 Linux
+            os: ubuntu-22.04-arm
+            arch: arm64
+            cibw-arch: aarch64
+          - name: arm64 macOS
+            os: macos-14
+            arch: arm64
+            cibw-arch: arm64
+          # set the right manylinux image to use
+          - {torch-version: '2.3', manylinux-version: "2014"}
+          - {torch-version: '2.4', manylinux-version: "2014"}
+          - {torch-version: '2.5', manylinux-version: "2014"}
+          # only torch >= 2.6 on arm64-linux needs the newer manylinux
+          - {torch-version: '2.6', arch: x86_64, manylinux-version: "2014"}
+          - {torch-version: '2.6', arch: arm64, manylinux-version: "_2_28"}
+
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.13"
+
+      - name: install dependencies
+        run: python -m pip install cibuildwheel
+
+      - name: build wheel
+        run: python -m cibuildwheel --output-dir ./wheelhouse
+        env:
+          CIBW_BUILD: cp312-*
+          CIBW_SKIP: "*musllinux*"
+          CIBW_ARCHS: "${{ matrix.cibw-arch }}"
+          CIBW_BUILD_VERBOSITY: 1
+          CIBW_MANYLINUX_X86_64_IMAGE: quay.io/pypa/manylinux${{ matrix.manylinux-version }}_x86_64
+          CIBW_MANYLINUX_AARCH64_IMAGE: quay.io/pypa/manylinux${{ matrix.manylinux-version }}_aarch64
+          CIBW_ENVIRONMENT: >
+            PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cpu
+            MACOSX_DEPLOYMENT_TARGET=11
+            PETNC_BUILD_WITH_TORCH_VERSION=${{ matrix.torch-version }}.*
+          # do not complain for missing libtorch.so
+          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}
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: single-version-wheel-${{ matrix.torch-version }}-${{ matrix.os }}-${{ matrix.arch }}
+          path: ./wheelhouse/*.whl
+
+  merge-wheels:
+    needs: build-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 Linux
+            os: ubuntu-22.04-arm
+            arch: arm64
+          - name: arm64 macOS
+            os: macos-14
+            arch: arm64
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Download wheels
+        uses: actions/download-artifact@v4
+        with:
+          pattern: 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.13"
+
+      - 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 {} "pet_neighbors_convert-*.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/pet_neighbors_convert-*/pet_neighbors_convert-*.dist-info/METADATA
+
+          echo "\n\n METADATA = \n\n"
+          cat dist/unpacked/pet_neighbors_convert-*/pet_neighbors_convert-*.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/pet_neighbors_convert-*/pet_neighbors_convert-*.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: wheel-${{ matrix.os }}-${{ matrix.arch }}
+          path: ./wheelhouse/*.whl
+
+  build-sdist:
+    name: sdist
+    runs-on: ubuntu-22.04
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.13"
+
+      - name: install dependencies
+        run: python -m pip install build
+
+      - name: build sdist
+        env:
+          PIP_EXTRA_INDEX_URL: "https://download.pytorch.org/whl/cpu"
+        run: python -m build . --outdir=dist/
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: sdist
+          path: dist/*.tar.gz
+
+  merge-and-release:
+    name: Merge and release wheels/sdists
+    needs: [merge-wheels, build-sdist]
+    runs-on: ubuntu-22.04
+    permissions:
+      contents: write
+      id-token: write
+      pull-requests: write
+    environment:
+      name: pypi
+      url: https://pypi.org/project/pet-neighbors-convert
+    steps:
+      - name: Download wheels
+        uses: actions/download-artifact@v4
+        with:
+          path: wheels
+          pattern: wheel-*
+          merge-multiple: true
+
+      - name: Download sdist
+        uses: actions/download-artifact@v4
+        with:
+          path: wheels
+          name: sdist
+
+      - name: Re-upload a single wheels artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: wheels
+          path: wheels/*
+
+      - name: Comment with download link
+        uses: PicoCentauri/comment-artifact@v1
+        if: github.event.pull_request.head.repo.fork == false
+        with:
+          name: wheels
+          description: ⚙️ Download Python wheels for this pull-request (you can install these with pip)
+
+      - name: Publish distribution to PyPI
+        if: startsWith(github.ref, 'refs/tags/v')
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: wheels
+
+      - name: upload to GitHub release
+        if: startsWith(github.ref, 'refs/tags/v')
+        uses: softprops/action-gh-release@v2
+        with:
+          files: |
+            wheels/*.tar.gz
+            wheels/*.whl
+          prerelease: ${{ contains(github.ref, '-rc') }}
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+

CMakeLists.txt

@@ -0,0 +1,58 @@
+cmake_minimum_required(VERSION 3.27)
+
+if (POLICY CMP0076)
+    #  target_sources() converts relative paths to absolute
+    cmake_policy(SET CMP0076 NEW)
+endif()
+
+project(neighbors_convert CXX)
+
+# Set a default build type if none was specified
+if (${CMAKE_CURRENT_SOURCE_DIR} STREQUAL ${CMAKE_SOURCE_DIR})
+    if("${CMAKE_BUILD_TYPE}" STREQUAL "" AND "${CMAKE_CONFIGURATION_TYPES}" STREQUAL "")
+        message(STATUS "Setting build type to 'relwithdebinfo' as none was specified.")
+        set(
+            CMAKE_BUILD_TYPE "relwithdebinfo"
+            CACHE STRING
+            "Choose the type of build, options are: none(CMAKE_CXX_FLAGS or CMAKE_C_FLAGS used) debug release relwithdebinfo minsizerel."
+            FORCE
+        )
+        set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS release debug relwithdebinfo minsizerel none)
+    endif()
+endif()
+
+# add path to the cmake configuration of the version of libtorch used
+# by the Python torch module. PYTHON_EXECUTABLE is provided by skbuild
+execute_process(
+    COMMAND ${PYTHON_EXECUTABLE} -c "import torch.utils; print(torch.utils.cmake_prefix_path)"
+    RESULT_VARIABLE TORCH_CMAKE_PATH_RESULT
+    OUTPUT_VARIABLE TORCH_CMAKE_PATH_OUTPUT
+    ERROR_VARIABLE TORCH_CMAKE_PATH_ERROR
+)
+
+if (NOT ${TORCH_CMAKE_PATH_RESULT} EQUAL 0)
+message(FATAL_ERROR "failed to find your pytorch installation\n${TORCH_CMAKE_PATH_ERROR}")
+endif()
+
+string(STRIP ${TORCH_CMAKE_PATH_OUTPUT} TORCH_CMAKE_PATH_OUTPUT)
+set(CMAKE_PREFIX_PATH "${CMAKE_PREFIX_PATH};${TORCH_CMAKE_PATH_OUTPUT}")
+
+find_package(Torch 2.3 REQUIRED)
+
+file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/_build_torch_version.py "BUILD_TORCH_VERSION = '${Torch_VERSION}'")
+
+add_library(neighbors_convert SHARED 
+    "src/pet_neighbors_convert/neighbors_convert.cpp"
+)
+
+# only link to `torch_cpu_library` instead of `torch`, which could also include
+# `libtorch_cuda`.
+target_link_libraries(neighbors_convert PUBLIC torch_cpu_library)
+target_include_directories(neighbors_convert PUBLIC "${TORCH_INCLUDE_DIRS}")
+target_compile_definitions(neighbors_convert PUBLIC "${TORCH_CXX_FLAGS}")
+
+target_compile_features(neighbors_convert PUBLIC cxx_std_17)
+
+install(TARGETS neighbors_convert
+    LIBRARY DESTINATION "lib"
+)

MANIFEST.in

@@ -2,12 +2,15 @@ graft src
 
 include LICENSE
 include README.rst
+include CMakeLists.txt
 
-prune tests
 prune .github
 prune .tox
+prune tests
+prune scripts
 
 exclude .gitignore
 exclude tox.ini
+recursive-include build-backend *.py
 
 global-exclude *.py[cod] __pycache__/* *.so *.dylib

build-backend/backend.py

@@ -0,0 +1,26 @@
+# This is a custom Python build backend wrapping setuptool's to add a build-time
+# dependencies on torch when building the wheel and not the sdist
+import os
+from setuptools import build_meta
+
+FORCED_TORCH_VERSION = os.environ.get("PETNC_BUILD_WITH_TORCH_VERSION")
+if FORCED_TORCH_VERSION is not None:
+    TORCH_DEP = f"torch =={FORCED_TORCH_VERSION}"
+else:
+    TORCH_DEP = "torch >=2.3"
+
+# ==================================================================================== #
+#                   Build backend functions definition                                 #
+# ==================================================================================== #
+
+# Use the default version of these
+prepare_metadata_for_build_wheel = build_meta.prepare_metadata_for_build_wheel
+get_requires_for_build_sdist = build_meta.get_requires_for_build_sdist
+build_wheel = build_meta.build_wheel
+build_sdist = build_meta.build_sdist
+
+
+# Special dependencies to build the wheels
+def get_requires_for_build_wheel(config_settings=None):
+    defaults = build_meta.get_requires_for_build_wheel(config_settings)
+    return defaults + [TORCH_DEP]

pyproject.toml

@@ -1,11 +1,12 @@
 [build-system]
-requires = ["setuptools >= 68", "setuptools_scm>=8", "wheel", "torch >= 2.3"]
-build-backend = "setuptools.build_meta"
+requires = ["setuptools >= 68", "setuptools_scm>=8", "wheel >= 0.36"]
+build-backend = "backend"
+backend-path = ["build-backend"]
 
 [project]
 authors = [{name = "lab-cosmo developers"}]
 classifiers = [
-    "Development Status :: 4 - Beta",
+    "Development Status :: 5 - Production/Stable",
     "Intended Audience :: Science/Research",
     "License :: OSI Approved :: BSD License",
     "Operating System :: POSIX",
@@ -15,9 +16,8 @@ classifiers = [
     "Topic :: Scientific/Engineering",
     "Topic :: Software Development :: Libraries :: Python Modules"
 ]
-dependencies = ["torch >= 2.3"]
-description = "Extension for PET model for reordering the neighbors during message passing."
-dynamic = ["version"]
+description = "PET model extension for processing neighbor lists"
+dynamic = ["version", "dependencies"]
 license = {text = "BSD-3-Clause"}
 name = "pet-neighbors-convert"
 readme = "README.rst"