TIME - MoE 模型代码 5——Time-MoE-main/time_moe/utils/log_util.py

源码：https://github.com/Time-MoE/Time-MoE

这段代码实现了一个自定义的日志配置模块，主要用于分布式训练环境下的日志管理，确保只有主进程（local_rank=0）输出日志，避免多进程环境下的日志冗余。

---

1. 日志基础配置

DEFAULT_LOGGER = "time_moe_logger"
DEFAULT_FORMATTER = logging.Formatter('%(asctime)s - %(filename)s[pid:%(process)d;line:%(lineno)d:%(funcName)s] - %(levelname)s: %(message)s'
)
_ch = logging.StreamHandler(stream=sys.stdout)
_ch.setFormatter(DEFAULT_FORMATTER)
_DEFAULT_HANDLERS = [_ch]

• 日志格式：
  定义了详细的日志格式，包含：

  • asctime：日志时间
  • filename：文件名
  • process：进程 ID
  • lineno：代码行号
  • funcName：函数名
  • levelname：日志级别（INFO/WARNING/ERROR 等）
  • message：日志内容
    这种格式便于定位日志对应的代码位置，适合调试和生产环境。

• 处理程序：
  创建一个输出到标准输出（stdout）的处理器（StreamHandler），并绑定上述格式器。

2. 分布式环境检测函数

def is_local_rank_0():local_rank = os.getenv('LOCAL_RANK')if local_rank is None or local_rank == '0':return Trueelse:return False

• 功能：检测当前进程是否是分布式环境中的主进程（local_rank=0）。
• 应用场景：在分布式训练（如使用 PyTorch DDP）中，通常只需要主进程输出日志，避免多进程重复打印。

3. 日志记录器工厂函数

def get_logger(name, level="INFO", handlers=None, update=False):if name in _LOGGER_CACHE and not update:return _LOGGER_CACHE[name]logger = logging.getLogger(name)logger.setLevel(level)logger.handlers = handlers or _DEFAULT_HANDLERSlogger.propagate = False  # 防止日志向上传播到根记录器_LOGGER_CACHE[name] = loggerreturn logger

• 单例模式：
  使用_LOGGER_CACHE缓存已创建的日志记录器，避免重复创建，提高性能。

  • 首次调用时创建新的记录器并缓存。
  • 后续调用直接从缓存获取（update=False时）。

• 参数说明：

  • name：日志记录器名称（如time_moe_logger）。
  • level：日志级别（默认INFO，支持DEBUG/WARNING/ERROR等）。
  • handlers：自定义处理器（默认使用_DEFAULT_HANDLERS，即标准输出）。
  • update：是否强制更新缓存（默认False）。

4. 分布式日志输出函数

def log_in_local_rank_0(*msg, type='info', used_logger=None):msg = ' '.join([str(s) for s in msg])used_logger = used_logger or logger  # 默认使用全局loggerif is_local_rank_0():  # 仅主进程输出日志if type == 'warn' or type == 'warning':used_logger.warning(msg)elif type == 'error':used_logger.error(msg)else:used_logger.info(msg)

• 核心逻辑：

  • 将输入参数拼接为字符串msg。
  • 通过is_local_rank_0()判断是否为主进程，若是则根据type记录对应级别的日志。
  • 支持三种日志类型：info（默认）、warn、error。

• 使用场景：
  在分布式训练中，确保只有主进程输出日志，避免多卡环境下日志重复（如 DDP 的 8 卡训练中，仅local_rank=0的卡打印日志）。

5. 全局日志记录器

logger = get_logger(DEFAULT_LOGGER)

• 创建一个默认的日志记录器，名称为time_moe_logger，使用前面定义的格式和处理器。
• 其他模块可直接导入该logger，通过log_in_local_rank_0函数控制输出。