Binary PyPI wheels

.github/workflows/docs_build.yaml → .github/workflows/docs.yml

@@ -1,12 +1,8 @@
 name: Build Documentation
 
 on:
-  push:
-    branches:
-      - "main"
-  pull_request:
-    branches:
-      - "main"
+  workflow_dispatch:
+  workflow_call:
 
 jobs:
   build-docs:

.github/workflows/publish.yml

@@ -0,0 +1,186 @@
+name: Build wheels and publish
+
+on:
+  workflow_dispatch:
+  workflow_call:
+
+jobs:
+  build:
+    name: Build wheels on ${{ matrix.os }}-${{ matrix.cuda }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, ubuntu-24.04-arm, windows-2019, macos-latest]
+        cuda: [cpu, cu118, cu124]
+        exclude:
+          - os: ubuntu-24.04-arm
+            cuda: cu118
+          - os: ubuntu-24.04-arm
+            cuda: cu124
+          - os: macos-latest
+            cuda: cu118
+          - os: macos-latest
+            cuda: cu124
+
+    steps:
+      - name: Free space of Github Runner (otherwise it will fail by running out of disk)
+        if: matrix.os == 'ubuntu-latest'
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf "/usr/local/share/boost"
+          sudo rm -rf "/usr/local/.ghcup"
+          sudo rm -rf "/usr/local/julia1.9.2"
+          sudo rm -rf "/usr/local/lib/android"
+          sudo rm -rf "$AGENT_TOOLSDIRECTORY"
+
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.13"
+
+      - name: Install cibuildwheel
+        run: python -m pip install cibuildwheel==2.23.0
+
+      - name: Activate MSVC
+        uses: ilammy/msvc-dev-cmd@v1
+        if: matrix.os == 'windows-2019'
+
+      - name: Build wheels
+        if: matrix.os != 'windows-2019'
+        shell: bash -l {0}
+        run: |
+          if [[ "${{ runner.os }}" == "Linux" ]]; then
+            export CIBW_ARCHS="${{ matrix.os == 'ubuntu-24.04-arm' && 'aarch64' || 'auto64' }}"
+          fi
+          python -m cibuildwheel --output-dir wheelhouse
+        env:
+          CIBW_BEFORE_ALL_LINUX: "cibuildwheel_support/prepare_build_environment_linux_${{ matrix.cuda }}.sh"
+          CIBW_BEFORE_ALL_MACOS: "cibuildwheel_support/prepare_build_environment_osx.sh"
+          CPU_TRAIN: ${{ runner.os == 'macOS' && 'true' || 'false' }}
+          WITH_CUDA: ${{ matrix.cuda != 'cpu' && '1' || '0' }}
+
+      - name: Build wheels
+        if: matrix.os == 'windows-2019'
+        shell: cmd # Use cmd on Windows to avoid bash environment taking priority over MSVC variables
+        run: python -m cibuildwheel --output-dir wheelhouse
+        env:
+          CIBW_BEFORE_ALL_WINDOWS: "bash cibuildwheel_support/prepare_build_environment_windows_${{ matrix.cuda }}.sh"
+          WITH_CUDA: ${{ matrix.cuda != 'cpu' && '1' || '0' }}
+          DISTUTILS_USE_SDK: "1" # Windows requires this to use vc for building
+          SKIP_TORCH_COMPILE: "true"
+          CIBW_BUILD_VERBOSITY_WINDOWS: "1"
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: ${{ matrix.cuda }}-cibw-wheels-${{ matrix.os }}-${{ strategy.job-index }}
+          path: ./wheelhouse/*.whl
+
+  publish-to-accelera-pypi:
+    name: >-
+      Publish Python 🐍 distribution 📦 to Acellera PyPI
+    needs:
+      - build
+    runs-on: ubuntu-latest
+    permissions: # Needed for GCP authentication
+      contents: "read"
+      id-token: "write"
+    strategy:
+      fail-fast: false
+      matrix:
+        cuda: [cpu, cu118, cu124]
+
+    steps:
+      - uses: actions/checkout@v4 # Needed for GCP authentication for some reason
+
+      - name: Set up Cloud SDK
+        uses: google-github-actions/auth@v2
+        with:
+          workload_identity_provider: ${{ secrets.GCP_WORKLOAD_IDENTITY_PROVIDER }}
+          service_account: ${{ secrets.GCP_PYPI_SERVICE_ACCOUNT }}
+
+      - name: Download all the dists
+        uses: actions/download-artifact@v4
+        with:
+          pattern: "${{ matrix.cuda }}-cibw-wheels*"
+          path: dist/
+          merge-multiple: true
+
+      - name: Publish distribution 📦 to Acellera PyPI
+        run: |
+          pip install build twine keyring keyrings.google-artifactregistry-auth
+          pip install -U packaging
+          twine upload --repository-url https://us-central1-python.pkg.dev/pypi-packages-455608/${{ matrix.cuda }} dist/* --verbose --skip-existing
+
+  publish-to-official-pypi:
+    name: >-
+      Publish Python 🐍 distribution 📦 to PyPI
+    needs:
+      - build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/torchmd-net
+    permissions:
+      id-token: write # IMPORTANT: mandatory for trusted publishing
+
+    steps:
+      - name: Download all the dists
+        uses: actions/download-artifact@v4
+        with:
+          pattern: "cu124-cibw-wheels*"
+          path: dist/
+          merge-multiple: true
+
+      - name: Publish distribution 📦 to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          password: ${{ secrets.TMDNET_PYPI_API_TOKEN }}
+          skip_existing: true
+
+  github-release:
+    name: >-
+      Sign the Python 🐍 distribution 📦 with Sigstore
+      and upload them to GitHub Release
+    needs:
+      - build
+    runs-on: ubuntu-latest
+
+    permissions:
+      contents: write # IMPORTANT: mandatory for making GitHub Releases
+      id-token: write # IMPORTANT: mandatory for sigstore
+
+    steps:
+      - name: Download all the dists
+        uses: actions/download-artifact@v4
+        with:
+          path: dist/
+          merge-multiple: true
+
+      - name: Sign the dists with Sigstore
+        uses: sigstore/[email protected]
+        with:
+          inputs: >-
+            ./dist/*.whl
+
+      - name: Create GitHub Release
+        env:
+          GITHUB_TOKEN: ${{ github.token }}
+        run: >-
+          gh release create
+          "$GITHUB_REF_NAME"
+          --repo "$GITHUB_REPOSITORY"
+          --notes ""
+
+      - name: Upload artifact signatures to GitHub Release
+        env:
+          GITHUB_TOKEN: ${{ github.token }}
+        # Upload to GitHub Release using the `gh` CLI.
+        # `dist/` contains the built packages, and the
+        # sigstore-produced signatures and certificates.
+        run: >-
+          gh release upload
+          "$GITHUB_REF_NAME" dist/**
+          --repo "$GITHUB_REPOSITORY"

.github/workflows/CI.yml → .github/workflows/test.yml

@@ -1,12 +1,8 @@
-name: CI
+name: Test
 
 on:
-  push:
-    branches:
-      - "main"
-  pull_request:
-    branches:
-      - "main"
+  workflow_dispatch:
+  workflow_call:
   schedule:
     - cron: "0 0 * * 0" # Run every Sunday at 0:00
 
@@ -17,13 +13,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [
-            "ubuntu-latest",
-            "ubuntu-22.04-arm",
-            "macos-latest",
-            # "macos-13",
-            "windows-2022",
-          ]
+        os:
+          ["ubuntu-latest", "ubuntu-22.04-arm", "macos-latest", "windows-2022"]
         python-version: ["3.10"]
 
     defaults: # Needed for conda
@@ -56,7 +47,7 @@ jobs:
           elif [[ "${{ runner.os }}" == "Linux" ]]; then
             conda install gxx --channel conda-forge --override-channels
           elif [[ "${{ runner.os }}" == "macOS" ]]; then
-            conda install clangxx llvm-openmp pybind11 --channel conda-forge --override-channels
+            conda install "clangxx<20" "llvm-openmp<20" pybind11 --channel conda-forge --override-channels
           elif [[ "${{ runner.os }}" == "Windows" ]]; then
             conda install vc vc14_runtime vs2015_runtime --channel conda-forge --override-channels
           fi

.github/workflows/workflow.yml

@@ -0,0 +1,25 @@
+name: Whole
+
+on: [push]
+
+jobs:
+  test:
+    uses: ./.github/workflows/test.yml
+    secrets: inherit
+
+  publish:
+    if: startsWith(github.event.ref, 'refs/tags/v')
+    needs: test
+    uses: ./.github/workflows/publish.yml
+    secrets: inherit
+
+  docs:
+    if: startsWith(github.event.ref, 'refs/tags/')
+    needs: publish
+    uses: ./.github/workflows/docs.yml
+    secrets: inherit
+
+  docs_update:
+    if: startsWith(github.event.ref, 'refs/tags/d')
+    uses: ./.github/workflows/docs.yml
+    secrets: inherit

MANIFEST.in

@@ -0,0 +1,9 @@
+prune benchmarks
+prune cibuildwheel_support
+prune docs
+prune examples
+prune scripts
+prune tests
+exclude .readthedocs.yaml
+exclude environment.yml
+exclude Makefile

Makefile

@@ -1,10 +1,24 @@
-ifndef version
-$(error version variable is not set. Call with `make release version=XXX`)
-endif
+# Check that given variables are set and all have non-empty values,
+# die with an error otherwise.
+#
+# Params:
+#   1. Variable name(s) to test.
+#   2. (optional) Error message to print.
+check_defined = \
+    $(strip $(foreach 1,$1, \
+        $(call __check_defined,$1,$(strip $(value 2)))))
+__check_defined = \
+    $(if $(value $1),, \
+      $(error Undefined variable: $1$(if $2, ($2))))
+
+.PHONY: release docs
 
 release:
-	git checkout main
-	git fetch
-	git pull
-	git tag -a $(version) -m "$(version) release"
-	git push --tags origin $(version)
+	@:$(call check_defined, version, The release version)
+	git tag -a v$(version) -m "v$(version) release"
+	git push --tags origin v$(version)
+
+docs:
+	@:$(call check_defined, version, The release version)
+	git tag -fa d$(version) -m "$(version) docs"
+	git push -f --tags origin d$(version)

README.md

@@ -21,11 +21,30 @@ Documentation is available at https://torchmd-net.readthedocs.io
 
 
 ## Installation  
-TorchMD-Net is available in [conda-forge](https://conda-forge.org/) and can be installed with:  
+TorchMD-Net is available as a pip installable wheel as well as in [conda-forge](https://conda-forge.org/)
+
+TorchMD-Net provides builds for CPU-only, CUDA 11.8 and CUDA 12.4. CPU versions are only provided as reference,
+as the performance will be extremely limited.
+Depending on which variant you wish to install, you can install it with one of the following commands:
+
+```sh
+# The following will install the CUDA 12.4 version by default
+pip install torchmd-net 
+# The following will install the CUDA 11.8 version
+pip install torchmd-net --extra-index-url https://download.pytorch.org/whl/cu118 --extra-index-url https://us-central1-python.pkg.dev/pypi-packages-455608/cu118/simple
+# The following will install the CUDA 12.4 version
+pip install torchmd-net --extra-index-url https://download.pytorch.org/whl/cu124 --extra-index-url https://us-central1-python.pkg.dev/pypi-packages-455608/cu124/simple
+# The following will install the CPU only version (not recommended)
+pip install torchmd-net --extra-index-url https://download.pytorch.org/whl/cpu --extra-index-url https://us-central1-python.pkg.dev/pypi-packages-455608/cpu/simple   
+```
+
+Alternatively it can be installed with conda or mamba with one of the following commands.
+We recommend using [Miniforge](https://github.com/conda-forge/miniforge/) instead of anaconda.
+
 ```shell
-mamba install torchmd-net
+mamba install torchmd-net cuda-version=11.8
+mamba install torchmd-net cuda-version=12.4
 ```
-We recommend using [Mamba](https://github.com/conda-forge/miniforge/#mambaforge) instead of conda.  
 
 ### Install from source  
 

cibuildwheel_support/prepare_build_environment_linux_cpu.sh

@@ -0,0 +1,4 @@
+# Configure pip to use PyTorch extra-index-url for CPU
+mkdir -p $HOME/.config/pip
+echo "[global]
+extra-index-url = https://download.pytorch.org/whl/cpu" > $HOME/.config/pip/pip.conf

cibuildwheel_support/prepare_build_environment_linux_cu118.sh

@@ -0,0 +1,37 @@
+#! /bin/bash
+
+set -e
+set -x
+
+
+if [ "$CIBW_ARCHS" != "aarch64" ]; then
+
+    # Install CUDA 11.8:
+    dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo
+
+    dnf install --setopt=obsoletes=0 -y \
+        cuda-nvcc-11-8-11.8.89-1 \
+        cuda-cudart-devel-11-8-11.8.89-1 \
+        libcurand-devel-11-8-10.3.0.86-1 \
+        libcudnn9-devel-cuda-11-9.8.0.87-1 \
+        libcublas-devel-11-8-11.11.3.6-1 \
+        libnccl-devel-2.15.5-1+cuda11.8 \
+        libcusparse-devel-11-8-11.7.5.86-1 \
+        libcusolver-devel-11-8-11.4.1.48-1 \
+        gcc-toolset-11
+
+    ln -s cuda-11.8 /usr/local/cuda
+    ln -s /opt/rh/gcc-toolset-11/root/usr/bin/gcc /usr/local/cuda/bin/gcc
+    ln -s /opt/rh/gcc-toolset-11/root/usr/bin/g++ /usr/local/cuda/bin/g++
+
+    export CUDA_HOME="/usr/local/cuda"
+    export WITH_CUDA=1
+
+
+    # Configure pip to use PyTorch extra-index-url for CUDA 11.8
+    mkdir -p $HOME/.config/pip
+    echo "[global]
+extra-index-url = https://download.pytorch.org/whl/cu118" > $HOME/.config/pip/pip.conf
+
+fi
+