build-wheels.yml

name: Build Python wheels

on:
  push:
    tags: ["*"]
  pull_request:
    paths:
      # build wheels in PR if this file changed
      - '.github/workflows/build-wheels.yml'
      # build wheels in PR if any of the build system files changed
      - '**/VERSION'
      - '**/setup.py'
      - '**/pyproject.toml'
      - '**/MANIFEST.in'
      - '**/Cargo.toml'
      - '**/CMakeLists.txt'
      - '**/build.rs'
  schedule:
    # check the build once a week on mondays
    - cron: '0 10 * * 1'


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

env:
  METATENSOR_NO_LOCAL_DEPS: "1"


jobs:
  build-core-wheels:
    runs-on: ${{ matrix.os }}
    name: ${{ matrix.name }}
    strategy:
      matrix:
        include:
          - name: x86_64 Linux
            os: ubuntu-22.04
            rust-target: x86_64-unknown-linux-gnu
            cibw-arch: x86_64
          - name: arm64 Linux
            os: ubuntu-22.04
            rust-target: aarch64-unknown-linux-gnu
            cibw-arch: aarch64
          - name: x86_64 macOS
            os: macos-12
            rust-target: x86_64-apple-darwin
            cibw-arch: x86_64
          - name: arm64 macOS
            os: macos-14
            rust-target: aarch64-apple-darwin
            cibw-arch: arm64
          - name: x86_64 Windows
            os: windows-2019
            # TODO: add a 32-bit windows builder?
            rust-target: x86_64-pc-windows-msvc
            cibw-arch: AMD64
    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: "3.12"

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

      - 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 build -t rustc-manylinux2014_${{ matrix.cibw-arch }} python/scripts/rustc-manylinux2014_${{ matrix.cibw-arch }}

      - name: build metatensor-core wheel
        run: python -m cibuildwheel python/metatensor-core
        env:
          CIBW_BUILD: cp311-*
          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
          CIBW_ENVIRONMENT: >
            METATENSOR_NO_LOCAL_DEPS=1
            MACOSX_DEPLOYMENT_TARGET=11

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

      - uses: actions/upload-artifact@v4
        with:
          name: core-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: ['1.12', '1.13', '2.0', '2.1', '2.2', '2.3', '2.4']
        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 arm64 on macos is only supported for torch >= 2.0
          - {os: macos-14, arch: arm64, torch-version: '1.12'}
          - {os: macos-14, arch: arm64, torch-version: '1.13'}
          # 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'}
        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: '1.12', python-version: '3.10', cibw-python: 'cp310-*'}
          - {torch-version: '1.13', python-version: '3.10', cibw-python: 'cp310-*'}
          - {torch-version: '2.0',  python-version: '3.11', cibw-python: 'cp311-*'}
          - {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-*'}
    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 metatensor-torch wheel
        run: python -m cibuildwheel python/metatensor-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
          # METATENSOR_NO_LOCAL_DEPS is set to 1 when building a tag of
          # metatensor-torch, which will force to use the version of
          # metatensor-core already released on PyPI. Otherwise, this will use
          # the version of metatensor-core from git checkout (in case there are
          # unreleased breaking changes). This means we can not release breaking
          # changes in metatensor-core and metatensor-torch by putting a tag on
          # the same commit. Instead metatensor-core must be fully released
          # before we start the build of metatensor-torch wheels.
          CIBW_ENVIRONMENT: >
            METATENSOR_NO_LOCAL_DEPS=${{ startsWith(github.ref, 'refs/tags/metatensor-torch-v') && '1' || '0' }}
            METATENSOR_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 or libmetatensor.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 libmetatensor.so --exclude libtorch.so --exclude libtorch_cpu.so --exclude libc10.so -w {dest_dir} {wheel}

      - uses: actions/upload-artifact@v4
        with:
          name: 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 {} "metatensor_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/metatensor_torch-*/metatensor_torch-*.dist-info/METADATA

          echo "\n\n METADATA = \n\n"
          cat dist/unpacked/metatensor_torch-*/metatensor_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/metatensor_torch-*/metatensor_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: torch-wheel-${{ matrix.os }}-${{ matrix.arch }}
          path: ./wheelhouse/*.whl

  build-others:
    name: Build other wheels/sdists
    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.12"

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

      - name: build metatensor-core sdist
        run: python -m build python/metatensor-core --sdist --outdir=dist/

      - name: build metatensor-operations sdist and wheel
        run: python -m build python/metatensor-operations --outdir=dist/

      - name: build metatensor-learn sdist and wheel
        run: python -m build python/metatensor-learn --outdir=dist/

      - name: build metatensor-torch sdist
        run: python -m build python/metatensor-torch --sdist --outdir=dist/

      - name: build metatensor sdist and wheel
        run: python -m build . --outdir=dist/

      - name: check sdist and wheels with twine
        run: twine check dist/*.tar.gz dist/*.whl

      - name: create C++ tarballs
        run: |
          ./scripts/package-core.sh dist/cxx/
          ./scripts/package-torch.sh dist/cxx/

      - uses: actions/upload-artifact@v4
        with:
          name: others
          path: |
            dist/*.tar.gz
            dist/cxx/*.tar.gz
            dist/*.whl

  merge-and-release:
    name: Merge and release wheels/sdists
    needs: [build-core-wheels, merge-torch-wheels, build-others]
    runs-on: ubuntu-22.04
    steps:
      - name: Download metatensor-core wheels
        uses: actions/download-artifact@v4
        with:
          path: wheels
          pattern: core-wheel-*
          merge-multiple: true

      - name: Download metatensor-torch wheels
        uses: actions/download-artifact@v4
        with:
          path: wheels
          pattern: torch-wheel-*
          merge-multiple: true

      - name: Download other wheels and sdists
        uses: actions/download-artifact@v4
        with:
          path: wheels
          name: others

      - name: Re-upload a single wheels artifact
        uses: actions/upload-artifact@v4
        with:
          name: wheels
          path: |
            wheels/*
            wheels/cxx/*

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

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

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

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

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