Spaces:
Running
Running
| from datetime import datetime | |
| from io import TextIOBase | |
| import logging | |
| from logging import FileHandler, Formatter, Handler, StreamHandler | |
| from pathlib import Path | |
| import sys | |
| import time | |
| from typing import Optional | |
| import colorama | |
| from .env_vars import dispatcher_env_vars, trial_env_vars | |
| handlers = {} | |
| log_format = '[%(asctime)s] %(levelname)s (%(name)s/%(threadName)s) %(message)s' | |
| time_format = '%Y-%m-%d %H:%M:%S' | |
| formatter = Formatter(log_format, time_format) | |
| 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. | |
| 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() | |
| if dispatcher_env_vars.SDK_PROCESS == 'dispatcher': | |
| _init_logger_dispatcher() | |
| return | |
| trial_platform = trial_env_vars.NNI_PLATFORM | |
| if trial_platform == 'unittest': | |
| return | |
| if trial_platform and not trial_env_vars.REUSE_MODE: | |
| _init_logger_trial() | |
| return | |
| _init_logger_standalone() | |
| logging.getLogger('filelock').setLevel(logging.WARNING) | |
| _exp_log_initialized = False | |
| def init_logger_experiment() -> None: | |
| """ | |
| Initialize logger for `nni.experiment.Experiment`. | |
| This function will get invoked after `init_logger()`. | |
| """ | |
| global _exp_log_initialized | |
| if not _exp_log_initialized: | |
| _exp_log_initialized = True | |
| colorful_formatter = Formatter(log_format, time_format) | |
| colorful_formatter.format = _colorful_format | |
| 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.WARNING | |
| _register_handler(FileHandler(log_path), log_level, experiment_id) | |
| def stop_experiment_log(experiment_id: str) -> None: | |
| if experiment_id in handlers: | |
| logging.getLogger().removeHandler(handlers.pop(experiment_id)) | |
| def _init_logger_dispatcher() -> None: | |
| log_level_map = { | |
| 'fatal': logging.CRITICAL, | |
| 'error': logging.ERROR, | |
| 'warning': logging.WARNING, | |
| 'info': logging.WARNING, | |
| 'debug': logging.DEBUG, | |
| 'trace': 0 | |
| } | |
| log_path = _prepare_log_dir(dispatcher_env_vars.NNI_LOG_DIRECTORY) / 'dispatcher.log' | |
| log_level = log_level_map.get(dispatcher_env_vars.NNI_LOG_LEVEL, logging.WARNING) | |
| _register_handler(FileHandler(log_path), log_level) | |
| def _init_logger_trial() -> None: | |
| log_path = _prepare_log_dir(trial_env_vars.NNI_OUTPUT_DIR) / 'trial.log' | |
| log_file = open(log_path, 'a') | |
| _register_handler(StreamHandler(log_file), logging.WARNING) | |
| if trial_env_vars.NNI_PLATFORM == 'local': | |
| sys.stdout = _LogFileWrapper(log_file) | |
| def _init_logger_standalone() -> None: | |
| _register_handler(StreamHandler(sys.stdout), logging.WARNING) | |
| def _prepare_log_dir(path: Optional[str]) -> Path: | |
| if path is None: | |
| return Path() | |
| ret = Path(path) | |
| ret.mkdir(parents=True, exist_ok=True) | |
| return ret | |
| def _register_handler(handler: Handler, level: int, tag: str = '_default_') -> None: | |
| assert tag not in handlers | |
| handlers[tag] = handler | |
| handler.setFormatter(formatter) | |
| logger = logging.getLogger() | |
| 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 | |
| elif record.levelno >= logging.WARNING: | |
| color = colorama.Fore.YELLOW | |
| elif record.levelno >= logging.WARNING: | |
| color = colorama.Fore.GREEN | |
| else: | |
| color = colorama.Fore.BLUE | |
| msg = color + (record.msg % record.args) + colorama.Style.RESET_ALL | |
| if record.levelno < logging.WARNING: | |
| 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 | |
| def __init__(self, log_file: TextIOBase): | |
| self.file: TextIOBase = log_file | |
| self.line_buffer: Optional[str] = None | |
| self.line_start_time: Optional[datetime] = None | |
| def write(self, s: str) -> int: | |
| cur_time = datetime.now() | |
| if self.line_buffer and (cur_time - self.line_start_time).total_seconds() > 0.1: | |
| self.flush() | |
| if self.line_buffer: | |
| self.line_buffer += s | |
| else: | |
| self.line_buffer = s | |
| self.line_start_time = cur_time | |
| if '\n' not in s: | |
| return len(s) | |
| time_str = cur_time.strftime(time_format) | |
| lines = self.line_buffer.split('\n') | |
| for line in lines[:-1]: | |
| self.file.write(f'[{time_str}] PRINT {line}\n') | |
| self.file.flush() | |
| self.line_buffer = lines[-1] | |
| self.line_start_time = cur_time | |
| return len(s) | |
| def flush(self) -> None: | |
| if self.line_buffer: | |
| time_str = self.line_start_time.strftime(time_format) | |
| self.file.write(f'[{time_str}] PRINT {self.line_buffer}\n') | |
| self.file.flush() | |
| self.line_buffer = None | |