Fix propagation of logging levels in MPIEvaluator and related modules

EwoutH · EwoutH · commit 89817aef3b8e · 2023-10-27T12:26:09.000+02:00
- Introduced logging level propagation to the `MPIEvaluator` to ensure consistent logging across MPI processes.
- Adjusted the `mpi_initializer` function to accept a logger level and configure the logging for each worker accordingly.
- Changed the logging statements within `MPIEvaluator` and `run_experiment_mpi` from DEBUG to INFO for better visibility.
- Removed a redundant TODO regarding logging in `run_experiment_mpi`.
- Modified the `log_to_stderr` method in `ema_logging.py` to allow for setting logging levels on all root module loggers.
- Simplified the logging setup in the `ema_model.py` script, and ensured proper propagation of logging levels.
diff --git a/ema_workbench/em_framework/evaluators.py b/ema_workbench/em_framework/evaluators.py
@@ -13,6 +13,7 @@
 import sys
 import threading
 import warnings
+import logging
 
 from ema_workbench.em_framework.samplers import AbstractSampler
 from .callbacks import DefaultCallback
@@ -419,10 +420,13 @@ def evaluate_experiments(self, scenarios, policies, callback, combine="factorial
 experiment_runner = None
 
 
-def mpi_initializer(models):
+def mpi_initializer(models, logger_level):
     global experiment_runner
     experiment_runner = ExperimentRunner(models)
 
+    # Configure logger based on the passed level and adjusted format
+    logging.basicConfig(level=logger_level, format="[%(processName)s/%(levelname)s] %(message)s")
+
 
 class MPIEvaluator(BaseEvaluator):
     """Evaluator for experiments using MPI Pool Executor from mpi4py"""
@@ -440,7 +444,7 @@ def initialize(self):
         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,))
+        self._pool = MPIPoolExecutor(initializer=mpi_initializer, initargs=(models, _logger.level))
         _logger.info(f"MPI pool started with {self._pool._max_workers} workers")
         if self._pool._max_workers <= 10:
             _logger.warning(
@@ -460,15 +464,15 @@ def evaluate_experiments(self, scenarios, policies, callback, combine="factorial
         packed = [(experiment, experiment.model_name) for experiment in experiments]
 
         # Use the pool to execute in parallel
-        _logger.debug(
+        _logger.info(
             f"MPIEvaluator: Starting {len(packed)} experiments using MPI pool with {self._pool._max_workers} workers"
         )
         results = self._pool.map(run_experiment_mpi, packed)
 
-        _logger.debug(f"MPIEvaluator: Completed all {len(packed)} experiments")
+        _logger.info(f"MPIEvaluator: Completed all {len(packed)} experiments")
         for experiment, outcomes in results:
             callback(experiment, outcomes)
-        _logger.debug(f"MPIEvaluator: Callback completed for all {len(packed)} experiments")
+        _logger.info(f"MPIEvaluator: Callback completed for all {len(packed)} experiments")
 
 
 def run_experiment_mpi(packed_data):
@@ -477,13 +481,12 @@ def run_experiment_mpi(packed_data):
     rank = COMM_WORLD.Get_rank()
 
     experiment, model_name = packed_data
-    # TODO logger: They don't seem to use the main module logger/syntax, fix this
-    _logger.debug(f"MPI Rank {rank}: starting {experiment}")
+    _logger.info(f"MPI Rank {rank}: starting {repr(experiment)}")
 
     # Use the global ExperimentRunner created by the initializer
     outcomes = experiment_runner.run_experiment(experiment)
 
-    _logger.debug(f"MPI Rank {rank}: completed {experiment}")
+    _logger.info(f"MPI Rank {rank}: completed {experiment}")
 
     return experiment, outcomes
 
diff --git a/ema_workbench/util/ema_logging.py b/ema_workbench/util/ema_logging.py
@@ -178,7 +178,7 @@ def get_rootlogger():
     return _rootlogger
 
 
-def log_to_stderr(level=None):
+def log_to_stderr(level=None, set_root_logger_levels=False):
     """
     Turn on logging and add a handler which prints to stderr
 
@@ -206,4 +206,8 @@ def log_to_stderr(level=None):
     logger.addHandler(handler)
     logger.propagate = False
 
+    if set_root_logger_levels:
+        for _, mod_logger in _module_loggers.items():
+            mod_logger.setLevel(level)
+
     return logger
diff --git a/scripts/ema_model.py b/scripts/ema_model.py
@@ -19,8 +19,7 @@ def some_model(x1=None, x2=None, x3=None):
 
 
 if __name__ == "__main__":
-    ema_logging.log_to_stderr(level=10)
-    # TODO: Test logging to file
+    ema_logging.log_to_stderr(level=20, set_root_logger_levels=True)
 
     model = Model("simpleModel", function=some_model)  # instantiate the model
 
