From 587a19e43c61fda915adbfc5e487a2e7bb89fa03 Mon Sep 17 00:00:00 2001
From: Markus Kowalewski <markus.kowalewski@gmail.com>
Date: Wed, 2 Aug 2023 21:38:51 +0200
Subject: [PATCH] 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.
---
 doc/hooks/index.md                            |  1 +
 doc/hooks/mpi-check-hook.section.md           | 24 +++++++++
 .../setup-hooks/mpi-check-hook/default.nix    |  5 ++
 .../mpi-check-hook/mpi-check-hook.sh          | 54 +++++++++++++++++++
 pkgs/top-level/all-packages.nix               |  1 +
 5 files changed, 85 insertions(+)
 create mode 100644 doc/hooks/mpi-check-hook.section.md
 create mode 100644 pkgs/build-support/setup-hooks/mpi-check-hook/default.nix
 create mode 100644 pkgs/build-support/setup-hooks/mpi-check-hook/mpi-check-hook.sh

diff --git a/doc/hooks/index.md b/doc/hooks/index.md
index 602febaf9d9b..8100e91c8b48 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 000000000000..e3fb5c40dada
--- /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 000000000000..2834cfcc44ff
--- /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 000000000000..fca1f7b7f932
--- /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 c89e688e2db6..17f9d307d59c 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 { };