Add a couple of utilities to simulate an MPI environment
@@ -57,8 +57,12 @@ logger = logging.getLogger(__name__)
# GLOBAL VARIABLES
# ==============================================================================
+# Force serial execution even in MPI environment.
disable_mpi = False
+# A dummy MPI communicator used to simulate an MPI environment in tests.
+_simulated_mpicomm = None
+
# ==============================================================================
# MAIN FUNCTIONS
@@ -84,6 +88,10 @@ def get_mpicomm():
if disable_mpi:
return None
+ # If MPI is simulated, return the Dummy implementation.
+ if _simulated_mpicomm is not None:
+ return _simulated_mpicomm
+
# If we have already initialized MPI, return the cached MPI communicator.
if get_mpicomm._is_initialized:
return get_mpicomm._mpicomm
@@ -422,8 +430,11 @@ def distribute(task, distributed_args, *other_args, send_results_to='all',
If not None, the ``distributed_args`` are distributed among groups of
nodes that are isolated from each other. This is particularly useful
if ``task`` also calls :func:`distribute`, since normally that would result
- in unexpected behavior. If an integer, the nodes are split into equal
- groups of ``group_size`` nodes. If a list of integers, the nodes are
+ in unexpected behavior.
+
+ If an integer, the nodes are split into equal groups of ``group_size``
+ nodes. If ``n_nodes % group_size != 0``, the first jobs are allocated
+ more nodes than the latest. If a list of integers, the nodes are
split in possibly unequal groups (see example below).
Other Parameters
@@ -588,6 +599,35 @@ def delayed_termination(func):
return _delayed_termination
+# ==============================================================================
+# MPI TEST CLASSES
+# ==============================================================================
+
+class _DummyMPIComm():
+ """A Dummy MPI Communicator."""
+
+ def __init__(self, rank=0, size=4):
+ self.rank = rank
+ self.size = size
+
+
+@contextmanager
+def _simulated_mpi_environment(**kwargs):
+ """Context manager to temporarily set a simulated MPI environment.
+
+ Parameters
+ ----------
+ **kwargs : dict
+ The parameters to pass to _DummyMPIComm constructor.
+
+ """
+ global _simulated_mpicomm
+ old_simulated_mpicomm = _simulated_mpicomm
+ _simulated_mpicomm = _DummyMPIComm(**kwargs)
+ yield
+ _simulated_mpicomm = old_simulated_mpicomm
+
+
# ==============================================================================
# MAIN AND TESTS
# ==============================================================================
