File size: 2,346 Bytes
0839907 | 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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import functools
from typing import Callable
from loguru import logger
import sys
from fastgen.utils.distributed import get_rank, is_rank0
# Default logging level — can be changed later
LOG_LEVEL = "INFO"
logger.remove(0)
logger.add(
sys.stdout,
format="[<green>{time:MMM D, YYYY - HH:mm:ss}</green> | "
"<level>{level}</level> | "
"<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> ] {message}",
)
def formatter(record):
"""
Custom formatter function to conditionally add the MPI rank.
"""
# Base format string for all levels
msg = (
"[<green>{time:MMM D, YYYY - HH:mm:ss}</green> | "
"<level>{level}</level> | "
"<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> ] {message}\n"
)
# For DEBUG messages, add the rank
if LOG_LEVEL == "DEBUG":
msg = f"[<magenta>RANK: {get_rank()}</magenta>] {msg}"
# For all other levels, use the base format
return msg
def set_log_level(level: str):
"""Change log level dynamically."""
global LOG_LEVEL
LOG_LEVEL = level.upper()
logger.remove() # Remove all existing handlers
logger.add(
sys.stdout,
format=formatter,
level=LOG_LEVEL,
)
def rank0_if_not_debug(func: Callable) -> Callable:
"""
A decorator factory that applies the rank0_only decorator if LOG_LEVEL is not DEBUG.
"""
@functools.wraps(func)
def wrapper(*args, **kwargs):
# At runtime, check the global flag to decide which version to run.
if LOG_LEVEL == "DEBUG" or is_rank0():
return func(*args, **kwargs)
return wrapper
@rank0_if_not_debug
def trace(msg: str):
logger.opt(depth=2).trace(msg)
@rank0_if_not_debug
def info(msg: str):
logger.opt(depth=2).info(msg)
@rank0_if_not_debug
def debug(msg: str):
logger.opt(depth=2).debug(msg)
@rank0_if_not_debug
def success(msg: str):
logger.opt(depth=2).success(msg)
@rank0_if_not_debug
def critical(msg: str):
logger.opt(depth=2).critical(msg)
@rank0_if_not_debug
def warning(msg: str):
logger.opt(depth=2).warning(msg)
@rank0_if_not_debug
def error(msg: str):
logger.opt(depth=2).error(msg)
|