HIP wheels not available via PyPI

[Lewih]

OpenMM-Torch-HIP-7 404 (not on PyPI)
OpenMM-Torch-HIP-6 404 (not on PyPI)
OpenMM-Torch-CUDA-12 OK 1.5

pip install openmmtorch[hip7] can't ever resolve as no HIP wheel is available.

[peastman]

I've yet to figure out how to get it to build against the HIP version of PyTorch. See #178 and pytorch/pytorch#181939 if you're curious and have a lot of time to spare!

[Lewih]

I was able to create a OpenMM Apptainer container starting from the official rocm/pytorch docker.

I built OpenMM-Torch from source using the hipcc compiler. Claude also suggested to install swig from pip instead of using the apt one.

Here the recipe, maybe it can be useful.

Bootstrap: docker
From: rocm/pytorch:rocm7.2.4_ubuntu24.04_py3.12_pytorch_release_2.10.0

# PLUMED is NOT built here: it links against the cluster's central 2.10.0,
# already visible inside the container via the site apptainer config. Set
# PLUMED_ROOT below to that path.
#
# Build command is: apptainer build  --bind /apps:/apps --fakeroot  /tmp/out.sif recipe.def 

%labels
    OpenMM        8.5.2
    OpenMM-ML     1.7
    mace-torch    0.3.15
    OpenMM-PLUMED 2.1 (built from source)
    OpenMM-Torch  1.5 (built from source)
    PLUMED        2.10.0 (host, bind-mounted)

%setup
    # Create directory for binding at build-time.
    mkdir -p "${APPTAINER_ROOTFS}/apps"

%post
    set -eux

    # ---- versions ----------------------------------------------------
    OPENMM_VERSION=8.5.2
    OPENMM_ML_VERSION=1.7
    OPENMM_PLUMED_VERSION=2.1
    OPENMM_TORCH_VERSION=1.5
    MACE_VERSION=0.3.15
    PLUMED_ROOT=/myplumed/installation
    # ------------------------------------------------------------------

    PY=$(command -v python3)

    # Match the HIP extra to the ROCm in this image (hip6 / hip7).
    ROCM_MAJOR=$(cut -d. -f1 /opt/rocm/.info/version)
    echo "ROCm major version: ${ROCM_MAJOR}"

    # Stop pip from replacing the image's ROCm torch with a CUDA wheel.
    # Remove the +rocm suffix to avoid pip's "not a supported wheel on this platform" error.
    printf 'torch==%s\n' "$($PY -c 'import torch; print(torch.__version__.split("+")[0])')" \
        > /opt/pip-constraints.txt
    export PIP_CONSTRAINT=/opt/pip-constraints.txt

    $PY -m pip install --no-cache-dir \
        "openmm[hip${ROCM_MAJOR}]==${OPENMM_VERSION}" \
        "openmmml==${OPENMM_ML_VERSION}" \
        "mace-torch==${MACE_VERSION}"

    # Fail the build now rather than on a GPU node in three weeks.
    $PY -c 'import torch, sys; sys.exit(0 if torch.version.hip else "torch is not a ROCm build")'

    # ---- OpenMM-PLUMED plugin (README steps 4-8) ----------------------
    # Gotchas encoded below, in case this ever breaks:
    #   * swig must be >=4.3 and from pip, not apt's 4.2: a mismatched SWIG
    #     runtime makes addForce(PlumedForce(...)) raise a TypeError.
    #   * the pip OpenMM wheel keeps its C++ SDK in a sibling OpenMM.libs/ dir.
    #   * the plugin's CMake looks for `python`, not `python3`.
    #   * its setup.py imports numpy, so pip build isolation must be off.

    # Install build tools and SWIG from PyPI, not apt.
    export DEBIAN_FRONTEND=noninteractive
    apt-get update
    apt-get install -y --no-install-recommends \
        build-essential cmake git ca-certificates
    # clean up apt cache to reduce image size; pip will install swig from PyPI
    rm -rf /var/lib/apt/lists/*
    $PY -m pip install --no-cache-dir 'swig>=4.3'

    # The plugin's CMake runs a bare `${PIP_EXECUTABLE} install .`, and pip 26
    # ignores PIP_NO_BUILD_ISOLATION, so force the flag through a wrapper.
    printf '#!/bin/sh\nexec %s -m pip "$@" --no-build-isolation\n' "$PY" \
        > /usr/local/bin/pip-nobi
    chmod +x /usr/local/bin/pip-nobi

    PLUMED_INC="${PLUMED_ROOT}/include/plumed"      # holds wrapper/Plumed.h
    PLUMED_LIB=$(dirname "$(ls "$PLUMED_ROOT"/lib64/libplumed.so "$PLUMED_ROOT"/lib/libplumed.so 2>/dev/null | head -1)")
    OPENMM_DIR=$($PY -c 'import openmm,os,glob;s=os.path.dirname(os.path.dirname(openmm.__file__));print(next(os.path.dirname(os.path.dirname(h)) for h in glob.glob(s+"/**/include/OpenMM.h",recursive=True)))')
    test -f "${PLUMED_INC}/wrapper/Plumed.h"
    test -f "${OPENMM_DIR}/include/OpenMM.h"
    echo "PLUMED_LIB=${PLUMED_LIB}  OPENMM_DIR=${OPENMM_DIR}"

    SRC=/opt/omm-plumed-build
    git clone --depth 1 -b "v${OPENMM_PLUMED_VERSION}" \
        https://github.com/openmm/openmm-plumed.git "$SRC"
    cmake -S "$SRC" -B "$SRC/build" \
        -DCMAKE_BUILD_TYPE=Release \
        -DCMAKE_INSTALL_PREFIX="${OPENMM_DIR}" \
        -DOPENMM_DIR="${OPENMM_DIR}" \
        -DPLUMED_INCLUDE_DIR="${PLUMED_INC}" \
        -DPLUMED_LIBRARY_DIR="${PLUMED_LIB}" \
        -DPYTHON_EXECUTABLE="${PY}" \
        -DPIP_EXECUTABLE=/usr/local/bin/pip-nobi \
        -DSWIG_EXECUTABLE="$(command -v swig)" \
        -DCMAKE_SHARED_LINKER_FLAGS="-Wl,-rpath,${PLUMED_LIB}"
    cmake --build "$SRC/build" -j "$(nproc)" --target install
    LD_LIBRARY_PATH="${OPENMM_DIR}/lib:${PLUMED_LIB}" \
        cmake --build "$SRC/build" --target PythonInstall
    rm -rf "$SRC"

    # ---- OpenMM-Torch (from source: PyPI has no HIP wheel) -------------
    # openmmtorch[hip7] resolves to OpenMM-Torch-HIP-7, which was never
    # published; the plain wheel is CUDA-linked (libnvrtc.so.12, libtorch_cuda).
    # Gotchas:
    #   * its CMakeLists overwrites CMAKE_PREFIX_PATH with PYTORCH_DIR, so HIP
    #     must be located via -DHIP_DIR, not CMAKE_PREFIX_PATH.
    #   * the HIP platform links libOpenMMHIP.so, shipped by the openmm-hip
    #     wheel into this same OpenMM.libs prefix -- nothing extra to install.
    #   * this setup.py imports no numpy, so pip-nobi is not needed here.
    SRC=/opt/omm-torch-build
    git clone --depth 1 -b "v${OPENMM_TORCH_VERSION}" \
        https://github.com/openmm/openmm-torch.git "$SRC"
    cmake -S "$SRC" -B "$SRC/build" \
        -DCMAKE_CXX_COMPILER=/opt/rocm/bin/hipcc \
        -DCMAKE_BUILD_TYPE=Release \
        -DCMAKE_INSTALL_PREFIX="${OPENMM_DIR}" \
        -DOPENMM_DIR="${OPENMM_DIR}" \
        -DPYTORCH_DIR="$($PY -c 'import torch; print(torch.utils.cmake_prefix_path)')" \
        -DHIP_DIR=/opt/rocm/lib/cmake/hip \
        -DNN_BUILD_HIP_LIB=ON \
        -DPYTHON_EXECUTABLE="${PY}" \
        -DSWIG_EXECUTABLE="$(command -v swig)"
    cmake --build "$SRC/build" -j "$(nproc)" --target install
    LD_LIBRARY_PATH="${OPENMM_DIR}/lib" \
        cmake --build "$SRC/build" --target PythonInstall
    rm -rf "$SRC"
    # Fail here if the HIP platform silently didn't get built.
    test -f "${OPENMM_DIR}/lib/plugins/libOpenMMTorchHIP.so"

    printf '%s\n' "${OPENMM_DIR}/lib" \
        "$($PY -c 'import torch,os;print(os.path.join(os.path.dirname(torch.__file__),"lib"))')" \
        > /etc/ld.so.conf.d/openmm.conf
    ldconfig
    # %environment runs at run time and can't see $PLUMED_ROOT, so bake it here.
    printf 'export PATH="%s/bin:$PATH"\n' "${PLUMED_ROOT}" \
        > /.singularity.d/env/91-plumed.sh

%environment
    export OPENMM_FFT_BACKEND=1
    # libOpenMMPlumed.so links libplumed.so directly and carries an RPATH to it,
    # so no LD_LIBRARY_PATH is needed. PLUMED_KERNEL is only for the standalone
    # plumed driver. PATH is set in 91-plumed.sh (written by %post).

%runscript
    if [ $# -eq 0 ]; then exec /bin/bash; else exec "$@"; fi

%test
    python3 -m openmm.testInstallation
    python3 -c "
import torch, openmm as mm
print('torch    ', torch.__version__, 'HIP', torch.version.hip)
print('openmm   ', mm.version.version)
print('platforms', [mm.Platform.getPlatform(i).getName() for i in range(mm.Platform.getNumPlatforms())])
print('failures ', list(mm.Platform.getPluginLoadFailures()) or 'none')
from openmmtorch import TorchForce
from openmmplumed import PlumedForce
import openmmml, mace
# exercise the plugin, not just the import -- this is what a stale SWIG breaks:
s = mm.System(); s.addParticle(1.0); s.addParticle(1.0)
s.addForce(PlumedForce('d: DISTANCE ATOMS=1,2'))
print('plugins  ok (forces =', s.getNumForces(), ')')
"