mpiCheckPhaseHook: add new setup hook for MPI aware check phases

Add this hook to checkPhase to allow for running MPI application in
the sandbox. It detects the MPI implementations and sets the respective
environment variables.

5 files changed, 85 insertions, 0 deletions

diff --git a/doc/hooks/index.md b/doc/hooks/index.md
index 602febaf9d9ba..8100e91c8b481 100644
--- a/doc/hooks/index.md
+++ b/doc/hooks/index.md
@@ -17,6 +17,7 @@ installShellFiles.section.md
 libiconv.section.md
 libxml2.section.md
 meson.section.md
+mpi-check-hook.section.md
 ninja.section.md
 patch-rc-path-hooks.section.md
 perl.section.md
diff --git a/doc/hooks/mpi-check-hook.section.md b/doc/hooks/mpi-check-hook.section.md
new file mode 100644
index 0000000000000..e3fb5c40dadaa
--- /dev/null
+++ b/doc/hooks/mpi-check-hook.section.md
@@ -0,0 +1,24 @@
+#  mpiCheckPhaseHook {#setup-hook-mpi-check}
+
+
+This hook can be used to setup a check phase that
+requires running a MPI application. It detects the
+used present MPI implementaion type and exports
+the neceesary environment variables to use
+`mpirun` and `mpiexec` in a Nix sandbox.
+
+
+Example:
+
+```nix
+  { mpiCheckPhaseHook, mpi, ... }:
+
+  ...
+
+  nativeCheckInputs = [
+    openssh
+    mpiCheckPhaseHook
+  ];
+```
+
+
diff --git a/pkgs/build-support/setup-hooks/mpi-check-hook/default.nix b/pkgs/build-support/setup-hooks/mpi-check-hook/default.nix
new file mode 100644
index 0000000000000..2834cfcc44ff1
--- /dev/null
+++ b/pkgs/build-support/setup-hooks/mpi-check-hook/default.nix
@@ -0,0 +1,5 @@
+{ callPackage, makeSetupHook }:
+
+makeSetupHook {
+  name = "mpi-checkPhase-hook";
+} ./mpi-check-hook.sh
diff --git a/pkgs/build-support/setup-hooks/mpi-check-hook/mpi-check-hook.sh b/pkgs/build-support/setup-hooks/mpi-check-hook/mpi-check-hook.sh
new file mode 100644
index 0000000000000..fca1f7b7f9327
--- /dev/null
+++ b/pkgs/build-support/setup-hooks/mpi-check-hook/mpi-check-hook.sh
@@ -0,0 +1,54 @@
+preCheckHooks+=('setupMpiCheck')
+preInstallCheckHooks+=('setupMpiCheck')
+
+
+setupMpiCheck() {
+  # Find out which MPI implementation we are using
+  # and set safe defaults that are guaranteed to run
+  # on any build machine
+
+  mpiType="NONE"
+
+  # OpenMPI signature
+  if command ompi_info &> /dev/null; then
+    mpiType="openmpi"
+  fi
+
+  # MPICH based implementations
+  if command mpichversion &> /dev/null; then
+    if [ "$mpiType" != "NONE" ]; then
+      echo "WARNING: found OpenMPI and MPICH/MVAPICH executables"
+    fi
+
+    version=$(mpichversion)
+    if [[ "$version" == *"MPICH"* ]]; then
+      mpiType="MPICH"
+    fi
+    if [[ "$version" == *"MVAPICH"* ]]; then
+      mpiType="MVAPICH"
+    fi
+  fi
+
+  echo "Found MPI implementation: $mpiType"
+
+  case $mpiType in
+    openmpi)
+      # make sure the test starts even if we have less than the requested amount of cores
+      export OMPI_MCA_rmaps_base_oversubscribe=1
+      # Disable CPU pinning
+      export OMPI_MCA_hwloc_base_binding_policy=none
+      ;;
+    MPICH)
+      # Fix to make mpich run in a sandbox
+      export HYDRA_IFACE=lo
+      ;;
+    MVAPICH)
+      # Disable CPU pinning
+      export MV2_ENABLE_AFFINITY=0
+      ;;
+  esac
+
+  # Limit number of OpenMP threads. Default is "all cores".
+  export OMP_NUM_THREADS=1
+}
+
diff --git a/pkgs/top-level/all-packages.nix b/pkgs/top-level/all-packages.nix
index c89e688e2db64..17f9d307d59ca 100644
--- a/pkgs/top-level/all-packages.nix
+++ b/pkgs/top-level/all-packages.nix
@@ -12258,6 +12258,7 @@ with pkgs;
   outils = callPackage ../tools/misc/outils { };
 
   mpi = openmpi; # this attribute should used to build MPI applications
+  mpiCheckPhaseHook = callPackage ../build-support/setup-hooks/mpi-check-hook { };
 
   ucc = callPackage ../development/libraries/ucc { };