mirror of https://github.com/hpcaitech/ColossalAI
27 lines
946 B
Python
27 lines
946 B
Python
from colossalai.core import global_context as gpc
|
|
from .logging import DistributedLogger
|
|
|
|
__all__ = ['get_global_dist_logger', 'get_dist_logger', 'DistributedLogger', 'init_global_dist_logger']
|
|
|
|
_GLOBAL_LOGGER: DistributedLogger = None
|
|
|
|
|
|
def get_dist_logger(name, level='INFO', root_path: str = None, mode='a'):
|
|
return DistributedLogger(name=name, level=level, root_path=root_path, mode=mode)
|
|
|
|
|
|
def get_global_dist_logger():
|
|
assert _GLOBAL_LOGGER is not None, 'Global distributed logger is not initialized'
|
|
return _GLOBAL_LOGGER
|
|
|
|
|
|
def init_global_dist_logger():
|
|
rank = gpc.get_global_rank()
|
|
if hasattr(gpc.config, 'logging'):
|
|
logger = get_dist_logger(name=f'rank_{rank}', **gpc.config.logging)
|
|
else:
|
|
logger = get_dist_logger(name=f'rank_{rank}', level='INFO')
|
|
global _GLOBAL_LOGGER
|
|
assert _GLOBAL_LOGGER is None, 'Global distributed logger has already been initialized'
|
|
_GLOBAL_LOGGER = logger
|