misc: Use mpi_switch_log for mpi rank 0 logging

devito/logger.py

@@ -4,7 +4,7 @@
 import sys
 from contextlib import contextmanager
 
-__all__ = ('set_log_level', 'set_log_noperf', 'is_log_enabled_for',
+__all__ = ('set_log_level', 'set_log_noperf', 'is_log_enabled_for', 'switch_log_level',
            'log', 'warning', 'error', 'perf', 'hint',
            'RED', 'GREEN', 'BLUE')
 
@@ -13,13 +13,13 @@
 stream_handler = logging.StreamHandler()
 logger.addHandler(stream_handler)
 
-# Add extra logging levels
-DEBUG = logging.DEBUG  # value=10
+# Add extra logging levels (note: INFO has value=20, WARNING has value=30)
+DEBUG = logging.DEBUG
 PERF = 19
-INFO = logging.INFO  # value=20
-WARNING = logging.WARNING  # value=30
-ERROR = logging.ERROR  # value=40
-CRITICAL = logging.CRITICAL  # value=50
+INFO = logging.INFO
+WARNING = logging.WARNING
+ERROR = logging.ERROR
+CRITICAL = logging.CRITICAL
 
 logging.addLevelName(PERF, "PERF")
 
@@ -77,15 +77,36 @@ def set_log_level(level, comm=None):
     """
     from devito import configuration
 
-    if comm is not None:
+    if comm is not None and configuration['mpi']:
         if comm.rank != 0:
             logger.removeHandler(stream_handler)
             logger.addHandler(logging.NullHandler())
+    else:
+        logger.addHandler(stream_handler)
 
     # Triggers a callback to `_set_log_level`
     configuration['log-level'] = level
 
 
+class switch_log_level(object):
+    """
+    A context manager to temporarily change MPI logging.
+    """
+
+    def __init__(self, comm):
+
+        from devito import configuration
+        self.level = configuration['log-level']
+        self.comm = comm
+
+    def __enter__(self):
+        # Limit logging to rank 0
+        set_log_level(self.level, self.comm)
+
+    def __exit__(self, *args):
+        set_log_level(self.level)
+
+
 def set_log_noperf():
     """Do not print performance-related messages."""
     logger.setLevel(WARNING)

devito/operator/operator.py

@@ -6,11 +6,10 @@
 from cached_property import cached_property
 from sympy import sympify
 
-from devito import switchconfig
 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.exceptions import InvalidOperator, ExecutionError
+from devito.logger import debug, info, perf, warning, is_log_enabled_for, switch_log_level
 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,
@@ -873,13 +872,8 @@ def apply(self, **kwargs):
         self._postprocess_arguments(args, **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)
-                return self._emit_apply_profiling(args)
-
-        return self._emit_apply_profiling(args)
+        with switch_log_level(comm=args.comm):
+            return self._emit_apply_profiling(args)
 
     # Performance profiling
 
@@ -916,7 +910,6 @@ def _emit_timings(timings, indent=''):
 
     def _emit_apply_profiling(self, args):
         """Produce a performance summary of the profiled sections."""
-
         # Rounder to 2 decimal places
         fround = lambda i: ceil(i * 100) / 100