Implement MPIEvaluator for multi-node HPC systems support

EwoutH · EwoutH · commit 0bc9e15b1851 · 2023-10-29T12:40:10.000+01:00
Adds a new MPIEvaluator to the EMAworkbench, enabling experiments to be executed on multi-node High-Performance Computing (HPC) systems leveraging the mpi4py library. This evaluator optimizes performance for distributed computing environments by parallelizing experiments across multiple nodes and processors.

Changes include:
- Definition of the MPIEvaluator class.
- Initialization function to set up the global ExperimentRunner for worker processes.
- Proper handling to pack and unpack experiments for efficient data transfer between nodes.

Note: This addition requires the mpi4py package only when the MPIEvaluator is explicitly used, preventing unnecessary dependencies for users not requiring this feature.
diff --git a/ema_workbench/__init__.py b/ema_workbench/__init__.py
@@ -15,6 +15,7 @@
     Constant,
     Scenario,
     Policy,
+    MPIEvaluator,
     MultiprocessingEvaluator,
     IpyparallelEvaluator,
     SequentialEvaluator,
diff --git a/ema_workbench/em_framework/__init__.py b/ema_workbench/em_framework/__init__.py
@@ -30,6 +30,7 @@
     "perform_experiments",
     "optimize",
     "IpyparallelEvaluator",
+    "MPIEvaluator",
     "MultiprocessingEvaluator",
     "SequentialEvaluator",
     "ReplicatorModel",
@@ -76,6 +77,7 @@
 from .evaluators import (
     perform_experiments,
     optimize,
+    MPIEvaluator,
     MultiprocessingEvaluator,
     SequentialEvaluator,
     Samplers,
diff --git a/ema_workbench/em_framework/evaluators.py b/ema_workbench/em_framework/evaluators.py
@@ -415,6 +415,62 @@ def evaluate_experiments(self, scenarios, policies, callback, combine="factorial
         add_tasks(self.n_processes, self._pool, ex_gen, callback)
 
 
+# Create a global ExperimentRunner that will be used by all the worker processes
+experiment_runner = None
+
+
+def mpi_initializer(models):
+    global experiment_runner
+    experiment_runner = ExperimentRunner(models)
+
+
+class MPIEvaluator(BaseEvaluator):
+    """Evaluator for experiments using MPI Pool Executor from mpi4py"""
+
+    def __init__(self, msis, **kwargs):
+        super().__init__(msis, **kwargs)
+        self._pool = None
+
+    def initialize(self):
+        # Only import mpi4py if the MPIEvaluator is used, to avoid unnecessary dependencies.
+        from mpi4py.futures import MPIPoolExecutor
+
+        # Instead of instantiating the ExperimentRunner for each experiment, instantiate it once here
+        models = NamedObjectMap(AbstractModel)
+        models.extend(self._msis)
+
+        # Use the initializer function to set up the ExperimentRunner for all the worker processes
+        self._pool = MPIPoolExecutor(initializer=mpi_initializer, initargs=(models,))
+        _logger.info(f"MPI pool started with {self._pool._max_workers} workers")
+        return self
+
+    def finalize(self):
+        self._pool.shutdown()
+        _logger.info("MPI pool has been shut down")
+
+    def evaluate_experiments(self, scenarios, policies, callback, combine="factorial"):
+        ex_gen = experiment_generator(scenarios, self._msis, policies, combine=combine)
+        experiments = list(ex_gen)  # Convert generator to list
+
+        # Instead of sending all models for each experiment, send only the model_name
+        packed = [(experiment, experiment.model_name) for experiment in experiments]
+
+        # Use the pool to execute in parallel
+        results = self._pool.map(run_experiment_mpi, packed)
+
+        for experiment, outcomes in results:
+            callback(experiment, outcomes)
+
+
+def run_experiment_mpi(packed_data):
+    experiment, model_name = packed_data
+
+    # Use the global ExperimentRunner created by the initializer
+    outcomes = experiment_runner.run_experiment(experiment)
+
+    return experiment, outcomes
+
+
 class IpyparallelEvaluator(BaseEvaluator):
     """evaluator for using an ipypparallel pool"""
 
