Acellera · stefdoerr · Apr 8, 2025 · Apr 8, 2025 · Apr 17, 2025 · Apr 17, 2025
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
@@ -0,0 +1,191 @@
+name: Build wheels and publish
+
+on:
+  workflow_dispatch:
+  workflow_call:
+
+jobs:
+  build:
+    name: Build wheels on ${{ matrix.os }}-${{ matrix.accelerator }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, ubuntu-24.04-arm, windows-2022, macos-latest]
+        accelerator: [cpu, cu118, cu126] #, cu128]
+        exclude:
+          - os: ubuntu-24.04-arm
+            accelerator: cu118
+          - os: ubuntu-24.04-arm
+            accelerator: cu126
+          # - os: ubuntu-24.04-arm
+          #   accelerator: cu128
+          - os: macos-latest
+            accelerator: cu118
+          - os: macos-latest
+            accelerator: cu126
+          # - os: macos-latest
+          #   accelerator: cu128
+
+    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
+        with:
+          toolset: 14.29
+        if: matrix.os == 'windows-2022'
+
+      - name: "Install dependencies with choco"
+        uses: crazy-max/ghaction-chocolatey@v3
+        if: matrix.os == 'windows-2022'
+        with:
+          args: install jom
+
+      - name: Build wheels
+        if: matrix.os != 'windows-2022'
+        shell: bash -l {0}
+        run: python -m cibuildwheel python/ --output-dir wheelhouse
+        env:
+          ACCELERATOR: ${{ matrix.accelerator }}
+
+      - name: Build wheels
+        if: matrix.os == 'windows-2022'
+        shell: cmd # Use cmd on Windows to avoid bash environment taking priority over MSVC variables
+        run: python -m cibuildwheel python/ --output-dir wheelhouse
+        env:
+          DISTUTILS_USE_SDK: "1" # Windows requires this to use vc for building
+          ACCELERATOR: ${{ matrix.accelerator }}
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: ${{ matrix.accelerator }}-cibw-wheels-${{ matrix.os }}-${{ strategy.job-index }}
+          path: ./wheelhouse/*.whl
+
+  publish-to-public-pypi:
+    name: >-
+      Publish Python 🐍 distribution 📦 to PyPI
+    needs:
+      - build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+    permissions:
+      id-token: write # IMPORTANT: mandatory for trusted publishing
+    strategy:
+      fail-fast: false
+      matrix:
+        accelerator: [cpu, cu118, cu126] #, hip, cu124, cu126, cu128]
+
+    steps:
+      - name: Download all the dists
+        uses: actions/download-artifact@v4
+        with:
+          pattern: "${{ matrix.accelerator }}-cibw-wheels*"
+          path: dist/
+          merge-multiple: true
+
+      - name: Publish distribution 📦 to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          password: ${{ secrets.PUBLIC_PYPI_API_TOKEN }}
+          skip-existing: true
+
+  # 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:
+  #       accelerator: [cpu, cu118, cu126] #, cu128]
+
+  #   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.accelerator }}-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.accelerator }} dist/* --verbose --skip-existing
+
+  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/gh-action-sigstore-python@v3.0.0
+        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"
diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
@@ -0,0 +1,9 @@
+name: Whole
+
+on: [push]
+
+jobs:
+  publish:
+    if: startsWith(github.event.ref, 'refs/tags/v')
+    uses: ./.github/workflows/publish.yml
+    secrets: inherit
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -92,6 +92,7 @@ SET_TARGET_PROPERTIES(${SHARED_NN_TARGET}
     LINK_FLAGS "${EXTRA_COMPILE_FLAGS}")
 TARGET_LINK_LIBRARIES(${SHARED_NN_TARGET} OpenMM)
 TARGET_LINK_LIBRARIES(${SHARED_NN_TARGET} "${TORCH_LIBRARIES}")
+SET_TARGET_PROPERTIES(${SHARED_NN_TARGET} PROPERTIES INSTALL_RPATH "$ORIGIN;$ORIGIN/../../torch/lib;$ORIGIN/../../nvidia/cuda_nvrtc/lib")
 INSTALL_TARGETS(/lib RUNTIME_DIRECTORY /lib ${SHARED_NN_TARGET})
 
 # install headers

diff --git a/platforms/cuda/CMakeLists.txt b/platforms/cuda/CMakeLists.txt
@@ -76,6 +76,7 @@ IF (APPLE)
     SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES LINK_FLAGS "-F/Library/Frameworks -framework CUDA ${EXTRA_COMPILE_FLAGS}")
 ENDIF (APPLE)
 
+SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES INSTALL_RPATH "$ORIGIN;$ORIGIN/..;$ORIGIN/../../../nvidia/cuda_runtime/lib/;$ORIGIN/../../../torch/lib;")
 INSTALL(TARGETS ${SHARED_TARGET} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/plugins)
 
 SUBDIRS (tests)
diff --git a/platforms/hip/CMakeLists.txt b/platforms/hip/CMakeLists.txt
@@ -60,6 +60,7 @@ SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES
     COMPILE_FLAGS "-DOPENMM_BUILDING_SHARED_LIBRARY ${EXTRA_COMPILE_FLAGS}"
     LINK_FLAGS "${EXTRA_COMPILE_FLAGS}")
 
+SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES INSTALL_RPATH "$ORIGIN;$ORIGIN/..;$ORIGIN/../../../torch/lib;")
 INSTALL(TARGETS ${SHARED_TARGET} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/plugins)
 
 SUBDIRS (tests)

diff --git a/platforms/opencl/CMakeLists.txt b/platforms/opencl/CMakeLists.txt
@@ -60,6 +60,7 @@ SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES
     COMPILE_FLAGS "-DOPENMM_BUILDING_SHARED_LIBRARY ${EXTRA_COMPILE_FLAGS}"
     LINK_FLAGS "${EXTRA_COMPILE_FLAGS}")
 
+SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES INSTALL_RPATH "$ORIGIN;$ORIGIN/..;$ORIGIN/../../../torch/lib;")
 INSTALL(TARGETS ${SHARED_TARGET} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/plugins)
 
 SUBDIRS (tests)

diff --git a/platforms/reference/CMakeLists.txt b/platforms/reference/CMakeLists.txt
@@ -46,6 +46,7 @@ SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES
     COMPILE_FLAGS "-DOPENMM_BUILDING_SHARED_LIBRARY ${EXTRA_COMPILE_FLAGS}"
     LINK_FLAGS "${EXTRA_COMPILE_FLAGS}")
 
+SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES INSTALL_RPATH "$ORIGIN;$ORIGIN/..;$ORIGIN/../../../torch/lib;")
 INSTALL(TARGETS ${SHARED_TARGET} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/plugins)
 
 SUBDIRS (tests)
diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
@@ -24,7 +24,7 @@ add_custom_target(PythonInstall DEPENDS "${WRAP_FILE}"  "${CMAKE_CURRENT_SOURCE_
 set(NN_PLUGIN_HEADER_DIR "${CMAKE_SOURCE_DIR}/openmmapi/include")
 set(NN_PLUGIN_LIBRARY_DIR "${CMAKE_BINARY_DIR}")
 configure_file(${CMAKE_CURRENT_SOURCE_DIR}/setup.py ${CMAKE_CURRENT_BINARY_DIR}/setup.py)
-add_custom_command(TARGET PythonInstall
-    COMMAND "${PYTHON_EXECUTABLE}" -m pip install .
-    WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}"
-)
+# add_custom_command(TARGET PythonInstall
+#     COMMAND "${PYTHON_EXECUTABLE}" -m pip install .
+#     WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}"
+# )
diff --git a/python/MANIFEST.in b/python/MANIFEST.in
@@ -0,0 +1,4 @@
+graft openmm/
+graft openmmtorch/
+prune tests
+prune cibuildwheel_support
diff --git a/python/cibuildwheel_support/before_all_linux.sh b/python/cibuildwheel_support/before_all_linux.sh
@@ -0,0 +1,123 @@
+#! /bin/bash
+
+set -e
+set -x
+
+# Install dependencies with yum
+dnf install -y zip opencl-headers ocl-icd tree
+
+# Check if we are running on aarch64
+if [ "$(uname -m)" == "aarch64" ]; then
+    dnf install -y libgfortran
+    ln -s /usr/lib64/libgfortran.so.5.0.0 /usr/lib64/libgfortran-0b50f350.so.5.0.0
+fi
+
+# 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
+
+if [ "$ACCELERATOR" == "cu118" ]; 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-compiler-11-8-11.8.0-1 \
+        cuda-libraries-11-8-11.8.0-1 \
+        cuda-libraries-devel-11-8-11.8.0-1 \
+        cuda-toolkit-11-8-11.8.0-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++
+    ln -s /usr/local/cuda/targets/x86_64-linux/lib/stubs/libcuda.so /usr/lib/libcuda.so.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
+
+    pip install openmm-unofficial-cu11 torch==2.7.1
+elif [ "$ACCELERATOR" == "cu126" ]; then
+    # Install CUDA 12.6
+    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-compiler-12-6-12.6.3-1\
+        cuda-libraries-12-6-12.6.3-1 \
+        cuda-libraries-devel-12-6-12.6.3-1 \
+        cuda-toolkit-12-6-12.6.3-1 \
+        gcc-toolset-13
+
+    ln -s cuda-12.6 /usr/local/cuda
+    ln -s /opt/rh/gcc-toolset-13/root/usr/bin/gcc /usr/local/cuda/bin/gcc
+    ln -s /opt/rh/gcc-toolset-13/root/usr/bin/g++ /usr/local/cuda/bin/g++
+    ln -s /usr/local/cuda/targets/x86_64-linux/lib/stubs/libcuda.so /usr/lib/libcuda.so.1
+
+    # Configure pip to use PyTorch extra-index-url for CUDA 12.6
+    mkdir -p $HOME/.config/pip
+    echo "[global]
+extra-index-url = https://download.pytorch.org/whl/cu126" > $HOME/.config/pip/pip.conf
+
+    pip install openmm-unofficial-cu12 torch==2.7.1
+elif [ "$ACCELERATOR" == "hip" ]; then
+    # Install HIP 6.2
+    dnf install -y https://repo.radeon.com/amdgpu-install/6.2.2/el/8.10/amdgpu-install-6.2.60202-1.el8.noarch.rpm
+    dnf install -y rocm-device-libs hip-devel hip-runtime-amd hipcc
+    pip install openmm-unofficial-cpu torch==2.7.1
+else
+    pip install openmm-unofficial-cpu torch==2.7.1
+fi
+
+#################
+SITE_PACKAGES=$(python -c 'import site; print(site.getsitepackages()[0])')
+
+CMAKE_FLAGS=""
+if [ "$ACCELERATOR" == "cu118" ] || [ "$ACCELERATOR" == "cu126" ]; then
+    ARCH_LIST=$(python -c "import torch; print(';'.join([f'{y[:-1]}.{y[-1]}' for y in [x[3:] for x in torch._C._cuda_getArchFlags().split() if x.startswith('sm_')]]))")
+    # CMakeLists.txt seems to ignore the CMAKE_CUDA_ARCHITECTURES variable, instead, it is overwritten by TORCH_CUDA_ARCH_LIST
+    ARCH_LIST_FMT=$(python -c "import torch; print(';'.join([y for y in [x[3:] for x in torch._C._cuda_getArchFlags().split() if x.startswith('sm_')]]))")
+    export CUDA_HOME="/usr/local/cuda"
+    CMAKE_FLAGS="    -DTORCH_CUDA_ARCH_LIST=${ARCH_LIST}"
+    CMAKE_FLAGS+="    -DCMAKE_CUDA_ARCHITECTURES=${ARCH_LIST_FMT}"
+    CMAKE_FLAGS+="    -DCUDA_TOOLKIT_ROOT_DIR=${CUDA_HOME}"
+    CMAKE_FLAGS+="    -DCMAKE_CUDA_COMPILER=${CUDA_HOME}/bin/nvcc"
+fi
+
+if [[ "$OSTYPE" == "linux-gnu"* ]]; then
+    CMAKE_FLAGS+=" -DOPENCL_INCLUDE_DIR=/usr/include/CL"
+    CMAKE_FLAGS+=" -DOPENCL_LIBRARY=/usr/lib64/libOpenCL.so.1"
+fi
+
+# Configure build with Cmake
+mkdir -p build
+mkdir -p install
+cd build
+
+echo $CMAKE_FLAGS
+export LD_LIBRARY_PATH=/usr/lib/:/usr/local/cuda/targets/x86_64-linux/lib/:$LD_LIBRARY_PATH
+
+cmake .. \
+    -DCMAKE_INSTALL_PREFIX=../install \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DOPENMM_DIR=${SITE_PACKAGES}/openmm \
+    -DPYTORCH_DIR=${SITE_PACKAGES}/torch \
+    -DTorch_DIR=${SITE_PACKAGES}/torch/share/cmake/Torch \
+    -DNN_BUILD_OPENCL_LIB=ON \
+    ${CMAKE_FLAGS}
+
+# Build OpenMMTorch
+make -j4 install
+make -j4 PythonInstall
+
+cd ..
+
+cp build/python/setup.py python/
+cp build/python/openmmtorch.py python/openmmtorch/
+cp build/python/TorchPluginWrapper.cpp python/openmmtorch/
+cp -r install/include python/openmmtorch/
+
+# Copy the libraries of openmm
+mkdir -p python/openmm/
+cp -r install/lib python/openmm/lib
diff --git a/python/cibuildwheel_support/before_all_osx.sh b/python/cibuildwheel_support/before_all_osx.sh
@@ -0,0 +1,13 @@
+#! /bin/bash
+
+set -e
+set -x
+
+brew install swig tree
+
+# 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
+
+cp python/setup.py $HOME/setup.py.bkp