File size: 1,100 Bytes
32b2aaa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import logging
from pathlib import Path

from rich.logging import RichHandler

from .distributed import global_leader_only


@global_leader_only
def setup_logging(run_dir):
    handlers = []
    stdout_handler = RichHandler()
    stdout_handler.setLevel(logging.INFO)
    handlers.append(stdout_handler)

    if run_dir is not None:
        filename = Path(run_dir) / f"log.txt"
        filename.parent.mkdir(parents=True, exist_ok=True)
        file_handler = logging.FileHandler(filename, mode="a")
        file_handler.setLevel(logging.DEBUG)
        handlers.append(file_handler)

    # Update all existing loggers
    for name in ["DeepSpeed"]:
        logger = logging.getLogger(name)
        if isinstance(logger, logging.Logger):
            for handler in list(logger.handlers):
                logger.removeHandler(handler)
            for handler in handlers:
                logger.addHandler(handler)

    # Set the default logger
    logging.basicConfig(
        level=logging.getLevelName("INFO"),
        format="%(message)s",
        datefmt="[%X]",
        handlers=handlers,
    )