pytorch将不同的卡的日志输出到不同的文件

import logging
import os
import torch.distributed as dist
from torch.distributed import init_process_group


def setup_logger(log_file):
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)

    # Create file handler which logs even debug messages
    fh = logging.FileHandler(log_file)
    fh.setLevel(logging.INFO)

    # Create console handler with a higher log level
    ch = logging.StreamHandler()
    ch.setLevel(logging.WARNING)

    # Create formatter and add it to the handlers
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    fh.setFormatter(formatter)
    ch.setFormatter(formatter)

    # Add the handlers to the logger
    logger.addHandler(fh)
    logger.addHandler(ch)
    return logger

if __name__ == "__main__":
    # Example log file: logs/log_rank_0.log, logs/log_rank_1.log, etc.
    # python -m torch.distributed.launch --nproc_per_node=4 test_data.py
    init_process_group(backend="nccl")
    rank = dist.get_rank()
    os.makedirs("logs", exist_ok=True)
    log_file = os.path.join("logs", f"log_rank_{rank}.log")
    logger = setup_logger(log_file)
    logger.info(f"Logging from rank {rank}")

然后命令行运行：

python -m torch.distributed.launch --nproc_per_node=4 test_data.py

结果展示：

其中一张卡的内容为：

2024-08-17 10:16:37,947 - root - INFO - Logging from rank 1