mirror of
https://git.datalinker.icu/ali-vilab/TeaCache
synced 2025-12-09 21:04:25 +08:00
33 lines
858 B
Python
33 lines
858 B
Python
import logging
|
|
|
|
import torch.distributed as dist
|
|
from rich.logging import RichHandler
|
|
|
|
|
|
def create_logger():
|
|
"""
|
|
Create a logger that writes to a log file and stdout.
|
|
"""
|
|
logger = logging.getLogger(__name__)
|
|
return logger
|
|
|
|
|
|
def init_dist_logger():
|
|
"""
|
|
Update the logger to write to a log file.
|
|
"""
|
|
global logger
|
|
if dist.get_rank() == 0:
|
|
logger = logging.getLogger(__name__)
|
|
handler = RichHandler(show_path=False, markup=True, rich_tracebacks=True)
|
|
formatter = logging.Formatter("VideoSys - %(levelname)s: %(message)s")
|
|
handler.setFormatter(formatter)
|
|
logger.addHandler(handler)
|
|
logger.setLevel(logging.INFO)
|
|
else: # dummy logger (does nothing)
|
|
logger = logging.getLogger(__name__)
|
|
logger.addHandler(logging.NullHandler())
|
|
|
|
|
|
logger = create_logger()
|