Merge pull request #246867 from markuskowa/add-mpi-hook

add mpiCheckPhaseHook

doc/hooks/index.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

doc/hooks/mpi-check-hook.section.md

···
+#  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
+  ];
+```
+
+

pkgs/applications/science/chemistry/cp2k/default.nix

···
-{ lib, stdenv, fetchFromGitHub, python3, gfortran, blas, lapack
+{ lib, stdenv, fetchFromGitHub, mpiCheckPhaseHook, python3, gfortran, blas, lapack
 , fftw, libint, libvori, libxc, mpi, gsl, scalapack, openssh, makeWrapper
 , libxsmm, spglib, which, pkg-config, plumed, zlib
 , enableElpa ? false
···
     EOF
   '';
 
+  nativeCheckInputs = [
+    mpiCheckPhaseHook
+    openssh
+  ];
+
   checkPhase = ''
-    export OMP_NUM_THREADS=1
+    runHook preCheck
 
-    export HYDRA_IFACE=lo  # Fix to make mpich run in a sandbox
-    export OMPI_MCA_rmaps_base_oversubscribe=1
     export CP2K_DATA_DIR=data
+    mpirun -np 2 exe/${arch}/libcp2k_unittest.${cp2kVersion}
 
-    mpirun -np 2 exe/${arch}/libcp2k_unittest.${cp2kVersion}
+    runHook postCheck
   '';
 
   installPhase = ''

pkgs/applications/science/chemistry/nwchem/default.nix

···
 , pkgs
 , fetchFromGitHub
 , fetchurl
+, mpiCheckPhaseHook
 , which
 , openssh
 , gcc
···
   doCheck = false;
 
   doInstallCheck = true;
+  nativeCheckInputs = [ mpiCheckPhaseHook ];
   installCheckPhase = ''
-    export OMP_NUM_THREADS=1
-
-    # Fix to make mpich run in a sandbox
-    export HYDRA_IFACE=lo
-    export OMPI_MCA_rmaps_base_oversubscribe=1
+    runHook preInstallCheck
 
     # run a simple water test
     mpirun -np 2 $out/bin/nwchem $out/share/nwchem/QA/tests/h2o/h2o.nw > h2o.out
     grep "Total SCF energy" h2o.out  | grep 76.010538
+
+    runHook postInstallCheck
   '';
 
   passthru = { inherit mpi; };

pkgs/build-support/setup-hooks/mpi-check-hook/default.nix

···
+{ callPackage, makeSetupHook }:
+
+makeSetupHook {
+  name = "mpi-checkPhase-hook";
+} ./mpi-check-hook.sh

pkgs/build-support/setup-hooks/mpi-check-hook/mpi-check-hook.sh

···
+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
+}
+

pkgs/development/libraries/elpa/default.nix

···
-{ lib, stdenv, fetchurl, autoreconfHook, gfortran, perl
-, mpi, blas, lapack, scalapack, openssh
+{ lib, stdenv, fetchurl, autoreconfHook, mpiCheckPhaseHook
+, gfortran, perl, mpi, blas, lapack, scalapack, openssh
 # CPU optimizations
 , avxSupport ? stdenv.hostPlatform.avxSupport
 , avx2Support ? stdenv.hostPlatform.avx2Support
···
     substituteInPlace Makefile.am --replace '#!/bin/bash' '#!${stdenv.shell}'
   '';
 
-  nativeBuildInputs = [ autoreconfHook perl openssh ];
+  nativeBuildInputs = [ autoreconfHook perl ];
 
   buildInputs = [ mpi blas lapack scalapack ]
     ++ lib.optional enableCuda cudatoolkit;
···
 
   doCheck = true;
 
+  nativeCheckInputs = [ mpiCheckPhaseHook openssh ];
   preCheck = ''
     #patchShebangs ./
-
-    # make sure the test starts even if we have less than 4 cores
-    export OMPI_MCA_rmaps_base_oversubscribe=1
-
-    # Fix to make mpich run in a sandbox
-    export HYDRA_IFACE=lo
 
     # Run dual threaded
     export OMP_NUM_THREADS=2

pkgs/development/libraries/science/math/dbcsr/default.nix

···
 , lib
 , fetchFromGitHub
 , cmake
+, mpiCheckPhaseHook
 , pkg-config
 , fypp
 , gfortran
···
     "-DUSE_MPI=ON"
   ];
 
-  checkInputs = [ openssh ];
+  checkInputs = [
+    openssh
+    mpiCheckPhaseHook
+  ];
 
   doCheck = true;
-  preCheck = ''
-    export HYDRA_IFACE=lo  # Fix to make mpich run in a sandbox
-    export OMPI_MCA_rmaps_base_oversubscribe=1
-  '';
 
   meta = with lib; {
     description = "Distributed Block Compressed Sparse Row matrix library";

pkgs/development/libraries/science/math/p4est-sc/default.nix

···
-{ lib, stdenv, fetchFromGitHub
+{ lib, stdenv, fetchFromGitHub, mpiCheckPhaseHook
 , autoreconfHook, pkg-config
 , p4est-sc-debugEnable ? true, p4est-sc-mpiSupport ? true
 , mpi, openssh, zlib
···
   enableParallelBuilding = true;
   makeFlags = [ "V=0" ];
 
-  preCheck = ''
-    export OMPI_MCA_rmaps_base_oversubscribe=1
-    export HYDRA_IFACE=lo
-  '';
+  nativeCheckInputs = lib.optionals mpiSupport [
+    mpiCheckPhaseHook
+    openssh
+  ];
 
   # disallow Darwin checks due to prototype incompatibility of qsort_r
   # to be fixed in a future version of the source code

pkgs/development/libraries/science/math/p4est/default.nix

···
     ++ lib.optional withMetis "--with-metis"
   ;
 
-  inherit (p4est-sc) makeFlags dontDisableStatic enableParallelBuilding preCheck doCheck;
+  inherit (p4est-sc) makeFlags dontDisableStatic enableParallelBuilding doCheck;
 
   meta = {
     branch = "prev3-develop";

pkgs/development/libraries/science/math/scalapack/default.nix

···
-{ lib, stdenv, fetchFromGitHub, fetchpatch, cmake, openssh
-, mpi, blas, lapack
+{ lib, stdenv, fetchFromGitHub, fetchpatch, cmake
+, openssh, mpiCheckPhaseHook, mpi, blas, lapack
 } :
 
 assert blas.isILP64 == lapack.isILP64;
···
   '';
 
   nativeBuildInputs = [ cmake ];
-  nativeCheckInputs = [ openssh ];
+  nativeCheckInputs = [ openssh mpiCheckPhaseHook ];
   buildInputs = [ blas lapack ];
   propagatedBuildInputs = [ mpi ];
   hardeningDisable = lib.optionals (stdenv.isAarch64 && stdenv.isDarwin) [ "stackprotector" ];
···
   # Increase individual test timeout from 1500s to 10000s because hydra's builds
   # sometimes fail due to this
   checkFlagsArray = [ "ARGS=--timeout 10000" ];
-
-  preCheck = ''
-    # make sure the test starts even if we have less than 4 cores
-    export OMPI_MCA_rmaps_base_oversubscribe=1
-
-    # Fix to make mpich run in a sandbox
-    export HYDRA_IFACE=lo
-
-    # Run single threaded
-    export OMP_NUM_THREADS=1
-  '';
 
   meta = with lib; {
     homepage = "http://www.netlib.org/scalapack/";

pkgs/development/python-modules/meep/default.nix

···
 , fetchFromGitHub
 , autoreconfHook
 , pkg-config
+, mpiCheckPhaseHook
 , gfortran
 , mpi
 , blas
···
   errors can be caught.
   */
   doCheck = true;
+  nativeCheckInputs = [ mpiCheckPhaseHook openssh ];
   checkPhase = ''
-    export PATH=$PATH:${openssh}/bin
-    export PYTHONPATH="$out/lib/${python.libPrefix}/site-packages:$PYTHONPATH"
-
-    export OMP_NUM_THREADS=1
+    runHook preCheck
 
-    # Fix to make mpich run in a sandbox
-    export HYDRA_IFACE=lo
-    export OMPI_MCA_rmaps_base_oversubscribe=1
+    export PYTHONPATH="$out/lib/${python.libPrefix}/site-packages:$PYTHONPATH"
 
     # Generate a python test script
     cat > test.py << EOF
···
     EOF
 
     ${mpi}/bin/mpiexec -np 2 python3 test.py
+
+    runHook postCheck
   '';
 
   meta = with lib; {

pkgs/development/python-modules/mpi4py/default.nix

···
-{ lib, fetchPypi, fetchpatch, python, buildPythonPackage, mpi, openssh }:
+{ lib, fetchPypi, fetchpatch, python, buildPythonPackage
+, mpi, mpiCheckPhaseHook, openssh
+}:
 
 buildPythonPackage rec {
   pname = "mpi4py";
···
     # sometimes packages specify where files should be installed outside the usual
     # python lib prefix, we override that back so all infrastructure (setup hooks)
     # work as expected
-
-    # Needed to run the tests reliably. See:
-    # https://bitbucket.org/mpi4py/mpi4py/issues/87/multiple-test-errors-with-openmpi-30
-    export OMPI_MCA_rmaps_base_oversubscribe=yes
   '';
 
   setupPyBuildFlags = ["--mpicc=${mpi}/bin/mpicc"];
···
 
   __darwinAllowLocalNetworking = true;
 
-  nativeCheckInputs = [ openssh ];
+  nativeCheckInputs = [ openssh mpiCheckPhaseHook ];
 
   meta = with lib; {
     description = "Python bindings for the Message Passing Interface standard";

pkgs/top-level/all-packages.nix

···
   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 { };