misc: Use mpi_switch_log for mpi rank 0 logging

devito/operator/operator.py

@@ -6,11 +6,11 @@
 from cached_property import cached_property
 from sympy import sympify
 
-from devito import switchconfig
+from devito import mpi_switch_log
 from devito.arch import compiler_registry, platform_registry
 from devito.data import default_allocator
 from devito.exceptions import InvalidOperator
-from devito.logger import debug, info, perf, warning, is_log_enabled_for, set_log_level
+from devito.logger import debug, info, perf, warning, is_log_enabled_for
 from devito.ir.equations import LoweredEq, lower_exprs
 from devito.ir.clusters import ClusterGroup, clusterize
 from devito.ir.iet import (Callable, CInterface, EntryFunction, FindSymbols, MetaCall,
@@ -874,15 +874,12 @@ def apply(self, **kwargs):
 
         # In case MPI is used restrict result logging to one rank only
         if configuration['mpi']:
-            # Only temporarily change configuration
-            with switchconfig(mpi=True):
-                set_log_level('DEBUG', comm=args.comm)
+            with mpi_switch_log(log_level='DEBUG', comm=args.comm):
                 return self._emit_apply_profiling(args)
-
-        return self._emit_apply_profiling(args)
+        else:
+            return self._emit_apply_profiling(args)
 
     # Performance profiling
-
     def _emit_build_profiling(self):
         if not is_log_enabled_for('PERF'):
             return

devito/parameters.py

@@ -4,11 +4,11 @@
 from os import environ
 from functools import wraps
 
-from devito.logger import info, warning
+from devito.logger import info, warning, logger, stream_handler, set_log_level
 from devito.tools import Signer, filter_ordered
 
 __all__ = ['configuration', 'init_configuration', 'print_defaults', 'print_state',
-           'switchconfig']
+           'switchconfig', 'mpi_switch_log']
 
 # Be EXTREMELY careful when writing to a Parameters dictionary
 # Read here for reference: http://wiki.c2.com/?GlobalVariablesAreBad
@@ -258,6 +258,27 @@ def wrapper(*args, **kwargs):
         return wrapper
 
 
+class mpi_switch_log(switchconfig):
+    """
+    A context manager subclassing `switchconfig` to temporarily change
+    MPI logging.
+    """
+
+    def __init__(self, **params):
+
+        self.params = {k.replace('_', '-'): v for k, v in params.items()}
+        self.previous = {}
+
+        # Limit logging to rank 0
+        comm = self.params.pop('comm')
+        set_log_level(self.params['log-level'], comm=comm)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # Reinstate logging upon exit
+        set_log_level(self.previous['log-level'])
+        logger.addHandler(stream_handler)
+
+
 def print_defaults():
     """Print the environment variables accepted by Devito, their default value,
     as well as all of the accepted values."""