logger_tools.py

import os
import inspect
import sys
import yaml
import wandb
from loguru import logger


def setup_logger(save_dir, distributed_rank=0, filename="log.txt", mode="a"):
    """setup logger for training and testing.
    Args:
        save_dir(str): location to save log file
        distributed_rank(int): device rank when multi-gpu environment
        filename (string): log save name.
        mode(str): log file write mode, `append` or `override`. default is `a`.

    Return:
        logger instance.
    """
    loguru_format = (
        "<blue>{time: MM-DD HH:mm:ss}</blue> | "
        #"<level>{level: <8}</level> | "
        #"<cyan>{name}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>"
        "<level>{message}</level>"
    )

    logger.remove()
    save_file = os.path.join(save_dir, filename)
    if mode == "o" and os.path.exists(save_file):
        os.remove(save_file)
    # only keep logger in rank0 process
    if distributed_rank == 0:
        logger.add(
            sys.stderr,
            format=loguru_format,
            level="INFO",
            enqueue=True,
        )
        logger.add(save_file,
            format=loguru_format,       
                  )


def set_args_and_logger(args, rank):
    """
    set logger file and print args
    """
    args_name_dir = args.root_path+args.out_root_path + "custom/" + args.name + args.notes + "/"
    if not os.path.exists(args_name_dir): os.makedirs(args_name_dir)
    args_name = args_name_dir + "/" + args.name +".yaml"
    
    os.environ["WANDB_MODE"] = "offline"
    
    if rank == 0:
        wandb.init(project=args.project, entity="liu1997", dir=args.root_path+args.out_root_path, name=args.name[12:]+ args.notes)
        wandb.config.update(args)

    if os.path.exists(args_name):
        s_add = 10
        logger.warning(f"Already exist args, add {s_add} to ran_seed to continue training")
        args.random_seed += s_add
    else:
        with open(args_name, "w+") as f:
            yaml.dump(args.__dict__, f) #bug for load list
            #json.dump(args.__dict__, f)
    setup_logger(args_name_dir, rank, filename=f"{args.name}.txt")