Build featomic-torch wheels on CI

.github/workflows/build-wheels.yml

@@ -98,6 +98,197 @@ jobs:
           name: featomic-wheel-${{ matrix.os }}-${{ matrix.cibw-arch }}
           path: ./wheelhouse/*.whl
 
+  build-torch-wheels:
+    runs-on: ${{ matrix.os }}
+    name: ${{ matrix.name }} (torch v${{ matrix.torch-version }})
+    strategy:
+      matrix:
+        torch-version: ['2.1', '2.2', '2.3', '2.4', '2.5']
+        arch: ['arm64', 'x86_64']
+        os: ['ubuntu-22.04', 'macos-13', 'macos-14', 'windows-2019']
+        exclude:
+          # remove mismatched arch for macOS
+          - {os: macos-14, arch: x86_64}
+          - {os: macos-13, arch: arm64}
+          # no arm64-windows build
+          - {os: windows-2019, arch: arm64}
+          # arch x86_64 on macos is only supported for torch <2.3
+          - {os: macos-13, arch: x86_64, torch-version: '2.3'}
+          - {os: macos-13, arch: x86_64, torch-version: '2.4'}
+          - {os: macos-13, arch: x86_64, torch-version: '2.5'}
+        include:
+          # add `cibw-arch` and `rust-target` to the different configurations
+          - name: x86_64 Linux
+            os: ubuntu-22.04
+            arch: x86_64
+            rust-target: x86_64-unknown-linux-gnu
+            cibw-arch: x86_64
+          - name: arm64 Linux
+            os: ubuntu-22.04
+            arch: arm64
+            rust-target: aarch64-unknown-linux-gnu
+            cibw-arch: aarch64
+          - name: x86_64 macOS
+            os: macos-13
+            arch: x86_64
+            rust-target: x86_64-apple-darwin
+            cibw-arch: x86_64
+          - name: arm64 macOS
+            os: macos-14
+            arch: arm64
+            rust-target: aarch64-apple-darwin
+            cibw-arch: arm64
+          - name: x86_64 Windows
+            os: windows-2019
+            arch: x86_64
+            rust-target: x86_64-pc-windows-msvc
+            cibw-arch: AMD64
+          # add the right python version for each torch version
+          - {torch-version: '2.1',  python-version: '3.11', cibw-python: 'cp311-*'}
+          - {torch-version: '2.2',  python-version: '3.12', cibw-python: 'cp312-*'}
+          - {torch-version: '2.3',  python-version: '3.12', cibw-python: 'cp312-*'}
+          - {torch-version: '2.4',  python-version: '3.12', cibw-python: 'cp312-*'}
+          - {torch-version: '2.5',  python-version: '3.12', cibw-python: 'cp312-*'}
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: setup rust
+        uses: dtolnay/rust-toolchain@master
+        with:
+          toolchain: stable
+          target: ${{ matrix.rust-target }}
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: install dependencies
+        run: python -m pip install cibuildwheel
+
+      - name: Set up QEMU for docker
+        if: matrix.os == 'ubuntu-22.04'
+        uses: docker/setup-qemu-action@v3
+
+      - name: build manylinux with rust docker image
+        if: matrix.os == 'ubuntu-22.04'
+        run: docker buildx build -t rustc-manylinux2014_${{ matrix.cibw-arch }} python/scripts/rustc-manylinux2014_${{ matrix.cibw-arch }}
+
+      - name: build featomic-torch wheel
+        run: python -m cibuildwheel python/featomic_torch
+        env:
+          CIBW_BUILD: ${{ matrix.cibw-python}}
+          CIBW_SKIP: "*musllinux*"
+          CIBW_ARCHS: ${{ matrix.cibw-arch }}
+          CIBW_BUILD_VERBOSITY: 1
+          CIBW_MANYLINUX_X86_64_IMAGE: rustc-manylinux2014_x86_64
+          CIBW_MANYLINUX_AARCH64_IMAGE: rustc-manylinux2014_aarch64
+          # FEATOMIC_NO_LOCAL_DEPS is set to 1 when building a tag of
+          # featomic-torch, which will force to use the version of featomic
+          # already released on PyPI. Otherwise, this will use the version of
+          # featomic from git checkout (in case there are unreleased breaking
+          # changes). This means we can not release breaking changes in featomic
+          # and v-torch by putting a tag on the same commit. Instead featomic
+          # must be fully released before we start the build of featomic-torch
+          # wheels.
+          CIBW_ENVIRONMENT: >
+            FEATOMIC_NO_LOCAL_DEPS=${{ startsWith(github.ref, 'refs/tags/featomic-torch-v') && '1' || '0' }}
+            FEATOMIC_TORCH_BUILD_WITH_TORCH_VERSION=${{ matrix.torch-version }}.*
+            PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cpu
+            MACOSX_DEPLOYMENT_TARGET=11
+          # do not complain for missing libtorch.so, libfeatomic.so, & co
+          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 libfeatomic.so \
+                              --exclude libmetatensor.so \
+                              --exclude libmetatensor_torch.so \
+                              --exclude libtorch.so \
+                              --exclude libtorch_cpu.so \
+                              --exclude libc10.so \
+                              -w {dest_dir} {wheel}
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: featomic-torch-single-version-wheel-${{ matrix.torch-version }}-${{ matrix.os }}-${{ matrix.arch }}
+          path: ./wheelhouse/*.whl
+
+  merge-torch-wheels:
+    needs: build-torch-wheels
+    runs-on: ubuntu-22.04
+    name: merge metatensor-torch ${{ matrix.name }}
+    strategy:
+      matrix:
+        include:
+          - name: x86_64 Linux
+            os: ubuntu-22.04
+            arch: x86_64
+          - name: arm64 Linux
+            os: ubuntu-22.04
+            arch: arm64
+          - name: x86_64 macOS
+            os: macos-13
+            arch: x86_64
+          - name: arm64 macOS
+            os: macos-14
+            arch: arm64
+          - name: x86_64 Windows
+            os: windows-2019
+            arch: x86_64
+    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 {} "featomic_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/featomic_torch-*/featomic_torch-*.dist-info/METADATA
+
+          echo "\n\n METADATA = \n\n"
+          cat dist/unpacked/featomic_torch-*/featomic_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/featomic_torch-*/featomic_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: featomic-torch-wheel-${{ matrix.os }}-${{ matrix.arch }}
+          path: ./wheelhouse/*.whl
 
   build-others:
     runs-on: ubuntu-22.04
@@ -145,7 +336,7 @@ jobs:
 
   merge-and-release:
     name: Merge and release wheels/sdists
-    needs: [build-wheels, build-others]
+    needs: [build-wheels, merge-torch-wheels, build-others]
     runs-on: ubuntu-22.04
     permissions:
       contents: write
@@ -157,6 +348,13 @@ jobs:
           pattern: featomic-wheel-*
           merge-multiple: true
 
+      - name: Download featomic-torch wheels
+        uses: actions/download-artifact@v4
+        with:
+          path: wheels
+          pattern: featomic-torch-wheel-*
+          merge-multiple: true
+
       - name: Download sdists
         uses: actions/download-artifact@v4
         with:

python/featomic_torch/build-backend/backend.py

@@ -24,6 +24,12 @@
     # we are building from a sdist
     FEATOMIC_DEP = "featomic >=0.1.0.dev0,<0.2.0"
 
+FORCED_TORCH_VERSION = os.environ.get("FEATOMIC_TORCH_BUILD_WITH_TORCH_VERSION")
+if FORCED_TORCH_VERSION is not None:
+    TORCH_DEP = f"torch =={FORCED_TORCH_VERSION}"
+else:
+    TORCH_DEP = "torch >=1.12"
+
 
 get_requires_for_build_sdist = build_meta.get_requires_for_build_sdist
 prepare_metadata_for_build_wheel = build_meta.prepare_metadata_for_build_wheel
@@ -35,7 +41,7 @@ def get_requires_for_build_wheel(config_settings=None):
     defaults = build_meta.get_requires_for_build_wheel(config_settings)
     return defaults + [
         "cmake",
-        "torch >= 1.12",
+        TORCH_DEP,
         "metatensor-torch >=0.6.0,<0.7.0",
         FEATOMIC_DEP,
     ]

scripts/create-torch-versions-range.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+
+"""
+This script updates the `Requires-Dist` information in featomic-torch wheel METADATA to
+contain the range of compatible torch versions. It expects newline separated
+`Requires-Dist: torch ==...` information (corresponding to wheels built against a single
+torch version) and will print `Requires-Dist: torch >=$MIN_VERSION,<${MAX_VERSION+1}` on
+the standard output.
+
+This output can the be used in the merged wheel containing the build against all torch
+versions.
+"""
+
+import re
+import sys
+
+
+if __name__ == "__main__":
+    torch_versions_raw = sys.argv[1]
+
+    torch_versions = []
+    for version in torch_versions_raw.split("\n"):
+        if version.strip() == "":
+            continue
+
+        match = re.match(r"Requires-Dist: torch[ ]?==(\d+)\.(\d+)\.\*", version)
+        if match is None:
+            raise ValueError(f"unexpected Requires-Dist format: {version}")
+
+        major, minor = match.groups()
+        major = int(major)
+        minor = int(minor)
+
+        version = (major, minor)
+
+        if version in torch_versions:
+            raise ValueError(f"duplicate torch version: {version}")
+
+        torch_versions.append(version)
+
+    torch_versions = list(sorted(torch_versions))
+
+    min_version = f"{torch_versions[0][0]}.{torch_versions[0][1]}"
+    max_version = f"{torch_versions[-1][0]}.{torch_versions[-1][1] + 1}"
+
+    print(f"Requires-Dist: torch >={min_version},<{max_version}")