File size: 1,048 Bytes
d7b3a74 | 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 40 41 42 | import logging
import wandb
from . import wandb_utils
from .tensorboard_utils import _TensorboardAdapter
_LOGGER_CONFIGURED = False
# ref: SGLang
def configure_logger(prefix: str = ""):
global _LOGGER_CONFIGURED
if _LOGGER_CONFIGURED:
return
_LOGGER_CONFIGURED = True
logging.basicConfig(
level=logging.INFO,
format=f"[%(asctime)s{prefix}] %(filename)s:%(lineno)d - %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
force=True,
)
def init_tracking(args, primary: bool = True, **kwargs):
if primary:
wandb_utils.init_wandb_primary(args, **kwargs)
else:
wandb_utils.init_wandb_secondary(args, **kwargs)
# TODO further refactor, e.g. put TensorBoard init to the "init" part
def log(args, metrics, step_key: str):
if args.use_wandb:
wandb.log(metrics)
if args.use_tensorboard:
metrics_except_step = {k: v for k, v in metrics.items() if k != step_key}
_TensorboardAdapter(args).log(data=metrics_except_step, step=metrics[step_key])
|