| 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, |
| ) -> None: |
| """Initializes a multi-GPU-friendly python command line logger that logs on all processes |
| with their rank prefixed in the log message. |
| |
| :param name: The name of the logger. Default is ``__name__``. |
| :param rank_zero_only: Whether to force all logs to only occur on the rank zero process. Default is `False`. |
| :param extra: (Optional) A dict-like object which provides contextual information. See `logging.LoggerAdapter`. |
| """ |
| logger = logging.getLogger(name) |
| 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) |
|
|
