| import logging |
| from typing import Mapping, Optional |
|
|
| from lightning_utilities.core.rank_zero import rank_prefixed_message, rank_zero_only |
|
|
|
|
| class RankedLogger(logging.LoggerAdapter): |
| """A multi-GPU-friendly python command line logger.""" |
|
|
| def __init__( |
| self, |
| name: str = __name__, |
| rank_zero_only: bool = False, |
| extra: Optional[Mapping[str, object]] = None, |
| log_file: str = "log.txt", |
| ) -> None: |
| logger = logging.getLogger(name) |
|
|
| |
| formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') |
|
|
| |
| file_handler = logging.FileHandler(log_file) |
| file_handler.setFormatter(formatter) |
| logger.addHandler(file_handler) |
|
|
| super().__init__(logger=logger, extra=extra) |
| self.rank_zero_only = rank_zero_only |
|
|
| def log(self, level: int, msg: str, rank: Optional[int] = None, *args, **kwargs) -> None: |
| """Delegate a log call to the underlying logger, after prefixing its message with the rank |
| of the process it's being logged from. If `'rank'` is provided, then the log will only |
| occur on that rank/process. |
| |
| :param level: The level to log at. Look at `logging.__init__.py` for more information. |
| :param msg: The message to log. |
| :param rank: The rank to log at. |
| :param args: Additional args to pass to the underlying logging function. |
| :param kwargs: Any additional keyword args to pass to the underlying logging function. |
| """ |
| if self.isEnabledFor(level): |
| msg, kwargs = self.process(msg, kwargs) |
| current_rank = getattr(rank_zero_only, "rank", None) |
| if current_rank is None: |
| raise RuntimeError("The `rank_zero_only.rank` needs to be set before use") |
| msg = rank_prefixed_message(msg, current_rank) |
| if self.rank_zero_only: |
| if current_rank == 0: |
| self.logger.log(level, msg, *args, **kwargs) |
| else: |
| if rank is None: |
| self.logger.log(level, msg, *args, **kwargs) |
| elif current_rank == rank: |
| self.logger.log(level, msg, *args, **kwargs) |
|
|
