Logging Utilities

fusion_bench.utils.rich_utils

display_available_styles()

Display all available styles in a grid.

Source code in fusion_bench/utils/rich_utils.py

def display_available_styles():
    """Display all available styles in a grid."""
    console = Console()
    style_samples = [
        Panel(f"Style: {style}", expand=False, border_style=style)
        for style in AVAILABLE_STYLES
    ]
    console.print(Columns(style_samples, equal=True, expand=False))

enforce_tags(cfg, save_to_file=False)

Prompts user to input tags from command line if no tags are provided in config.

:param cfg: A DictConfig composed by Hydra. :param save_to_file: Whether to export tags to the hydra output folder. Default is False.

Source code in fusion_bench/utils/rich_utils.py

@rank_zero_only
def enforce_tags(cfg: DictConfig, save_to_file: bool = False) -> None:
    """Prompts user to input tags from command line if no tags are provided in config.

    :param cfg: A DictConfig composed by Hydra.
    :param save_to_file: Whether to export tags to the hydra output folder. Default is ``False``.
    """
    if not cfg.get("tags"):
        if "id" in HydraConfig().cfg.hydra.job:
            raise ValueError("Specify tags before launching a multirun!")

        log.warning("No tags provided in config. Prompting user to input tags...")
        tags = Prompt.ask("Enter a list of comma separated tags", default="dev")
        tags = [t.strip() for t in tags.split(",") if t != ""]

        with open_dict(cfg):
            cfg.tags = tags

        log.info(f"Tags: {cfg.tags}")

    if save_to_file:
        with open(Path(cfg.paths.output_dir, "tags.log"), "w") as file:
            rich.print(cfg.tags, file=file)

print_bordered(message, title=None, style='blue', code_style=None)

Print a message with a colored border.

Args: message (str): The message to print. title (str, optional): The title of the panel. Defaults to None. style (str, optional): The color style for the border. Defaults to "cyan". code_style (str, optional): The syntax highlighting style if the message is code. Set to None for plain text. Defaults to "python".

Source code in fusion_bench/utils/rich_utils.py

def print_bordered(message, title=None, style="blue", code_style=None):
    """
    Print a message with a colored border.

    Args:
    message (str): The message to print.
    title (str, optional): The title of the panel. Defaults to None.
    style (str, optional): The color style for the border. Defaults to "cyan".
    code_style (str, optional): The syntax highlighting style if the message is code.
                                Set to None for plain text. Defaults to "python".
    """
    if code_style:
        content = Syntax(message, code_style, theme="monokai", word_wrap=True)
    else:
        content = Text(message)

    panel = Panel(content, title=title, border_style=style)
    print(panel)

print_config_tree(cfg, print_order=('data', 'model', 'callbacks', 'logger', 'trainer', 'paths', 'extras'), resolve=False, save_to_file=False)

Prints the contents of a DictConfig as a tree structure using the Rich library.

:param cfg: A DictConfig composed by Hydra. :param print_order: Determines in what order config components are printed. Default is ("data", "model", "callbacks", "logger", "trainer", "paths", "extras"). :param resolve: Whether to resolve reference fields of DictConfig. Default is False. :param save_to_file: Whether to export config to the hydra output folder. Default is False.

Source code in fusion_bench/utils/rich_utils.py

@rank_zero_only
def print_config_tree(
    cfg: DictConfig,
    print_order: Sequence[str] = (
        "data",
        "model",
        "callbacks",
        "logger",
        "trainer",
        "paths",
        "extras",
    ),
    resolve: bool = False,
    save_to_file: bool = False,
) -> None:
    """Prints the contents of a DictConfig as a tree structure using the Rich library.

    :param cfg: A DictConfig composed by Hydra.
    :param print_order: Determines in what order config components are printed. Default is ``("data", "model",
    "callbacks", "logger", "trainer", "paths", "extras")``.
    :param resolve: Whether to resolve reference fields of DictConfig. Default is ``False``.
    :param save_to_file: Whether to export config to the hydra output folder. Default is ``False``.
    """
    style = "tree"
    tree = rich.tree.Tree("CONFIG", style=style, guide_style=style)

    queue = []

    # add fields from `print_order` to queue
    for field in print_order:
        (
            queue.append(field)
            if field in cfg
            else log.warning(
                f"Field '{field}' not found in config. Skipping '{field}' config printing..."
            )
        )

    # add all the other fields to queue (not specified in `print_order`)
    for field in cfg:
        if field not in queue:
            queue.append(field)

    # generate config tree from queue
    for field in queue:
        branch = tree.add(field, style=style, guide_style=style)

        config_group = cfg[field]
        if isinstance(config_group, DictConfig):
            branch_content = OmegaConf.to_yaml(config_group, resolve=resolve)
        else:
            branch_content = str(config_group)

        branch.add(rich.syntax.Syntax(branch_content, "yaml"))

    # print config tree
    rich.print(tree)

    # save config tree to file
    if save_to_file:
        with open(Path(cfg.paths.output_dir, "config_tree.log"), "w") as file:
            rich.print(tree, file=file)

setup_colorlogging(force=False, level=logging.INFO, **kwargs)

Sets up color logging for the application.

Source code in fusion_bench/utils/rich_utils.py

def setup_colorlogging(
    force=False,
    level=logging.INFO,
    **kwargs,
):
    """
    Sets up color logging for the application.
    """
    FORMAT = "%(message)s"

    logging.basicConfig(
        level=level,
        format=FORMAT,
        datefmt="[%X]",
        handlers=[RichHandler()],
        force=force,
        **kwargs,
    )

fusion_bench.utils.pylogger

RankZeroLogger

Bases: Logger

A logger that logs only on rank zero and works just like logging.Logger

Source code in fusion_bench/utils/pylogger.py

class RankZeroLogger(logging.Logger):
    """A logger that logs only on rank zero and works just like logging.Logger"""

    @rank_zero_only
    def _log(self, *args, **kwargs):
        if "stacklevel" in kwargs:
            kwargs["stacklevel"] += 1
        else:
            kwargs["stacklevel"] = 2
        return super()._log(*args, **kwargs)

    def is_global_zero(self):
        return rank_zero_only.rank == 0

RankedLogger

Bases: LoggerAdapter

A multi-GPU-friendly python command line logger.

Source code in fusion_bench/utils/pylogger.py

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)

__init__(name=__name__, rank_zero_only=False, extra=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.

Source code in fusion_bench/utils/pylogger.py

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

log(level, msg, rank=None, *args, **kwargs)

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.

Source code in fusion_bench/utils/pylogger.py

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)

get_rankzero_logger(name=None)

Return a logger with the specified name, creating it if necessary.

If no name is specified, return the root logger.

Source code in fusion_bench/utils/pylogger.py

def get_rankzero_logger(name=None):
    """
    Return a logger with the specified name, creating it if necessary.

    If no name is specified, return the root logger.
    """
    if not name or isinstance(name, str) and name == logging.root.name:
        return logging.root
    return RankZeroLogger.manager.getLogger(name)