Support for Intel MPI installed in virtual envs and user site

mpi4py · Oct 18, 2023 · 75dab35 · 75dab35
1 parent ee3b9ea
commit 75dab35
Show file tree

Hide file tree

Showing 4 changed files with 153 additions and 36 deletions.
diff --git a/.cibw/mpi4py_mpiabi.py b/.cibw/mpi4py_mpiabi.py
@@ -14,14 +14,54 @@ def _verbose_info(message, verbosity=1):
         print(f"# [{__spec__.parent}] {message}", file=sys.stderr)
 
 
-_rpath_libmpi = []
-if sys.platform == "darwin":
-    _rpath_libmpi.extend([
-        "@rpath",
-        "/usr/local/lib",
-        "/opt/homebrew/lib",
-        "/opt/local/lib",
-    ])
+def _dlopen_rpath():  # noqa: C901
+    rpath = []
+
+    def add_rpath(*paths):
+        rpath.extend(paths)
+
+    def add_rpath_prefix(prefix):
+        if sys.platform == "linux":
+            add_rpath(os.path.join(prefix, "lib"))
+        if sys.platform == "win32":
+            add_rpath(os.path.join(prefix, "DLLs"))
+            add_rpath(os.path.join(prefix, "Library", "bin"))
+
+    site = sys.modules.get("site")
+    if site is not None and site.ENABLE_USER_SITE:
+        user_base = os.path.abspath(site.USER_BASE)
+        user_site = os.path.abspath(site.USER_SITE)
+        site_pkgs = os.path.commonpath((user_site, __file__))
+        if site_pkgs == user_site:
+            add_rpath_prefix(user_base)
+
+    if sys.exec_prefix != sys.base_exec_prefix:
+        add_rpath_prefix(sys.exec_prefix)
+
+    if sys.platform == "win32":
+        i_mpi_root = os.environ.get("I_MPI_ROOT")
+        if i_mpi_root:
+            library_kind = (
+                os.environ.get("I_MPI_LIBRARY_KIND") or
+                os.environ.get("library_kind") or
+                "release"
+            )
+            add_rpath(os.path.join(i_mpi_root, "bin", library_kind))
+            add_rpath(os.path.join(i_mpi_root, "bin"))
+        msmpi_bin = os.environ.get("MSMPI_BIN")
+        if msmpi_bin:
+            rpath.append(msmpi_bin)
+
+    add_rpath("")
+
+    if sys.platform == "darwin":
+        add_rpath(
+            "/usr/local/lib",
+            "/opt/homebrew/lib",
+            "/opt/local/lib",
+        )
+
+    return rpath
 
 
 def _dlopen_libmpi(libmpi=None):  # noqa: C901
@@ -54,10 +94,12 @@ def libmpi_names():
             yield "msmpi.dll"
 
     def libmpi_paths(path):
+        rpath = "@rpath" if sys.platform == "darwin" else ""
         for entry in path:
+            entry = entry or rpath
             entry = os.path.expandvars(entry)
             entry = os.path.expanduser(entry)
-            if os.path.isdir(entry) or entry == "@rpath":
+            if entry == rpath or os.path.isdir(entry):
                 for name in libmpi_names():
                     yield os.path.join(entry, name)
             else:
@@ -71,13 +113,9 @@ def libmpi_paths(path):
     if libmpi is not None:
         path = libmpi.split(os.pathsep)
     else:
-        path = _rpath_libmpi or None
-    if path is not None:
-        libmpi_iterable = libmpi_paths(path)
-    else:
-        libmpi_iterable = libmpi_names()
+        path = _libmpi_rpath or _dlopen_rpath() or [""]
     errors = ["cannot load MPI library"]
-    for filename in libmpi_iterable:
+    for filename in libmpi_paths(path):
         try:
             return dlopen(filename)
         except OSError as exc:
@@ -87,6 +125,9 @@ def libmpi_paths(path):
     raise RuntimeError("\n".join(errors))
 
 
+_libmpi_rpath = []
+
+
 def _get_mpiabi_from_libmpi(libmpi=None):
     # pylint: disable=import-outside-toplevel
     import ctypes as ct

diff --git a/.github/actions/mpi4py-test-basic/action.yml b/.github/actions/mpi4py-test-basic/action.yml
@@ -8,31 +8,39 @@ inputs:
     description: "shell"
     required: false
     default: 'bash'
+  python:
+    description: "python"
+    required: false
+    default: 'python'
+  mpiexec:
+    description: "mpiexec"
+    required: false
+    default: 'mpiexec'
 
 runs:
   using: 'composite'
   steps:
 
   - name:  Test mpi4py prefix
-    run:   python -m mpi4py --prefix
+    run:   ${{ inputs.python }} -m mpi4py --prefix
     shell: ${{ inputs.shell }}
 
   - name:  Test mpi4py version
-    run:   python -m mpi4py --version
+    run:   ${{ inputs.python }} -m mpi4py --version
     shell: ${{ inputs.shell }}
 
   - name:  Test mpi4py MPI standard version
-    run:   python -m mpi4py --mpi-std-version
+    run:   ${{ inputs.python }} -m mpi4py --mpi-std-version
     shell: ${{ inputs.shell }}
 
   - name:  Test mpi4py MPI library version
-    run:   python -m mpi4py --mpi-lib-version
+    run:   ${{ inputs.python }} -m mpi4py --mpi-lib-version
     shell: ${{ inputs.shell }}
 
   - name:  Test mpi4py helloworld
-    run:   mpiexec -n 2 python -m mpi4py.bench helloworld
+    run:   ${{ inputs.mpiexec }} -n 2 ${{ inputs.python }} -m mpi4py.bench helloworld
     shell: ${{ inputs.shell }}
 
   - name:  Test mpi4py ringtest
-    run:   mpiexec -n 2 python -m mpi4py.bench ringtest
+    run:   ${{ inputs.mpiexec }} -n 2 ${{ inputs.python }} -m mpi4py.bench ringtest
     shell: ${{ inputs.shell }}
diff --git a/.github/workflows/cd-wheel.yml b/.github/workflows/cd-wheel.yml
@@ -301,6 +301,11 @@ jobs:
 
     - uses: actions/checkout@v4
 
+    - uses: actions/download-artifact@v3
+      with:
+        name: wheel-${{ runner.os }}
+        path: dist
+
     - uses: mamba-org/setup-micromamba@v1
       with:
         init-shell: bash
@@ -341,11 +346,6 @@ jobs:
           ;;
         esac
 
-    - uses: actions/download-artifact@v3
-      with:
-        name: wheel-${{ runner.os }}
-        path: dist
-
     - run: python -m pip install mpi4py --no-index --find-links=dist
 
     - uses: ./.github/actions/mpi4py-test-basic
@@ -378,23 +378,81 @@ jobs:
 
     - uses: actions/checkout@v4
 
-    - uses: mpi4py/setup-mpi@v1
+    - uses: actions/download-artifact@v3
       with:
-        mpi: ${{ matrix.mpi }}
-
-    - if: ${{ matrix.mpi == 'mpich' && runner.os == 'Linux' }}
-      run: sudo ln -sr /usr/lib/$(arch)-linux-gnu/libmpi{ch,}.so.12
+        name: wheel-${{ runner.os }}
+        path: dist
 
     - uses:  actions/setup-python@v4
       with:
         python-version: ${{ matrix.py }}
 
-    - run: python -m pip install --upgrade pip
+    - if: ${{ matrix.mpi == 'impi' }}
+      id: user
+      shell: bash
+      run: |
+        userdir=$(python -m site --user-base)
+        if [[ $(uname) =~ NT ]]; then
+            userdir=$(cygpath -u "$userdir")
+            mpiexec="$userdir/Library/bin/mpiexec.exe"
+        else
+            mpiexec="$userdir/bin/mpiexec"
+        fi
+        echo "mpiexec=$mpiexec"    >> "$GITHUB_OUTPUT"
+        set -x
+        python -m pip install --user --upgrade pip
+        python -m pip install --user impi-rt
+        python -m pip install --user mpi4py --no-index --find-links=dist
+      env:
+        PYTHONUSERBASE: ${{ runner.temp }}/user
 
-    - uses: actions/download-artifact@v3
+    - if: ${{ matrix.mpi == 'impi' }}
+      uses: ./.github/actions/mpi4py-test-basic
+      timeout-minutes: 2
       with:
-        name: wheel-${{ runner.os }}
-        path: dist
+        mpiexec: ${{steps.user.outputs.mpiexec }}
+      env:
+        PYTHONUSERBASE: ${{ runner.temp }}/user
+        I_MPI_FABRICS: shm
+
+    - if: ${{ matrix.mpi == 'impi' }}
+      id: venv
+      shell: bash
+      run: |
+        venvdir="${{ runner.temp }}/venv"
+        python -m venv "$venvdir"
+        if [[ $(uname) =~ NT ]]; then
+            venvdir=$(cygpath -u "$venvdir")
+            python="$venvdir/Scripts/python.exe"
+            mpiexec="$venvdir/Library/bin/mpiexec.exe"
+        else
+            python="$venvdir/bin/python"
+            mpiexec="$venvdir/bin/mpiexec"
+        fi
+        echo "python=$python"   >> "$GITHUB_OUTPUT"
+        echo "mpiexec=$mpiexec" >> "$GITHUB_OUTPUT"
+        set -x
+        $python -m pip install --upgrade pip
+        $python -m pip install impi-rt
+        $python -m pip install mpi4py --no-index --find-links=dist
+
+    - if: ${{ matrix.mpi == 'impi' }}
+      uses: ./.github/actions/mpi4py-test-basic
+      timeout-minutes: 2
+      with:
+        python: ${{ steps.venv.outputs.python }}
+        mpiexec: ${{steps.venv.outputs.mpiexec }}
+      env:
+        I_MPI_FABRICS: shm
+
+    - uses: mpi4py/setup-mpi@v1
+      with:
+        mpi: ${{ matrix.mpi }}
+
+    - if: ${{ matrix.mpi == 'mpich' && runner.os == 'Linux' }}
+      run: sudo ln -sr /usr/lib/$(arch)-linux-gnu/libmpi{ch,}.so.12
+
+    - run: python -m pip install --upgrade pip
 
     - run: python -m pip install mpi4py --no-index --find-links=dist
 

diff --git a/README.md b/README.md
@@ -13,7 +13,7 @@ mpi4py wheels are uploaded to the [Anaconda.org](https://anaconda.org/mpi4py)
 package server. These wheels can be installed with `pip` specifying the
 alternative index URL:
 
-```
+```sh
 python -m pip install -i https://pypi.anaconda.org/mpi4py/simple mpi4py
 ```
 
@@ -25,6 +25,16 @@ Python virtual environments and use an externally-provided MPI runtime coming
 from the system package manager, sysadmin-maintained builds accessible via
 module files, or customized user builds.
 
+[Intel MPI](https://software.intel.com/intel-mpi-library) distributes [Linux
+and Windows wheels](https://pypi.org/project/impi-rt/#files) for Intel-based
+processor architectures (`x86_64`/`AMD64`). These wheels are hosted on
+[PyPI](https://pypi.org/project/impi-rt). mpi4py and Intel MPI wheels can be
+installed together with `pip` to get a ready-to-use Python+MPI environment:
+
+```sh
+python -m pip install impi-rt mpi4py \
+   --extra-index-url https://pypi.anaconda.org/mpi4py/simple
+```
 
 ## Linux (`x86_64`, `aarch64`, `ppc64le`):