64 safely checking for mpi4py import

rsopt/mpi.py

@@ -1,4 +1,8 @@
+active_env = None
+
 def get_mpi_environment():
+    global active_env
+
     # Test for mpi4py install
     try:
         import mpi4py
@@ -8,6 +12,10 @@ def get_mpi_environment():
         # mpi4py not installed so it can't be used
         return None
 
+    # If we already ran this process, return the active environment
+    if active_env:
+        return active_env
+
     from inspect import currentframe, getframeinfo
     frameinfo = getframeinfo(currentframe())
     print(f"Initializing MPI from {frameinfo.filename}:L{frameinfo.lineno}", flush=True)
@@ -38,4 +46,7 @@ def get_mpi_environment():
     is_manager = MPI.COMM_WORLD.Get_rank() == 0
     mpi_environment = {'mpi_comm': MPI.COMM_WORLD, 'comms': 'mpi', 'nworkers': nworkers, 'is_manager': is_manager}
 
+    # Save global environment
+    active_env = mpi_environment
+
     return mpi_environment

rsopt/util.py

@@ -4,6 +4,8 @@
 import numpy as np
 import pickle
 from libensemble.tools import save_libE_output
+from .mpi import active_env as MPI_ENV
+from .mpi import get_mpi_environment
 
 SLURM_PREFIX = 'nid'
 
@@ -54,15 +56,9 @@ def return_nodelist(nodelist_string):
 
 def return_used_nodes():
     """Returns all used processor names to rank 0 or an empty list if MPI not used. For ranks != 0  returns None."""
-    try:
-        import mpi4py
-        mpi4py.rc.initialize = False
-    except ModuleNotFoundError:
-        # If MPI not being used to start rsopt then no nodes will have srun executed yet
-        return []
-
-    from mpi4py import MPI
-    MPI.Init()
+    if not MPI_ENV:
+        get_mpi_environment()
+
     rank = MPI.COMM_WORLD.Get_rank()
     name = MPI.Get_processor_name()
     all_names = MPI.COMM_WORLD.gather(name, root=0)
@@ -95,15 +91,9 @@ def return_unused_node():
 
 def broadcast(data, root_rank=0):
     """broadcast, or don't bother"""
-    try:
-        import mpi4py
-        mpi4py.rc.initialize = False
-    except ModuleNotFoundError:
-        # If MPI not available for import then assume it isn't needed
-        return data
+    if not MPI_ENV:
+        get_mpi_environment()
 
-    from mpi4py import MPI
-    MPI.Init()
     if MPI.COMM_WORLD.Get_size() == 1:
         return data