|
4 | 4 | import sys
|
5 | 5 | from contextlib import contextmanager
|
6 | 6 |
|
7 |
| -__all__ = ('set_log_level', 'set_log_noperf', 'is_log_enabled_for', |
| 7 | +__all__ = ('set_log_level', 'set_log_noperf', 'is_log_enabled_for', 'switch_log_level', |
8 | 8 | 'log', 'warning', 'error', 'perf', 'hint',
|
9 | 9 | 'RED', 'GREEN', 'BLUE')
|
10 | 10 |
|
@@ -71,21 +71,42 @@ def set_log_level(level, comm=None):
|
71 | 71 | comm : MPI communicator, optional
|
72 | 72 | An MPI communicator the logger should be collective over. If provided, only
|
73 | 73 | rank-0 on that communicator will write to the registered handlers, other
|
74 |
| - ranks will use a `logging.NullHandler`. By default, ``comm`` is set |
75 |
| - to ``None``, so all ranks will use the default handlers. This could be |
| 74 | + ranks will use a `logging.NullHandler`. By default, ``comm`` is set |
| 75 | + to ``None``, so all ranks will use the default handlers. This could be |
76 | 76 | used, for example, if one wants to log to one file per rank.
|
77 | 77 | """
|
78 | 78 | from devito import configuration
|
79 | 79 |
|
80 |
| - if comm is not None: |
| 80 | + if comm is not None and configuration['mpi']: |
81 | 81 | if comm.rank != 0:
|
82 | 82 | logger.removeHandler(stream_handler)
|
83 | 83 | logger.addHandler(logging.NullHandler())
|
| 84 | + else: |
| 85 | + logger.addHandler(stream_handler) |
84 | 86 |
|
85 | 87 | # Triggers a callback to `_set_log_level`
|
86 | 88 | configuration['log-level'] = level
|
87 | 89 |
|
88 | 90 |
|
| 91 | +class switch_log_level(object): |
| 92 | + """ |
| 93 | + A context manager to temporarily change MPI logging. |
| 94 | + """ |
| 95 | + |
| 96 | + def __init__(self, comm): |
| 97 | + |
| 98 | + from devito import configuration |
| 99 | + self.level = configuration['log-level'] |
| 100 | + self.comm = comm |
| 101 | + |
| 102 | + def __enter__(self): |
| 103 | + # Limit logging to rank 0 |
| 104 | + set_log_level(self.level, self.comm) |
| 105 | + |
| 106 | + def __exit__(self, *args): |
| 107 | + set_log_level(self.level) |
| 108 | + |
| 109 | + |
89 | 110 | def set_log_noperf():
|
90 | 111 | """Do not print performance-related messages."""
|
91 | 112 | logger.setLevel(WARNING)
|
|
0 commit comments