Merge pull request #108983 from markuskowa/dev-mpi

Use mpi attribute consistently to provide a default MPI implementation

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

@@ -3,7 +3,7 @@
 , pkgs
 , numpy
 , scipy
-, openmpi
+, mpi
 , enum34
 , protobuf
 , pip
@@ -17,8 +17,8 @@ in
 buildPythonPackage {
   inherit (cntk) name version src;
 
-  nativeBuildInputs = [ swig openmpi ];
-  buildInputs = [ cntk openmpi ];
+  nativeBuildInputs = [ swig mpi ];
+  buildInputs = [ cntk mpi ];
   propagatedBuildInputs = [ numpy scipy enum34 protobuf pip ];
 
   CNTK_LIB_PATH = "${cntk}/lib";
@@ -28,7 +28,7 @@ buildPythonPackage {
 
   postPatch = ''
     cd bindings/python
-    sed -i 's,"libmpi.so.12","${openmpi}/lib/libmpi.so",g' cntk/train/distributed.py
+    sed -i 's,"libmpi.so.12","${mpi}/lib/libmpi.so",g' cntk/train/distributed.py
 
     # Remove distro and libs checks; they aren't compatible with NixOS and besides we guarantee
     # compatibility by providing a package.

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

@@ -1,7 +1,7 @@
 { stdenv, lib, fetchFromGitHub, fetchpatch, buildPythonPackage, python,
   cudaSupport ? false, cudatoolkit ? null, cudnn ? null, nccl ? null, magma ? null,
   mklDnnSupport ? true, useSystemNccl ? true,
-  openMPISupport ? false, openmpi ? null,
+  MPISupport ? false, mpi,
   buildDocs ? false,
   cudaArchList ? null,
 
@@ -29,8 +29,6 @@
 
   isPy3k, pythonOlder }:
 
-assert !openMPISupport || openmpi != null;
-
 # assert that everything needed for cuda is present and that the correct cuda versions are used
 assert !cudaSupport || cudatoolkit != null;
 assert cudnn == null || cudatoolkit != null;
@@ -38,7 +36,7 @@ assert !cudaSupport || (let majorIs = lib.versions.major cudatoolkit.version;
                         in majorIs == "9" || majorIs == "10" || majorIs == "11");
 
 # confirm that cudatoolkits are sync'd across dependencies
-assert !(openMPISupport && cudaSupport) || openmpi.cudatoolkit == cudatoolkit;
+assert !(MPISupport && cudaSupport) || mpi.cudatoolkit == cudatoolkit;
 assert !cudaSupport || magma.cudatoolkit == cudatoolkit;
 
 let
@@ -224,7 +222,7 @@ in buildPythonPackage rec {
     typing-extensions
     # the following are required for tensorboard support
     pillow six future tensorflow-tensorboard protobuf
-  ] ++ lib.optionals openMPISupport [ openmpi ]
+  ] ++ lib.optionals MPISupport [ mpi ]
     ++ lib.optionals (pythonOlder "3.7") [ dataclasses ];
 
   checkInputs = [ hypothesis ninja psutil ];

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

@@ -11,7 +11,7 @@
 # Common deps
 , git, pybind11, which, binutils, glibcLocales, cython, perl
 # Common libraries
-, jemalloc, openmpi, gast, grpc, sqlite, boringssl, jsoncpp
+, jemalloc, mpi, gast, grpc, sqlite, boringssl, jsoncpp
 , curl, snappy, flatbuffers-core, lmdb-core, icu, double-conversion, libpng, libjpeg_turbo, giflib
 # Upsteam by default includes cuda support since tensorflow 1.15. We could do
 # that in nix as well. It would make some things easier and less confusing, but
@@ -129,7 +129,7 @@ let
 
     buildInputs = [
       jemalloc
-      openmpi
+      mpi
       glibcLocales
       git