Logging refactor (step 1) - experiment handlers (#4792)

nni/__init__.py

@@ -6,12 +6,13 @@ try:
 except ModuleNotFoundError:
     __version__ = '999.dev0'
 
-from .runtime.log import init_logger
-init_logger()
+from .runtime.log import _init_logger
+_init_logger()
 
 from .common.serializer import trace, dump, load
 from .experiment import Experiment
 from .runtime.env_vars import dispatcher_env_vars
+from .runtime.log import enable_global_logging, silence_stdout
 from .utils import ClassArgsValidator
 
 if dispatcher_env_vars.SDK_PROCESS != 'dispatcher':

nni/experiment/experiment.py

@@ -10,13 +10,12 @@ from pathlib import Path
 import socket
 from subprocess import Popen
 import time
-from typing import Any
+from typing import Any, cast
 
-import colorama
 import psutil
 from typing_extensions import Literal
 
-import nni.runtime.log
+from nni.runtime.log import start_experiment_logging, stop_experiment_logging
 
 from .config import ExperimentConfig
 from .data import TrialJob, TrialMetricData, TrialResult
@@ -73,8 +72,6 @@ class Experiment:
     """
 
     def __init__(self, config_or_platform: ExperimentConfig | str | list[str] | None):
-        nni.runtime.log.init_logger_for_command_line()
-
         self.config: ExperimentConfig | None = None
         self.id: str = management.generate_experiment_id()
         self.port: int | None = None
@@ -109,11 +106,10 @@ class Experiment:
         if config.use_annotation:
             raise RuntimeError('NNI annotation is not supported by Python experiment API.')
 
-        if config.experiment_working_directory is not None:
-            log_dir = Path(config.experiment_working_directory, self.id, 'log')
-        else:  # this should never happen in latest version, keep it until v2.7 for potential compatibility
-            log_dir = Path.home() / f'nni-experiments/{self.id}/log'
-        nni.runtime.log.start_experiment_log(self.id, log_dir, debug)
+        log_file = Path(config.experiment_working_directory, self.id, 'log', 'experiment.log')
+        log_file.parent.mkdir(parents=True, exist_ok=True)
+        log_level = 'debug' if (debug or config.log_level == 'trace') else config.log_level
+        start_experiment_logging(self.id, log_file, cast(str, log_level))
 
         self._proc = launcher.start_experiment(self._action, self.id, config, port, debug, run_mode, self.url_prefix)
         assert self._proc is not None
@@ -126,7 +122,7 @@ class Experiment:
                 if interface.family == socket.AF_INET:
                     ips.append(interface.address)
         ips = [f'http://{ip}:{port}' for ip in ips if ip]
-        msg = 'Web portal URLs: ' + colorama.Fore.CYAN + ' '.join(ips) + colorama.Style.RESET_ALL
+        msg = 'Web portal URLs: ${CYAN}' + ' '.join(ips)
         _logger.info(msg)
 
     def stop(self) -> None:
@@ -136,7 +132,7 @@ class Experiment:
         _logger.info('Stopping experiment, please wait...')
         atexit.unregister(self.stop)
 
-        nni.runtime.log.stop_experiment_log(self.id)
+        stop_experiment_logging(self.id)
         if self._proc is not None:
             try:
                 rest.delete(self.port, '/experiment', self.url_prefix)

nni/experiment/launcher.py

@@ -15,7 +15,6 @@ import sys
 import time
 from typing import Any, TYPE_CHECKING, cast
 
-import colorama
 from typing_extensions import Literal
 
 from .config import ExperimentConfig
@@ -118,9 +117,7 @@ def start_experiment(
 
     proc = None
     try:
-        _logger.info(
-            'Creating experiment, Experiment ID: %s', colorama.Fore.CYAN + exp_id + colorama.Style.RESET_ALL
-        )
+        _logger.info('Creating experiment, Experiment ID: ${CYAN}%s', exp_id)
         proc = _start_rest_server(nni_manager_args, run_mode)
         start_time = int(time.time() * 1000)
 

nni/retiarii/experiment/pytorch.py

@@ -348,11 +348,10 @@ class RetiariiExperiment(Experiment):
 
         self.id = management.generate_experiment_id()
 
-        if self.config.experiment_working_directory is not None:
-            log_dir = Path(self.config.experiment_working_directory, self.id, 'log')
-        else:
-            log_dir = Path.home() / f'nni-experiments/{self.id}/log'
-        nni.runtime.log.start_experiment_log(self.id, log_dir, debug)
+        log_file = Path(self.config.experiment_working_directory, self.id, 'log', 'experiment.log')
+        log_file.parent.mkdir(parents=True, exist_ok=True)
+        log_level = 'debug' if (debug or self.config.log_level == 'trace') else self.config.log_level
+        nni.runtime.log.start_experiment_logging(self.id, log_file, cast(str, log_level))
 
         ws_url = f'ws://localhost:{port}/tuner'
         self._proc = launcher.start_experiment('create', self.id, self.config, port, debug,  # type: ignore
@@ -465,7 +464,7 @@ class RetiariiExperiment(Experiment):
             self._dispatcher_thread.join(timeout=1)
 
         if self.id is not None:
-            nni.runtime.log.stop_experiment_log(self.id)
+            nni.runtime.log.stop_experiment_logging(self.id)
         if self._proc is not None:
             try:
                 # this if is to deal with the situation that

nni/runtime/log.py

 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+"""
+NNI's logging system.
+
+For end users you only need to care about :func:`silence_stdout` and :func:`enable_global_logging`.
+Following stuffs are for NNI contributors.
+
+The logging system is initialized on importing ``nni``.
+By design it should not have side effects on non-NNI modules' logs,
+unless the user explicitly invokes :func:`enable_global_logging`.
+
+The logging system is divided into experiment (user code) part, trial part, and dispatcher part.
+
+Experiment Part
+===============
+
+There are two kinds of log handlers here, stdout handler and experiment file handler.
+
+Console
+-------
+
+NNI prints log messages of ``INFO`` and above levels to stdout, in colorful format.
+This can be turned off with :func:`silence_stdout`.
+
+Logs are automatically colored according to their log level.
+
+One can manually alter the text color with escape sequence.
+For example ``logger.error('hello ${GREEN}world')`` will print "hello" in red (error's color) and "world" in green.
+
+The escape sequence affects the words from its position to end of line.
+It supports all colors in ``colorama.Fore``.
+
+Please use NNI's escape sequences rather than colorama's because the latter will pollute log files.
+
+Log Files
+---------
+
+NNI save log messages to ``~/nni-experiments/<ID>/log/experiment.log``.
+
+It is the experiment classes' role to invoke :func:`start_experiment_logging` and :func:`stop_experiment_logging`.
+
+Only the messages logged during an experiment's life span will be saved to log file.
+Logs written before `exp.start()` and after `exp.stop()` will not be saved.
+
+If there are multiple experiments running concurrently,
+all logs will be saved to all experiments' log files.
+
+Trial Part
+==========
+
+WIP
+
+Dispatcher Part
+===============
+
+WIP
+"""
+
 from __future__ import annotations
 
+__all__ = ['enable_global_logging', 'silence_stdout']
+
 import logging
 import sys
 from datetime import datetime
 from io import TextIOBase
-from logging import FileHandler, Formatter, Handler, StreamHandler
+from logging import FileHandler, Formatter, Handler, LogRecord, Logger, StreamHandler
 from pathlib import Path
+import string
 from typing import Optional
 
 import colorama
 
 from .env_vars import dispatcher_env_vars, trial_env_vars
 
+_colorama_initialized: bool = False
+_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:
+        return
+
+    handler = FileHandler(log_file, encoding='utf_8')
+    handler.setFormatter(_LogFileFormatter())
+    handler.setLevel(level.upper())
+
+    _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. Useful for debugging.
+
+    By default only NNI's logs are printed to stdout and saved to ``nni-experiments`` log files.
+    The function will extend these settings to all modules' logs.
+
+    Use ``enable_global_logging(False)`` to reverse it.
+    The log level of root logger will not be reversed though.
+    """
+    global _global_logging_enabled, _root_logger
+
+    if enable == _global_logging_enabled:
+        return
+
+    if enable:
+        level = logging.getLogger('nni').getEffectiveLevel()
+        logging.getLogger().setLevel(level)
+
+    new_root = logging.getLogger() if enable else logging.getLogger('nni')
+    for handler in _handlers.values():
+        _root_logger.removeHandler(handler)
+        new_root.addHandler(handler)
+
+    _root_logger = new_root
+    _global_logging_enabled = enable
+
+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:
+        global _colorama_initialized
+
+        if '${' in record.message:  # contains colorful text: "hello ${GREEN}world"
+            if not _colorama_initialized:
+                colorama.init()
+                _colorama_initialized = True
+            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}'
+        else:
+            return f'[{record.asctime}] {record.threadName}:{record.name} {content}'
+
+class _LogFileFormatter(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(_colorless)
+        else:
+            message = record.message
+
+        if record.threadName == 'MainThread':
+            return f'[{record.asctime}] {record.levelname} ({record.name}) {message}'
+        else:
+            return f'[{record.asctime}] {record.levelname} ({record.threadName}:{record.name}) {message}'
+
+
+## legacy part ##
+
 handlers = {}
 
 log_format = '[%(asctime)s] %(levelname)s (%(name)s/%(threadName)s) %(message)s'
@@ -22,7 +232,7 @@ time_format = '%Y-%m-%d %H:%M:%S'
 formatter = Formatter(log_format, time_format)
 
 
-def init_logger() -> None:
+def _init_logger() -> None:
     """
     This function will (and should only) get invoked on the first time of importing nni (no matter which submodule).
     It will try to detect the running environment and setup logger accordingly.
@@ -30,7 +240,9 @@ def init_logger() -> None:
     The detection should work in most cases but for `nnictl` and `nni.experiment`.
     They will be identified as "standalone" mode and must configure the logger by themselves.
     """
-    colorama.init()
+    # I think dispatcher and trial do not need colorful stdout.
+    # Remove this when we finish refactor.
+    #colorama.init()
 
     if dispatcher_env_vars.SDK_PROCESS == 'dispatcher':
         _init_logger_dispatcher()
@@ -45,41 +257,7 @@ def init_logger() -> None:
         _init_logger_trial()
         return
 
-    _init_logger_standalone()
-
-    logging.getLogger('filelock').setLevel(logging.WARNING)
-
-_cli_log_initialized = False
-
-def init_logger_for_command_line() -> None:
-    """
-    Initialize logger for command line usage.
-    This means that NNI is used as "main function" rather than underlying library or background service,
-    so it should print log to stdout.
-
-    It is used by nnictl and `nni.experiment.Experiment`.
-
-    This function will get invoked after `init_logger()`.
-    """
-    global _cli_log_initialized
-    if not _cli_log_initialized:
-        _cli_log_initialized = True
-        colorful_formatter = Formatter(log_format, time_format)
-        colorful_formatter.format = _colorful_format
-        if '_default_' not in handlers:  # this happens when building sphinx gallery
-            _register_handler(StreamHandler(sys.stdout), logging.INFO)
-        handlers['_default_'].setFormatter(colorful_formatter)
-
-def start_experiment_log(experiment_id: str, log_directory: Path, debug: bool) -> None:
-    log_path = _prepare_log_dir(log_directory) / 'dispatcher.log'
-    log_level = logging.DEBUG if debug else logging.INFO
-    _register_handler(FileHandler(log_path), log_level, experiment_id)
-
-def stop_experiment_log(experiment_id: str) -> None:
-    if experiment_id in handlers:
-        handler = handlers.pop(experiment_id, None)
-        if handler is not None:
-            logging.getLogger().removeHandler(handler)
+    start_stdout_logging()
 
 
 def _init_logger_dispatcher() -> None:
@@ -106,10 +284,6 @@ def _init_logger_trial() -> None:
         sys.stdout = _LogFileWrapper(log_file)
 
 
-def _init_logger_standalone() -> None:
-    _register_handler(StreamHandler(sys.stdout), logging.INFO)
-
-
 def _prepare_log_dir(path: Path | str) -> Path:
     if path is None:
         return Path()
@@ -125,28 +299,6 @@ def _register_handler(handler: Handler, level: int, tag: str = '_default_') -> N
     logger.addHandler(handler)
     logger.setLevel(level)
 
-def _colorful_format(record):
-    time = formatter.formatTime(record, time_format)
-    if not record.name.startswith('nni.'):
-        return '[{}] ({}) {}'.format(time, record.name, record.msg % record.args)
-    if record.levelno >= logging.ERROR:
-        color = colorama.Fore.RED
-        level = 'ERROR: '
-    elif record.levelno >= logging.WARNING:
-        color = colorama.Fore.YELLOW
-        level = 'WARNING: '
-    elif record.levelno >= logging.INFO:
-        color = colorama.Fore.GREEN
-        level = ''
-    else:
-        color = colorama.Fore.BLUE
-        level = ''
-    msg = color + level + (record.msg % record.args) + colorama.Style.RESET_ALL
-    if record.levelno < logging.INFO:
-        return '[{}] {}:{} {}'.format(time, record.threadName, record.name, msg)
-    else:
-        return '[{}] {}'.format(time, msg)
-
 class _LogFileWrapper(TextIOBase):
     # wrap the logger file so that anything written to it will automatically get formatted
 

nni/tools/nnictl/launcher.py

@@ -11,7 +11,6 @@ import yaml
 
 from nni.experiment import Experiment, RunMode
 from nni.experiment.config import ExperimentConfig, convert, utils
-from nni.runtime.log import init_logger_for_command_line
 from nni.tools.annotation import expand_annotations, generate_search_space
 
 # used for v1-only legacy setup, remove them later
@@ -28,11 +27,6 @@ def create_experiment(args):
     url_prefix = args.url_prefix
     foreground = args.foreground
 
-    # it should finally be done in nnictl main function
-    # but for now don't break routines without logging support
-    init_logger_for_command_line()
-    logging.getLogger('nni').setLevel(logging.INFO)
-
     if not config_file.is_file():
         _logger.error(f'"{config_file}" is not a valid file.')
         exit(1)
@@ -104,9 +98,6 @@ def resume_experiment(args):
     foreground = args.foreground
     exp_dir = args.experiment_dir
 
-    init_logger_for_command_line()
-    logging.getLogger('nni').setLevel(logging.INFO)
-
     config_json = get_stopped_experiment_config_json(exp_id, exp_dir)
     if config_json.get('trainingServicePlatform'):
         legacy_launcher.resume_experiment(args)
@@ -121,9 +112,6 @@ def view_experiment(args):
     port = args.port
     exp_dir = args.experiment_dir
 
-    init_logger_for_command_line()
-    logging.getLogger('nni').setLevel(logging.INFO)
-
     config_json = get_stopped_experiment_config_json(exp_id, exp_dir)
     if config_json.get('trainingServicePlatform'):
         legacy_launcher.view_experiment(args)