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)

get_syntax_theme()

Return (theme_name, background_color) suited to the current terminal.

Use this instead of hard-coding "monokai" when creating :class:rich.syntax.Syntax objects.

Source code in fusion_bench/utils/rich_utils.py

def get_syntax_theme() -> Tuple[str, Optional[str]]:
    """Return ``(theme_name, background_color)`` suited to the current terminal.

    Use this instead of hard-coding ``"monokai"`` when creating
    :class:`rich.syntax.Syntax` objects.
    """
    global _terminal_theme_cache

    if _terminal_theme_cache is None:
        _terminal_theme_cache = _detect_terminal_theme()

    if _terminal_theme_cache == "light":
        return _LIGHT_THEME, "default"
    return _DARK_THEME, "default"

print_bordered(message, title=None, style='blue', code_style=None, *, expand=True, theme=None, background_color='default', print_fn=print, format_code=True)

Print a message with a colored border.

Parameters:

• message (str) –

  The message to print.

• title (str, default: None ) –

  The title of the panel. Defaults to None.

• style (str, default: 'blue' ) –

  The color style for the border. Defaults to "cyan".

• code_style (str, default: None ) –

  The syntax highlighting style if the message is code. Set to None for plain text. Defaults to "python".

• theme (str, default: None ) –

  Syntax highlighting theme. Defaults to auto-detection based on terminal background (dark theme for dark terminals, light theme for light terminals).

Source code in fusion_bench/utils/rich_utils.py

def print_bordered(
    message,
    title=None,
    style="blue",
    code_style=None,
    *,
    expand: bool = True,
    theme: Optional[str] = None,
    background_color: Optional[str] = "default",
    print_fn=print,
    format_code: bool = True,
):
    """
    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".
        theme (str, optional): Syntax highlighting theme. Defaults to auto-detection
            based on terminal background (dark theme for dark terminals, light theme
            for light terminals).
    """
    if theme is None:
        theme, background_color = get_syntax_theme()
    if code_style:
        if format_code:
            message = format_code_str(message, code_style)
        content = Syntax(
            message,
            code_style,
            word_wrap=True,
            theme=theme,
            background_color=background_color,
        )
    else:
        content = Text(message)

    panel = Panel(content, title=title, border_style=style, expand=expand)
    print_fn(panel)

print_code(message, title=None, code_style=None, *, expand=True, theme=None, background_color='default', print_fn=print)

Print code or plain text with optional syntax highlighting.

Parameters:

• message (str) –

  The message or code to print.

• title (str, default: None ) –

  Optional title associated with this output. Currently not used by this function, but kept for API compatibility. Defaults to None.

• code_style (str, default: None ) –

  The language/lexer name for syntax highlighting (for example, "python"). If None, the message is rendered as plain text without syntax highlighting. Defaults to None.

• expand (bool, default: True ) –

  Placeholder flag for API symmetry with other printing helpers. It is not used in the current implementation. Defaults to True.

• theme (str, default: None ) –

  Name of the Rich syntax highlighting theme to use when code_style is provided. Defaults to auto-detection based on terminal background.

• background_color (str, default: 'default' ) –

  Background color style to apply to the code block when using syntax highlighting. Defaults to "default".

• print_fn (Callable, default: print ) –

  Function used to render the resulting Rich object. Defaults to :func:rich.print.

Source code in fusion_bench/utils/rich_utils.py

def print_code(
    message,
    title=None,
    code_style=None,
    *,
    expand: bool = True,
    theme: Optional[str] = None,
    background_color: Optional[str] = "default",
    print_fn=print,
):
    """
    Print code or plain text with optional syntax highlighting.

    Args:
        message (str): The message or code to print.
        title (str, optional): Optional title associated with this output. Currently
            not used by this function, but kept for API compatibility. Defaults to None.
        code_style (str, optional): The language/lexer name for syntax highlighting
            (for example, ``"python"``). If ``None``, the message is rendered as plain
            text without syntax highlighting. Defaults to ``None``.
        expand (bool, optional): Placeholder flag for API symmetry with other printing
            helpers. It is not used in the current implementation. Defaults to True.
        theme (str, optional): Name of the Rich syntax highlighting theme to use when
            ``code_style`` is provided. Defaults to auto-detection based on terminal
            background.
        background_color (str, optional): Background color style to apply to the code
            block when using syntax highlighting. Defaults to ``"default"``.
        print_fn (Callable, optional): Function used to render the resulting Rich
            object. Defaults to :func:`rich.print`.
    """
    if theme is None:
        theme, background_color = get_syntax_theme()
    if code_style:
        content = Syntax(
            message,
            code_style,
            word_wrap=True,
            theme=theme,
            background_color=background_color,
        )
    else:
        content = Text(message)

    print_fn(content)

print_config_tree(cfg, print_order=('data', 'model', 'callbacks', 'logger', 'trainer', 'path', 'extras'), resolve=False, save_to_file=False, *, theme=None, background_color='default')

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

Parameters:

• cfg (DictConfig) –

  A DictConfig composed by Hydra.

• print_order (Sequence[str], default: ('data', 'model', 'callbacks', 'logger', 'trainer', 'path', 'extras') ) –

  Determines in what order config components are printed. Defaults to ("data", "model", "callbacks", "logger", "trainer", "paths", "extras").

• resolve (bool, default: False ) –

  Whether to resolve reference fields of DictConfig. Defaults to False.

• save_to_file (bool, default: False ) –

  Whether to export config to the hydra output folder. Defaults to False.

• theme (str, default: None ) –

  The theme to use for syntax highlighting. Defaults to auto-detection based on terminal background.

• background_color (str, default: 'default' ) –

  The background color to use for syntax highlighting. Defaults to "default".

Returns:

• None –

  None

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",
        "path",
        "extras",
    ),
    resolve: bool = False,
    save_to_file: bool = False,
    *,
    theme: Optional[str] = None,
    background_color: Optional[str] = "default",
) -> None:
    """Prints the contents of a DictConfig as a tree structure using the Rich library.

    Args:
        cfg (DictConfig): A DictConfig composed by Hydra.
        print_order (Sequence[str], optional): Determines in what order config components are printed.
            Defaults to ``("data", "model", "callbacks", "logger", "trainer", "paths", "extras")``.
        resolve (bool, optional): Whether to resolve reference fields of DictConfig.
            Defaults to ``False``.
        save_to_file (bool, optional): Whether to export config to the hydra output folder.
            Defaults to ``False``.
        theme (str, optional): The theme to use for syntax highlighting. Defaults to
            auto-detection based on terminal background.
        background_color (str, optional): The background color to use for syntax highlighting.
            Defaults to "default".

    Returns:
        None
    """
    if theme is None:
        theme, background_color = get_syntax_theme()
    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..."
            )
        )

    # 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).strip()
        else:
            branch_content = str(config_group)

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

    # add all the other fields to queue (not specified in `print_order`)
    other_fields = [field for field in cfg if field not in queue]
    if other_fields:
        others_branch = tree.add(Text("[others]"), style=style, guide_style=style)

        other_cfg = OmegaConf.create({field: cfg[field] for field in other_fields})
        branch_content = OmegaConf.to_yaml(other_cfg, resolve=resolve).strip()

        others_branch.add(
            rich.syntax.Syntax(
                branch_content, "yaml", theme=theme, background_color=background_color
            )
        )

    # print config tree
    rich.print(tree)

    # save config tree to file
    if save_to_file:
        if not cfg.get("paths") or not cfg.paths.get("output_dir"):
            log.error(
                "Cannot save config tree to file. 'paths.output_dir' is not specified in the config."
            )
        else:
            with open(Path(cfg.path.output_dir, "config_tree.log"), "w") as file:
                rich.print(tree, file=file)

print_config_yaml(cfg, resolve=False, output_path=False, *, theme=None, background_color='default')

Prints the contents of a DictConfig as a YAML string using the Rich library.

Parameters:

• cfg (DictConfig) –

  A DictConfig composed by Hydra.

• resolve (bool, default: False ) –

  Whether to resolve reference fields of DictConfig. Default is False.

• output_path (Optional[str], default: False ) –

  Optional path to export the config YAML to. If provided, the file is written to this path.

• theme (str, default: None ) –

  The theme to use for syntax highlighting. Defaults to auto-detection based on terminal background.

Source code in fusion_bench/utils/rich_utils.py

@rank_zero_only
def print_config_yaml(
    cfg: DictConfig,
    resolve: bool = False,
    output_path: Optional[str] = False,
    *,
    theme: Optional[str] = None,
    background_color: Optional[str] = "default",
) -> None:
    """
    Prints the contents of a DictConfig as a YAML string using the Rich library.

    Args:
        cfg: A DictConfig composed by Hydra.
        resolve: Whether to resolve reference fields of DictConfig. Default is ``False``.
        output_path: Optional path to export the config YAML to. If provided, the file is written to this path.
        theme (str, optional): The theme to use for syntax highlighting. Defaults to
            auto-detection based on terminal background.
    """
    if theme is None:
        theme, background_color = get_syntax_theme()
    config_yaml = OmegaConf.to_yaml(cfg, resolve=resolve)
    syntax = rich.syntax.Syntax(
        config_yaml, "yaml", theme=theme, background_color=background_color
    )
    rich.print(syntax)

    if output_path:
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        with open(Path(output_path), "w") as file:
            rich.print(syntax, 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)