Add OpenMPI host injection script

scripts/mpi_support/install_openmpi_host_injection.sh

@@ -0,0 +1,265 @@
+#!/usr/bin/env bash
+
+# This script can be used to install the host MPI libraries under the `.../host_injections` directory.
+# It allows EESSI software to use the MPI stack from the host.
+#
+# The `host_injections` directory is a variant symlink that by default points to
+# `/opt/eessi`, unless otherwise defined in the local CVMFS configuration (see
+# https://cvmfs.readthedocs.io/en/stable/cpt-repo.html#variant-symlinks). For the
+# installation to be successful, this directory needs to be writeable by the user
+# executing this script.
+
+# Initialise our bash functions
+TOPDIR=$(dirname $(realpath $BASH_SOURCE))
+source "$TOPDIR"/../utils.sh
+
+
+# Function to display help message
+show_help() {
+    echo "Usage: $0 [OPTIONS]"
+    echo "Options:"
+    echo "  --help                           Display this help message"
+    echo "  --mpi-path /path/to/mpi          Specify the path to the MPI host installation (Required)"
+    echo "  -t, --temp-dir /path/to/tmpdir   Specify a location to use for temporary"
+    echo "                                   storage during the mpi injection"
+    echo "  --noclean                        Do not remove the temporary directory and the host injected libraries after finishing injection"
+    echo "  --force                          Force MPI injection even if it is already done"
+}
+
+
+parse_cmdline() {
+    while [[ $# -gt 0 ]]; do
+        case "$1" in
+            --help)
+                show_help
+                exit 0
+                ;;
+            --mpi-path)
+                if [ -n "$2" ]; then
+                    readonly MPI_PATH="$2"
+                    shift 2
+                else
+                    echo_red "Error: Argument required for $1"
+                    show_help
+                    exit 1
+                fi
+                ;;
+            -t|--temp-dir)
+                if [ -n "$2" ]; then
+                    readonly TEMP_DIR="$2"
+                    shift 2
+                else
+                    echo_red "Error: Argument required for $1"
+                    show_help
+                    exit 1
+                fi
+                ;;
+            --noclean)
+                CLEAN=false
+                shift 1
+                ;;
+            --force)
+                FORCE=true
+                shift 1
+                ;;
+            *)
+                echo_red "Error: Unknown option: $1"
+                show_help
+                exit 1
+                ;;
+        esac
+    done
+    if [ -z "${MPI_PATH}" ]; then
+        echo_yellow "MPI path was not specified and it is required"
+        show_help
+        exit 0
+    fi
+
+    readonly CLEAN=${CLEAN:=true}
+    readonly FORCE=${FORCE:=false}
+}
+
+
+# ****Warning: patchelf v0.18.0 (currently shipped with EESSI) does not work.****
+# We get v0.17.2
+download_patchelf() {
+    # Temporary directory to save patchelf
+    local tmpdir=$1
+
+    local patchelf_version="0.17.2"
+    local url
+    local curl_opts="-L --silent --show-error --fail"
+
+    url="https://github.com/NixOS/patchelf/releases/download/${patchelf_version}/"
+    url+="patchelf-${patchelf_version}-${EESSI_CPU_FAMILY}.tar.gz"
+
+    local patchelf_path=${tmpdir}/patchelf
+    mkdir ${patchelf_path}
+
+    curl ${url} ${curl_opts} -o ${patchelf_path}/patchelf.tar.gz
+    tar -xf ${patchelf_path}/patchelf.tar.gz -C ${patchelf_path}
+    PATCHELF_BIN=${patchelf_path}/bin/patchelf
+}
+
+
+inject_mpi() {
+    # Temporary directory for injection
+    local tmpdir=$1
+
+    local eessi_ldd="${EESSI_EPREFIX}/usr/bin/ldd"
+    local system_ldd="/usr/bin/ldd"
+
+    local host_injection_mpi_path
+
+    host_injection_mpi_path=${EESSI_SOFTWARE_PATH/versions/host_injections}
+    host_injection_mpi_path+="/rpath_overrides/OpenMPI/system/lib"
+
+    if [ -d ${host_injection_mpi_path} ]; then
+        if [ -n "$(ls -A ${host_injection_mpi_path})" ]; then
+            echo "MPI was already injected"
+            if ${FORCE}; then
+                echo "Forcing new MPI injection"
+            else
+                return 0
+            fi
+        fi
+    fi
+
+    mkdir -p ${host_injection_mpi_path}
+
+    local temp_inject_path="${tmpdir}/mpi_inject"
+    mkdir ${temp_inject_path}
+
+    # Get all library files from openmpi dir
+    find ${MPI_PATH} -maxdepth 1 -type f -name "*.so*" -exec cp {} ${temp_inject_path} \;
+
+    # Copy library links to host injection path
+    find ${MPI_PATH} -maxdepth 1 -type l -name "*.so*" -exec cp -P {} ${host_injection_mpi_path} \;
+
+    # Get MPI libs dependencies from system ldd
+    local libname libpath pmixpath
+    local -A libs_dict
+    local -a dlopen_libs
+
+    readarray -d '' dlopen_libs < <(find ${MPI_PATH} -mindepth 2 -name "*.so*")
+
+    # Get all library names and paths in associative array
+    # If library is libfabric, libpmix, or from the MPI path
+    # modify libpath in assoc array to point to host_injection_mpi_path
+    while read -r libname libpath; do
+        if [[ ${libname} =~ libfabric\.so ]] && [[ ! -f ${temp_inject_path}/${libname} ]]; then
+            local libdir="$(dirname ${libpath})/"     # without trailing slash the find does not work
+            find ${libdir} -maxdepth 1 -type f -name "libfabric.so*" -exec cp {} ${temp_inject_path} \;
+            find ${libdir} -maxdepth 1 -type l -name "libfabric.so*" -exec cp -P {} ${host_injection_mpi_path} \;
+
+            local depname deppath
+            while read -r depname deppath; do
+                libs_dict[${depname}]=${deppath}
+            done < <(${system_ldd} ${libpath} | awk '/=>/ {print $1, $3}' | sort | uniq)
+
+            libpath=${host_injection_mpi_path}/$(basename ${libpath})
+        fi
+
+        if [[ ${libname} =~ libpmix\.so ]] && [[ ! -f ${temp_inject_path}/${libname} ]]; then
+            local libdir="$(dirname ${libpath})/"     # without trailing slash the find does not work
+            [ -n "${PMIX_PATH}" ] && pmixpath="${PMIX_PATH}/pmix" || pmixpath="$(dirname ${libpath})/pmix"
+            find ${libdir} -maxdepth 1 -type f -name "libpmix.so*" -exec cp {} ${temp_inject_path} \;
+            find ${libdir} -maxdepth 1 -type l -name "libpmix.so*" -exec cp -P {} ${host_injection_mpi_path} \;
+
+            libpath=${host_injection_mpi_path}/$(basename ${libpath})
+        fi
+
+        if [[ ${libpath} =~ ${MPI_PATH} ]]; then
+            libpath=${host_injection_mpi_path}/$(basename ${libpath})
+        fi
+
+        libs_dict[${libname}]=${libpath}
+
+    done < <(cat <(find ${temp_inject_path} -maxdepth 1 -type f -name "*.so*" -exec ${system_ldd} {} \;) \
+                 <(for dlopen in ${dlopen_libs[@]}; do ${system_ldd} ${dlopen}; done) \
+            | awk '/=>/ {print $1, $3}' | sort | uniq)
+
+    # Do library injection to openmpi libs, libfabric and libpmix
+    local lib
+    while read -r lib; do
+        local dep
+
+        # Force system libefa, librdmacm, libibverbs and libpsm2 (present in the EESSI compat layer)
+        # Must be done before the injection of unresolved dependencies
+        if [[ ${lib} =~ libfabric\.so ]]; then
+            while read -r dep; do
+                ${PATCHELF_BIN} --replace-needed ${dep} ${libs_dict[${dep}]} ${lib}
+            done < <(${system_ldd} ${lib} | awk '/libefa/ || /libibverbs/ || /libpsm2/ || /librdmacm/ {print $1}' | sort | uniq)
+        fi
+
+        # Do injection of unresolved libraries
+        ${PATCHELF_BIN} --set-rpath "${host_injection_mpi_path}" ${lib}
+        while read -r dep; do
+            if ${PATCHELF_BIN} --print-needed ${lib} | grep -q "${dep}"; then
+                ${PATCHELF_BIN} --replace-needed ${dep} ${libs_dict[${dep}]} ${lib}
+            else
+                ${PATCHELF_BIN} --add-needed ${libs_dict[${dep}]} ${lib}
+            fi
+        done < <(${eessi_ldd} ${lib} | awk '/not found/ {print $1}' | sort | uniq)
+
+        # Inject into libmpi.so non resolved dependencies from dlopen libraries that are not already present in libmpi.so
+        if [[ ${lib} =~ libmpi\.so ]]; then
+            while read -r dep; do
+                if ! ${PATCHELF_BIN} --print-needed ${lib} | grep -q "${dep}"; then
+                    ${PATCHELF_BIN} --add-needed ${libs_dict[${dep}]} ${lib}
+                fi
+            done < <(for dlopen in ${dlopen_libs[@]}; do ${eessi_ldd} ${dlopen}; done \
+                     | awk '/not found/ && !/libmpi\.so.*/ {print $1}' | sort | uniq)
+        fi
+
+    done < <(find ${temp_inject_path} -type f)
+
+    # Sanity check MPI injection
+    local sanity=1
+    if ${eessi_ldd} ${temp_inject_path}/* &> /dev/null; then
+        cp ${temp_inject_path}/* -t ${host_injection_mpi_path}
+        if ${eessi_ldd} ${temp_inject_path}/* | grep -q "not found"; then
+            ${CLEAN} && rm -f ${host_injection_mpi_path}/*.so*
+        else
+            sanity=0
+        fi
+    fi
+
+    return ${sanity}
+}
+
+
+main() {
+    parse_cmdline "$@"
+    check_eessi_initialised
+
+    # Create directory linked by host_injections
+    local inject_dir=$(readlink -f /cvmfs/software.eessi.io/host_injections)
+    [[ ! -d ${inject_dir} ]] && mkdir -p ${inject_dir}
+
+    # we need a directory we can use for temporary storage
+    if [[ -z "${TEMP_DIR}" ]]; then
+        tmpdir=$(mktemp -d)
+    else
+        tmpdir="${TEMP_DIR}"/temp
+        if ! mkdir -p "$tmpdir" ; then
+            fatal_error "Could not create directory ${tmpdir}"
+        fi
+    fi
+
+    echo "Temporary directory for injection: ${tmpdir}"
+
+    download_patchelf ${tmpdir}
+
+    if inject_mpi ${tmpdir}; then
+        echo_green "MPI injection was successful"
+    else
+        fatal_error "MPI host injection failed"
+    fi
+
+    if ${CLEAN}; then
+        rm -rf "${tmpdir}"
+    fi
+}
+
+main "$@"