File size: 1,312 Bytes
feba2ad |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
"""
Miscellaneous logging utilities.
"""
from io import StringIO
import yaml
from lightning.fabric.utilities.rank_zero import rank_zero_only
from rich.console import Console
from rich.panel import Panel
@rank_zero_only
def pretty_print_yaml_config(logger, config: dict) -> None:
"""
Pretty print config with rich formatting. Assumes that the config is already saved as a
dictionary - this can be done by calling `asdict` on the dataclass or loading in the config
from a yaml file.
NOTE: this function is only called on rank 0.
Args:
logger: Logger object to log the formatted output to.
config: Dictionary containing the config to pretty print.
"""
# Create string buffer
output = StringIO()
console = Console(file=output, force_terminal=False)
# Convert to YAML string first
yaml_str = yaml.dump(
config, default_flow_style=False, sort_keys=False, Dumper=yaml.SafeDumper
)
# Create formatted panel
panel = Panel(
yaml_str,
border_style="blue",
padding=(0, 1), # Reduced padding
expand=False, # Don't expand to terminal width
)
# Print to buffer
console.print(panel)
# Log the formatted output
for line in output.getvalue().splitlines():
logger.info(line)
|