如何在WandB项目中自定义日志输出方案

背景介绍

WandB是一个流行的机器学习实验跟踪工具，它默认会向终端输出大量运行信息。在实际项目中，我们经常需要将这些输出信息整合到自己的日志系统中，以便统一管理和分析。

问题分析

WandB默认使用自己的输出函数(termlog/termwarn/termerror)而非Python标准logging模块，这导致开发者无法直接通过Python的logging配置来捕获这些输出。虽然可以通过获取"wandb"日志器来添加处理器，但这种方法会捕获到大量调试信息，而无法精确控制需要记录的内容。

解决方案

我们可以通过重写WandB的输出函数来实现自定义日志记录。以下是具体实现步骤：

1. 创建自定义日志器

首先需要创建一个独立的日志器实例，并配置所需的处理器：

import logging
from rich.logging import RichHandler

# 创建自定义日志器
custom_logger = logging.getLogger("wandb_terminal")
custom_logger.setLevel(logging.INFO)
custom_logger.propagate = False  # 防止日志向上传播

# 添加Rich处理器用于控制台输出
rich_handler = RichHandler()
custom_logger.addHandler(rich_handler)

# 添加文件处理器用于持久化日志
file_handler = logging.FileHandler("wandb_redirected_terminal.log")
formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
file_handler.setFormatter(formatter)
custom_logger.addHandler(file_handler)

2. 重写WandB输出函数

保存原始函数引用，然后创建自定义函数来同时调用原始函数和自定义日志记录：

import wandb

# 保存原始函数引用
original_termlog = wandb.termlog
original_termwarn = wandb.termwarn
original_termerror = wandb.termerror

def custom_termlog(string="", newline=True, repeat=True):
    original_termlog(string)  # 保持原有终端输出
    if string:
        custom_logger.info(string)  # 记录到自定义日志器
        
def custom_termwarn(string="", newline=True, repeat=True):
    original_termwarn(string)
    if string:
        custom_logger.warning(string)
        
def custom_termerror(string="", newline=True, repeat=True):
    original_termerror(string)
    if string:
        custom_logger.error(string)

3. 替换WandB默认函数

# 替换WandB的默认输出函数
wandb.termlog = custom_termlog
wandb.termwarn = custom_termwarn
wandb.termerror = custom_termerror

4. 初始化WandB并测试

# 初始化WandB
wandb.init(project="misc-test")

# 测试日志输出
for i in range(10):
    wandb.log({"step": i})

wandb.finish()

技术要点解析

1. 日志器隔离：创建独立的日志器实例可以避免干扰其他模块的日志配置。

2. 函数重写：通过保存原始函数引用并创建包装函数，我们可以在保持原有功能的同时添加自定义行为。

3. 日志级别控制：可以根据需要调整日志级别，只记录重要信息。

4. 格式统一：通过自定义Formatter可以统一所有日志的输出格式。

扩展应用

这种技术方案不仅适用于WandB，也可以应用于其他使用自定义输出函数的库。通过类似的方法，我们可以：

1. 将不同来源的日志统一到同一个系统中
2. 添加额外的日志处理逻辑（如过滤敏感信息）
3. 实现日志的多目的地输出（控制台、文件、网络等）
4. 集成更丰富的日志展示工具（如Rich）

注意事项

1. 确保在WandB初始化前完成函数替换
2. 考虑线程安全性，特别是在多线程环境中
3. 对于生产环境，建议添加日志轮转机制防止日志文件过大
4. 可以根据项目需要调整日志格式和内容

通过这种方案，开发者可以灵活地将WandB的输出集成到自己的日志系统中，实现更专业的日志管理和分析。