Support for Intel MPI installed in virtual environment

mpi4py · Oct 17, 2023 · bad9cf6 · bad9cf6
1 parent ee3b9ea
commit bad9cf6
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 24 deletions.
diff --git a/.cibw/mpi4py_mpiabi.py b/.cibw/mpi4py_mpiabi.py
@@ -14,14 +14,28 @@ def _verbose_info(message, verbosity=1):
         print(f"# [{__spec__.parent}] {message}", file=sys.stderr)
 
 
+_rpath_prefix = "@rpath" if sys.platform == "darwin" else ""
 _rpath_libmpi = []
+if sys.prefix != sys.base_prefix:
+    if sys.platform == "linux":
+        _rpath_libmpi.extend((
+            f"{sys.exec_prefix}/lib",
+            _rpath_prefix,
+        ))
+    if sys.platform == "win32":
+        _rpath_libmpi.extend((
+            f"{sys.exec_prefix}\\DLLs",
+            f"{sys.exec_prefix}\\Lib",
+            f"{sys.exec_prefix}\\Library\\bin",
+            _rpath_prefix,
+        ))
 if sys.platform == "darwin":
-    _rpath_libmpi.extend([
-        "@rpath",
+    _rpath_libmpi.extend((
+        _rpath_prefix,
         "/usr/local/lib",
         "/opt/homebrew/lib",
         "/opt/local/lib",
-    ])
+    ))
 
 
 def _dlopen_libmpi(libmpi=None):  # noqa: C901
@@ -57,7 +71,8 @@ def libmpi_paths(path):
         for entry in path:
             entry = os.path.expandvars(entry)
             entry = os.path.expanduser(entry)
-            if os.path.isdir(entry) or entry == "@rpath":
+            entry = os.path.normpath(entry)
+            if entry == _rpath_prefix or os.path.isdir(entry):
                 for name in libmpi_names():
                     yield os.path.join(entry, name)
             else:

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,49 @@ 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: venv
+      shell: bash
+      run: |
+        venvdir=/tmp/venv
+        python -m venv $venvdir
+        if [[ $(uname) =~ NT ]]; then
+          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"
+        $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 }}
 
-    - uses: actions/download-artifact@v3
+    - uses: mpi4py/setup-mpi@v1
       with:
-        name: wheel-${{ runner.os }}
-        path: dist
+        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