johannesschmude's picture
Initial commit
b73936d
raw
history blame
3.92 kB
import functools
import logging
import os
import sys
from time import time
from packaging.version import Version
import wandb
from typing import Dict, Optional, Any
if Version(wandb.__version__) < Version("0.20.0"):
WANDB_USE_SYNC = True
else:
WANDB_USE_SYNC = False
def log(
run,
data: Dict[str, Any],
step: Optional[int] = None,
commit: Optional[bool] = None,
sync: Optional[bool] = None,
) -> None:
if run is not None:
# Note: wandb changed the .log API with version 0.20.0.
# This includes: "Removed no-op sync argument from wandb.Run::log function"
# We didn't test whether sync has any function here. But since we did
# all our development with it, let's keep it here for now.
# See https://github.com/wandb/wandb/releases/tag/v0.20.0
if WANDB_USE_SYNC:
run.log(data, step, commit, sync)
else:
run.log(data, step, commit)
else:
print(data)
# See: https://github.com/microsoft/Swin-Transformer/blob/main/logger.py
# See: https://github.com/Meituan-AutoML/Twins/blob/main/logger.py
def create_logger(output_dir: str, dist_rank: int, name: str) -> logging.Logger:
# create logger
logger = logging.getLogger(name)
logger.setLevel(logging.DEBUG)
logger.propagate = False
# create formatter
fmt = "[%(asctime)s %(name)s]: %(levelname)s %(message)s"
# create console handlers
if name.endswith("main"):
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(logging.INFO)
console_handler.setFormatter(
logging.Formatter(fmt=fmt, datefmt="%Y-%m-%d %H:%M:%S")
)
logger.addHandler(console_handler)
# create file handlers
file_handler = logging.FileHandler(
os.path.join(output_dir, f"{name}.log"), mode="a"
)
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(logging.Formatter(fmt=fmt, datefmt="%Y-%m-%d %H:%M:%S"))
logger.addHandler(file_handler)
return logger
def log_decorator(logger, _func=None):
def log_decorator_info(func):
@functools.wraps(func)
def log_decorator_wrapper(*args, **kwargs):
"""Create a list of the positional arguments passed to function.
- Using repr() for string representation for each argument. repr() is similar to str() only
difference being it prints with a pair of quotes and if we calculate a value we get more
precise value than str().
"""
# py_file_caller = getframeinfo(stack()[1][0])
local_rank = os.environ.get("LOCAL_RANK", default=None)
rank = os.environ.get("LOCAL_RANK", default=None)
try:
"""log return value from the function"""
start_time = time()
value = func(*args, **kwargs)
if local_rank is None or rank is None:
logger.info(
f"Function '{func.__name__}' - Execution time: {(time() - start_time):.1f} seconds."
)
else:
logger.info(
f"Function '{func.__name__}' - Execution time: {(time() - start_time):.1f} "
f"seconds on rank {os.environ['RANK']} and local_rank {os.environ['LOCAL_RANK']}."
)
except Exception as err:
logger.error(f"Exception: {err}")
raise
return value
# Return the pointer to the function
return log_decorator_wrapper
# Decorator was called with arguments, so return a decorator function that can read and return a function
if _func is None:
return log_decorator_info
# Decorator was called without arguments, so apply the decorator to the function immediately
else:
return log_decorator_info(_func)