mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
d224ac7f77
Fixes https://github.com/pytorch/torchdynamo/issues/1932 Discussed with @mlazos: if we still want to separate streams for code logging and the rest of info, we can use a separate logger object with a unique name. Pull Request resolved: https://github.com/pytorch/pytorch/pull/90234 Approved by: https://github.com/ezyang
84 lines
2.1 KiB
Python
84 lines
2.1 KiB
Python
import itertools
|
|
import logging
|
|
import os
|
|
|
|
# Return all loggers that torchdynamo/torchinductor is responsible for
|
|
def get_loggers():
|
|
return [
|
|
logging.getLogger("torch._dynamo"),
|
|
logging.getLogger("torch._inductor"),
|
|
]
|
|
|
|
|
|
# Set the level of all loggers that torchdynamo is responsible for
|
|
def set_loggers_level(level):
|
|
for logger in get_loggers():
|
|
logger.setLevel(level)
|
|
|
|
|
|
LOGGING_CONFIG = {
|
|
"version": 1,
|
|
"formatters": {
|
|
"torchdynamo_format": {
|
|
"format": "[%(asctime)s] %(name)s: [%(levelname)s] %(message)s"
|
|
},
|
|
},
|
|
"handlers": {
|
|
"torchdynamo_console": {
|
|
"class": "logging.StreamHandler",
|
|
"level": "DEBUG",
|
|
"formatter": "torchdynamo_format",
|
|
"stream": "ext://sys.stderr",
|
|
},
|
|
},
|
|
"loggers": {
|
|
"torch._dynamo": {
|
|
"level": "DEBUG",
|
|
"handlers": ["torchdynamo_console"],
|
|
"propagate": False,
|
|
},
|
|
"torch._inductor": {
|
|
"level": "DEBUG",
|
|
"handlers": ["torchdynamo_console"],
|
|
"propagate": False,
|
|
},
|
|
},
|
|
"disable_existing_loggers": False,
|
|
}
|
|
|
|
|
|
# initialize torchdynamo loggers
|
|
def init_logging(log_level, log_file_name=None):
|
|
if "PYTEST_CURRENT_TEST" not in os.environ:
|
|
logging.config.dictConfig(LOGGING_CONFIG)
|
|
if log_file_name is not None:
|
|
log_file = logging.FileHandler(log_file_name)
|
|
log_file.setLevel(log_level)
|
|
for logger in get_loggers():
|
|
logger.addHandler(log_file)
|
|
|
|
set_loggers_level(log_level)
|
|
|
|
|
|
# Creates a logging function that logs a message with a step # prepended.
|
|
# get_step_logger should be lazily called (i.e. at runtime, not at module-load time)
|
|
# so that step numbers are initialized properly. e.g.:
|
|
|
|
# @functools.lru_cache(None)
|
|
# def _step_logger():
|
|
# return get_step_logger(logging.getLogger(...))
|
|
|
|
# def fn():
|
|
# _step_logger()(logging.INFO, "msg")
|
|
|
|
_step_counter = itertools.count(1)
|
|
|
|
|
|
def get_step_logger(logger):
|
|
step = next(_step_counter)
|
|
|
|
def log(level, msg):
|
|
logger.log(level, f"Step {step}: {msg}")
|
|
|
|
return log
|