Support for OFI and UCX

* MPICH: Configure --with-device=ch4:ofi,ucx
* Open MPI: Configure --with-ofi --with-ucx
* Remove references to source/build directories
* Strip binaries to reduce size
* Remove unneeded binaries

Author: dalcinl

.github/workflows/cd-wheel.yml

@@ -76,6 +76,7 @@ jobs:
       shell: python
       name: 'setup build matrix'
       run: |
+        # setup build matrix # "
         keys = ("os", "arch", "runner")
         rows = [
             ("Linux", "aarch64", "ubuntu-22.04"),
@@ -93,6 +94,7 @@ jobs:
         import os, json
         with open(os.getenv("GITHUB_OUTPUT"), "w") as out:
             print(f"matrix={json.dumps(matrix)}", file=out)
+        # "
 
   build:
     needs: setup
@@ -117,12 +119,12 @@ jobs:
         cd $(brew --prefix)/bin
         gfortran=$(ls gfortran-* | sort | head -n 1)
         sudo ln -s $gfortran gfortran
+        # unlink libevent
+        brew unlink libevent || true
         # install autotools
         brew install autoconf
         brew install automake
         brew install libtool
-        # unlink libevent
-        brew unlink libevent || true
         # install uv
         brew install uv
 
@@ -142,9 +144,11 @@ jobs:
 
     - id: source-date-epoch
       run: |
-        SOURCE_DATE_EPOCH=$(git log -1 --pretty=%ct)
+        sde=$(cat source-date-epoch 2>/dev/null || true)
+        SOURCE_DATE_EPOCH=${sde:-$(git log -1 --pretty=%ct)}
         echo SOURCE_DATE_EPOCH=$SOURCE_DATE_EPOCH >> $GITHUB_ENV
-        echo $(git log -1 --pretty=%ci) [timestamp=$SOURCE_DATE_EPOCH]
+        echo $(date -d @$SOURCE_DATE_EPOCH) [timestamp=$SOURCE_DATE_EPOCH]
+      working-directory: package/source
 
     - id: build
       uses: pypa/cibuildwheel@v2.21.3
@@ -174,7 +178,7 @@ jobs:
         CIBW_REPAIR_WHEEL_COMMAND_MACOS: >
           delocate-wheel
           --ignore-missing-dependencies
-          --exclude libmpi --exclude libpmpi
+          --exclude libmpi --exclude libpmpi --exclude libfabric
           --require-archs {delocate_archs}
           -w {dest_dir} -v {wheel}
 

.gitignore

@@ -5,4 +5,5 @@
 /venv/
 *.egg-info/
 *.tar.gz
+*.tar.bz2
 *~

Makefile

@@ -10,10 +10,11 @@ lint:
 	yamllint .github/
 
 clean:
+	$(RM) -r package/METADATA
+	$(RM) -r package/LICENSE*
 	$(RM) -r package/build
-	$(RM) -r package/LICENSE
-	$(RM) -r package/install
 	$(RM) -r package/source
 	$(RM) -r package/workdir
+	$(RM) -r package/install
 	$(RM) -r package/*.egg-info
 	$(RM) -r .*_cache

bootstrap.sh

@@ -8,6 +8,11 @@ case "$mpiname" in
 esac
 version=${VERSION:-$version}
 
+ucxversion=1.15.0
+ofiversion=1.21.1
+ucxversion=${UCXVERSION:-$ucxversion}
+ofiversion=${OFIVERSION:-$ofiversion}
+
 PROJECT=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
 PACKAGE=$PROJECT/package
 SOURCE=$PACKAGE/source
@@ -27,7 +32,7 @@ fi
 if test ! -d "$SOURCE"; then
     if test ! -f "$tarball"; then
         echo downloading "$urlbase"/"$tarball"...
-        curl -fsO "$urlbase"/"$tarball"
+        curl -fsSLO "$urlbase"/"$tarball"
     else
         echo reusing "$tarball"...
     fi
@@ -38,8 +43,32 @@ if test ! -d "$SOURCE"; then
     if test -f "$patch"; then
         patch -p1 -i "$patch" -d "$SOURCE"
     fi
-    if test "$mpiname-$(uname)" = "openmpi-Darwin"; then
-        if test -d "$SOURCE"/3rd-party; then
+    if test "$mpiname" = "mpich"; then
+        if test "${version}" \> "4.2.2" && test "$(uname)" = "Linux"; then
+            conf_args='s/\(ofi_subdir_args\)=\("--enable-embedded"\)/\1=""/'
+            configure="$SOURCE"/configure
+            sed -i.orig "$conf_args" "$configure"
+        fi
+        if test "${version}" \< "4.2.0"; then
+            disable_doc='s/^\(install-data-local:\s\+\)\$/\1#\$/'
+            Makefile="$SOURCE"/Makefile.in
+            sed -i.orig "$disable_doc" "$Makefile"
+        fi
+    fi
+    if test "$mpiname" = "openmpi"; then
+        makefiles=(
+            "$SOURCE"/3rd-party/openpmix/src/util/keyval/Makefile.in
+            "$SOURCE"/3rd-party/prrte/src/mca/rmaps/rank_file/Makefile.in
+            "$SOURCE"/3rd-party/prrte/src/util/hostfile/Makefile.in
+        )
+        for makefile in "${makefiles[@]}"; do
+            test -f "$makefile" || continue
+            echo "PMIX_CFLAGS_BEFORE_PICKY = @CFLAGS@" >> "$makefile"
+            echo "PRTE_CFLAGS_BEFORE_PICKY = @CFLAGS@" >> "$makefile"
+        done
+    fi
+    if test "$mpiname" = "openmpi" && test "${version}" \< "5.0.5"; then
+        if test -d "$SOURCE"/3rd-party && test "$(uname)" = "Darwin"; then
             cd "$SOURCE"/3rd-party
             tar xf libevent-*.tar.gz && cd libevent-*
             echo running autogen.sh on "$(basename "$(pwd)")"
@@ -50,5 +79,74 @@ if test ! -d "$SOURCE"; then
 else
     echo reusing directory "$SOURCE"...
 fi
-echo copying license file...
+
+if test "$(uname)" = "Linux"; then
+    if test "$mpiname" = "mpich"; then
+        ofidestdir="$SOURCE"/modules/libfabric
+        ucxdestdir="$SOURCE"/modules/ucx
+    fi
+    if test "$mpiname" = "openmpi"; then
+        ofigithub="https://github.com/ofiwg/libfabric"
+        ofiurlbase="$ofigithub/releases/download/v$ofiversion"
+        ofitarball="libfabric-$ofiversion.tar.bz2"
+        ofidestdir="$SOURCE"/3rd-party/"${ofitarball%%.tar.bz2}"
+        if test ! -d "$ofidestdir"; then
+            if test ! -f "$ofitarball"; then
+                echo downloading "$ofiurlbase"/"$ofitarball"...
+                curl -fsSLO "$ofiurlbase"/"$ofitarball"
+            else
+                echo reusing "$ofitarball"...
+            fi
+            echo extracting "$ofitarball" to "$ofidestdir"...
+            tar xf "$ofitarball"
+            mkdir -p "$(dirname "$ofidestdir")"
+            mv "$(basename "$ofidestdir")" "$ofidestdir"
+        else
+            echo reusing directory "$ofidestdir"...
+        fi
+        ucxgithub="https://github.com/openucx/ucx"
+        ucxurlbase="$ucxgithub/releases/download/v$ucxversion"
+        ucxtarball="ucx-$ucxversion.tar.gz"
+        ucxdestdir="$SOURCE"/3rd-party/"${ucxtarball%%.tar.gz}"
+        if test ! -d "$ucxdestdir"; then
+            if test ! -f "$ucxtarball"; then
+                echo downloading "$ucxurlbase"/"$ucxtarball"...
+                curl -fsSLO "$ucxurlbase"/"$ucxtarball"
+            else
+                echo reusing "$ucxtarball"...
+            fi
+            echo extracting "$ucxtarball" to "$ucxdestdir"...
+            tar xf "$ucxtarball"
+            mkdir -p "$(dirname "$ucxdestdir")"
+            mv "$(basename "$ucxdestdir")" "$ucxdestdir"
+        else
+            echo reusing directory "$ucxdestdir"...
+        fi
+    fi
+fi
+
+if test "$mpiname" = "mpich"; then
+    mpidate=$(sed -nE "s/MPICH_RELEASE_DATE=\"(.*)\"/\1/p" "$SOURCE/configure")
+fi
+if test "$mpiname" = "openmpi"; then
+    mpidate=$(sed -nE "s/date=\"(.*)\"/\1/p" "$SOURCE/VERSION")
+fi
+if test -n "${mpidate:-}"; then
+    echo writing source-date-epoch ...
+    date -d "$mpidate" +%s > "$SOURCE/source-date-epoch"
+fi
+
+echo writing package metadata ...
+echo "Name: $mpiname" > "$PACKAGE/METADATA"
+echo "Version: $version" >> "$PACKAGE/METADATA"
+
+echo copying MPI license file...
 cp "$SOURCE"/"$license" "$PACKAGE/LICENSE"
+if test -n "${ofidestdir:-}"; then
+    echo copying OFI license file...
+    cp "$ofidestdir/COPYING" "$PACKAGE/LICENSE.ofi"
+fi
+if test -n "${ucxdestdir:-}"; then
+    echo copying UCX license file...
+    cp "$ucxdestdir/LICENSE" "$PACKAGE/LICENSE.ucx"
+fi

build-wheel.sh

@@ -13,14 +13,14 @@ export CIBW_ARCHS=$ARCHLIST
 export CIBW_BEFORE_ALL='bash {project}/cibw-build-mpi.sh'
 export CIBW_TEST_COMMAND='bash {project}/cibw-check-mpi.sh'
 export CIBW_ENVIRONMENT_PASS='MPINAME RELEASE SOURCE WORKDIR DESTDIR'
-export CIBW_REPAIR_WHEEL_COMMAND_MACOS='delocate-wheel --ignore-missing-dependencies --exclude libmpi --exclude libpmpi --require-archs {delocate_archs -w {dest_dir} -v {wheel}'
+export CIBW_REPAIR_WHEEL_COMMAND_MACOS='delocate-wheel --ignore-missing-dependencies --exclude libmpi --exclude libpmpi --exclude libfabric --require-archs {delocate_archs -w {dest_dir} -v {wheel}'
 
 if test "$(uname)" = Linux; then
     containerengine=$(basename "$(command -v podman || command -v docker)")
     export CIBW_CONTAINER_ENGINE=$containerengine
     export SOURCE="/project/$SOURCE"
-    export WORKDIR="/host/$PWD/$WORKDIR"
-    export DESTDIR="/host/$PWD/$DESTDIR"
+    export WORKDIR="/host$PWD/$WORKDIR"
+    export DESTDIR="/host$PWD/$DESTDIR"
     platform=linux
 fi
 if test "$(uname)" = Darwin; then

check-wheel.sh

@@ -16,7 +16,9 @@ unzip -qq "$wheelfile" -d "$workdir"
 cd "$workdir"
 
 whlname=$(basename "$wheelfile")
-pkgname=${whlname%%-*}
+whlinfo=${whlname%%-py*}
+pkgname=${whlinfo%%-*}
+version=${whlinfo##*-}
 mpiname=${pkgname}
 
 data=$(ls -d "$pkgname"-*.data/data)
@@ -29,9 +31,12 @@ if test "$(uname)" = Linux; then
     libsdir=.libs
     print-runpath() { patchelf --print-rpath  "$1"; }
     print-needed()  { patchelf --print-needed "$1"; }
-    if test -f "$data"/lib/libucp.so; then
+    if test -f "$data"/lib/*/libucp.so.*; then
         runlibs=$runlibs'|libuc(m|p|s|t)'$soregex
     fi
+    if test -f "$data"/lib/*/libfabric.so.*; then
+        runlibs=$runlibs'|libfabric'$soregex
+    fi
 fi
 if test "$(uname)" = Darwin; then
     runpath='@executable_path/../lib/|@loader_path/'
@@ -41,9 +46,12 @@ if test "$(uname)" = Darwin; then
     libsdir=.dylibs
     print-runpath() { otool -l "$1" | sed -n '/RPATH/{n;n;p;}'; }
     print-needed()  { otool -L "$1" | sed 1,1d; }
-    if test -f "$data"/lib/libucp.dylib; then
+    if test -f "$data"/lib/libucp.*.dylib; then
         runlibs=$runlibs'|libuc(m|p|s|t)'$soregex
     fi
+    if test -f "$data"/lib/libfabric.*.dylib; then
+        runlibs=$runlibs'|libfabric'$soregex
+    fi
 fi
 
 if test "$mpiname" = "mpich"; then
@@ -65,12 +73,17 @@ if test "$mpiname" = "mpich"; then
         "$data"/bin/hydra_*
     )
     libraries=(
-        "$data"/lib/libmpi.*
+        "$data"/lib/lib*mpi.*
     )
-    if test -d "$data"/lib/ucx; then
+    if ls "$data"/lib/*/libfabric.* > /dev/null 2>&1; then
+        libraries+=(
+            "$data"/lib/*/libfabric.*
+        )
+    fi
+    if ls "$data"/lib/*/libucp.* > /dev/null 2>&1; then
         libraries+=(
-            "$data"/lib/libuc[mpst]*.*
-            "$data"/lib/ucx/libuc*.*
+            "$data"/lib/*/libuc[mpst]*.*
+            "$data"/lib/*/ucx/libuc[mpst]*.*
         )
     fi
 fi
@@ -96,33 +109,53 @@ if test "$mpiname" = "openmpi"; then
         "$data"/lib/libmpi.*
         "$data"/lib/libopen-*.*
     )
+    if test "${version%%.*}" -ge 5; then
+        libraries+=(
+            "$data"/lib/libevent*.*
+            "$data"/lib/libhwloc.*
+            "$data"/lib/libpmix.*
+            "$data"/lib/libprrte.*
+        )
+    fi
+    if ls "$data"/lib/*/libucp.* > /dev/null 2>&1; then
+        libraries+=(
+            "$data"/lib/*/libuc[mpst]*.*
+            "$data"/lib/*/ucx/libuc[mpst]*.*
+        )
+    fi
+    if ls "$data"/lib/*/libfabric.* > /dev/null 2>&1; then
+        libraries+=(
+            "$data"/lib/*/libfabric.*
+        )
+    fi
     runlibs+='|lib(z|util|event.*|hwloc)'$soregex
     runlibs+='|lib(open-(pal|rte)|pmix|prrte)'$soregex
 fi
 
 check-binary() {
     local dso=$1 out1="" out2=""
     echo checking "$dso"...
-    test -f "$dso" || printf "ERROR: file not found\n"
+    test -f "$dso" || (printf "ERROR: file not found\n"; exit 1)
     out1="$(print-runpath "$dso" | grep -vE "$runpath" || true)"
     test -z "$out1" || printf "ERROR: RUNPATH\n%s\n" "$out1"
     out2="$(print-needed  "$dso" | grep -vE "$runlibs" || true)"
     test -z "$out2" || printf "ERROR: NEEDED\n%s\n" "$out2"
-    test -z "$out1" -a -z "$out2"
+    test -z "$out1"
+    test -z "$out2"
 }
 
 for header in "${headers[@]-}"; do
     test -n "$header" || break
     echo checking "$header"...
-    test -f "$header"
+    test -f "$header" || (printf "ERROR: file not found\n"; exit 1)
     out=$(grep -E '^#include\s+"mpicxx\.h"' "$header" || true)
     test -z "$out" || printf "ERROR: include\n%s\n" "$out"
     test -z "$out"
 done
 for script in "${scripts[@]-}"; do
     test -n "$script" || break
     echo checking "$script"...
-    test -f "$script"
+    test -f "$script" || (printf "ERROR: file not found\n"; exit 1)
     out=$(grep -E "/opt/$mpiname" "$script" || true)
     test -z "$out" || printf "ERROR: prefix\n%s\n" "$out"
     test -z "$out"