Merge pull request #3416 from jsiirola/mpi4py_available

Update `mpi4py_available` to work around `conda-forge/openmpi`

pyomo/common/dependencies.py

@@ -998,6 +998,12 @@ def _finalize_matplotlib(module, available):
     import matplotlib.backends
 
 
+def _finalize_mpi4py(module, available):
+    if not available:
+        return
+    import mpi4py.MPI
+
+
 def _finalize_numpy(np, available):
     if not available:
         return
@@ -1081,7 +1087,9 @@ def _pyutilib_importer():
 
     # Commonly-used optional dependencies
     dill, dill_available = attempt_import('dill')
-    mpi4py, mpi4py_available = attempt_import('mpi4py')
+    mpi4py, mpi4py_available = attempt_import(
+        'mpi4py', deferred_submodules=['MPI'], callback=_finalize_mpi4py
+    )
     networkx, networkx_available = attempt_import('networkx')
     numpy, numpy_available = attempt_import('numpy', callback=_finalize_numpy)
     pandas, pandas_available = attempt_import('pandas')

pyomo/common/tests/test_dependencies.py

@@ -27,6 +27,7 @@
     check_min_version,
     dill,
     dill_available,
+    mpi4py_available,
 )
 
 import pyomo.common.tests.dep_mod as dep_mod
@@ -450,6 +451,12 @@ class A_Class(UnavailableClass(module_obj)):
         ):
             A_Class.method()
 
+    @unittest.pytest.mark.mpi
+    def test_mpi4py_available(self):
+        from mpi4py import MPI
+
+        self.assertTrue(bool(mpi4py_available))
+
 
 if __name__ == '__main__':
     unittest.main()

pyomo/contrib/benders/benders_cuts.py

@@ -9,27 +9,22 @@
 #  This software is distributed under the 3-clause BSD License.
 #  ___________________________________________________________________________
 
+import logging
+
+from pyomo.common.collections import ComponentSet
+from pyomo.common.dependencies import (
+    mpi4py,
+    mpi4py_available,
+    numpy as np,
+    numpy_available,
+)
 from pyomo.core.base.block import BlockData, declare_custom_block
-import pyomo.environ as pyo
-from pyomo.solvers.plugins.solvers.persistent_solver import PersistentSolver
 from pyomo.core.expr.visitor import identify_variables
-from pyomo.common.collections import ComponentSet
-
-try:
-    from mpi4py import MPI
-
-    mpi4py_available = True
-except:
-    mpi4py_available = False
-try:
-    import numpy as np
-
-    numpy_available = True
-except:
-    numpy_available = False
-import logging
+from pyomo.solvers.plugins.solvers.persistent_solver import PersistentSolver
 
+import pyomo.environ as pyo
 
+MPI = mpi4py.MPI
 logger = logging.getLogger(__name__)
 
 

pyomo/contrib/pynumero/sparse/tests/test_mpi_block_matrix.py

@@ -12,6 +12,7 @@
 import warnings
 import pyomo.common.unittest as unittest
 
+from pyomo.common.dependencies import mpi4py, mpi4py_available
 from pyomo.contrib.pynumero.dependencies import (
     numpy_available,
     scipy_available,
@@ -24,15 +25,13 @@
 else:
     SKIPTESTS.append("Pynumero needs scipy and numpy>=1.13.0 to run BlockMatrix tests")
 
-try:
-    from mpi4py import MPI
-
-    comm = MPI.COMM_WORLD
+if mpi4py_available:
+    comm = mpi4py.MPI.COMM_WORLD
     if comm.Get_size() < 3:
         SKIPTESTS.append(
             "Pynumero needs at least 3 processes to run BlockMatrix MPI tests"
         )
-except ImportError:
+else:
     SKIPTESTS.append("Pynumero needs mpi4py to run BlockMatrix MPI tests")
 
 if not SKIPTESTS:

pyomo/contrib/pynumero/sparse/tests/test_mpi_block_vector.py

@@ -10,6 +10,7 @@
 #  ___________________________________________________________________________
 import pyomo.common.unittest as unittest
 
+from pyomo.common.dependencies import mpi4py, mpi4py_available
 from pyomo.contrib.pynumero.dependencies import (
     numpy_available,
     scipy_available,
@@ -22,15 +23,13 @@
 else:
     SKIPTESTS.append("Pynumero needs scipy and numpy>=1.13.0 to run BlockMatrix tests")
 
-try:
-    from mpi4py import MPI
-
-    comm = MPI.COMM_WORLD
+if mpi4py_available:
+    comm = mpi4py.MPI.COMM_WORLD
     if comm.Get_size() < 3:
         SKIPTESTS.append(
             "Pynumero needs at least 3 processes to run BlockVector MPI tests"
         )
-except ImportError:
+else:
     SKIPTESTS.append("Pynumero needs mpi4py to run BlockVector MPI tests")
 
 if not SKIPTESTS: