Repository: microsoft/nni
File: nni/runtime/log.py (Python)

Code before:
    _handlers[experiment_id] = handler
    _root_logger.addHandler(handler)

def stop_experiment_logging(experiment_id: str) -> None:
    """
    Unregister an experiment's ``experiment.log`` handler.
    """
    handler = _handlers.pop(experiment_id, None)
    if handler is not None:
        _root_logger.removeHandler(handler)

def enable_global_logging(enable: bool = True) -> None:
    """
    Let NNI to handle all logs in this process.

    Use ``enable_global_logging(False)`` to reverse the setting.
    (The root log level will not be restored.)
    """
    global _global_logging_enabled
    _global_logging_enabled = enable
    _root_logger.setLevel(logging.INFO)

def _nni_filter(record: LogRecord) -> bool:
    return _global_logging_enabled or record.name.startswith('nni')

class _StdoutFormatter(Formatter):
    def __init__(self):
        fmt = '[%(asctime)s] (%(threadName)s:%(name)s) %(message)s'
        datefmt = '%Y-%m-%d %H:%M:%S'
        super().__init__(fmt, datefmt)

    def formatMessage(self, record: LogRecord) -> str:
        if '${' in record.message:
            message = string.Template(record.message).safe_substitute(_colorful)
        else:
            message = record.message

        if record.levelno >= logging.ERROR:
            level = colorama.Fore.RED + 'ERROR: '
        elif record.levelno >= logging.WARNING:
            level = colorama.Fore.YELLOW + 'WARNING: '
        elif record.levelno >= logging.INFO:
            level = colorama.Fore.GREEN
        else:
            level = colorama.Fore.BLUE

        content = level + message + colorama.Style.RESET_ALL
        if record.levelno >= logging.INFO:
            return f'[{record.asctime}] {content}'
        elif record.threadName == 'MainThread':
            return f'[{record.asctime}] {record.name} {content}'

Review comment (bug):
I think adding a handler to `logging.getLogger('nni')` enables log-handling on every logger prefixed with NNI. There's no need to manipulate the root logger. Maybe I'm wrong.

Code after:
_global_logging_enabled: bool = False

_root_logger: Logger = logging.getLogger('nni')
_root_logger.setLevel(logging.INFO)

_handlers: dict[str, Handler] = {}

_colorful: dict[str, str] = colorama.Fore.__dict__
_colorless: dict[str, str] = {key: '' for key in _colorful.keys()}

def start_stdout_logging() -> None:
    """
    Register the stdout handler.

    This function should be invoked on importing nni.

    It is safe to call it multiple times.
    """
    if '_stdout_' in _handlers:
        return

    handler = StreamHandler(sys.stdout)
    handler.setFormatter(_StdoutFormatter())

    _handlers['_stdout_'] = handler
    _root_logger.addHandler(handler)

def silence_stdout() -> None:
    """
    Stop NNI from printing to stdout.

    By default NNI prints log messages of ``INFO`` and higher levels to console.
    Use this function if you want a clean stdout, or if you want to handle logs by yourself.
    """
    handler = _handlers.pop('_stdout_', None)
    if handler is not None:
        _root_logger.removeHandler(handler)

def start_experiment_logging(experiment_id: str, log_file: Path, level: str) -> None:
    """
    Register the log file handler for an experiment's ``experiment.log``.

    This function should be invoked on starting experiment.
    We don't want to create the experiment folder if the user does not launch it.

    If there are multiple experiments running concurrently,
    log messages will be written to all running experiments' log files.

    It is safe to call it multiple times.
    """
    if experiment_id in _handlers: