Mike0021/zonos2

1from __future__ import annotations2 3from functools import partial4from typing import TYPE_CHECKING5 6_LOG_LEVEL = None7 8_LEVEL_MAP = {9    "DEBUG": 10,  # logging.DEBUG10    "INFO": 20,  # logging.INFO11    "WARNING": 30,  # logging.WARNING12    "ERROR": 40,  # logging.ERROR13    "CRITICAL": 50,  # logging.CRITICAL14}15 16 17def set_log_level(level: int | str) -> None:18    """Set the global log level for all zonos2 loggers.19 20    Args:21        level: Either a logging level int (e.g. logging.DEBUG) or a string22               like "DEBUG", "INFO", etc.23    """24    import logging25 26    global _LOG_LEVEL27    if isinstance(level, str):28        level = _LEVEL_MAP.get(level.upper(), logging.INFO)29    _LOG_LEVEL = level30 31    # Update all existing loggers under the zonos2 namespace32    for name, obj in logging.Logger.manager.loggerDict.items():33        if isinstance(obj, logging.Logger) and name.startswith("zonos2"):34            obj.setLevel(level)35 36 37def init_logger(38    name: str,39    suffix: str = "",40    *,41    strip_file: bool = True,42    level: str | None = None,43    use_pid: bool | None = None,44    use_tp_rank: bool | None = None,45):46    """Initialize the logger for the module with colors and pretty formatting."""47    import logging48    import os49    import sys50 51    global _LOG_LEVEL52    if _LOG_LEVEL is None:53        level = level or os.getenv("LOG_LEVEL", "").upper()54        _LOG_LEVEL = _LEVEL_MAP.get(level, logging.INFO)55 56    if strip_file:57        suffix = os.path.basename(suffix)58 59    if suffix:60        suffix = f"|{suffix}"61 62    if use_pid is None:63        use_pid = os.getenv("LOG_PID", "0").lower() in ("1", "true", "yes")64 65    if use_pid:66        pid = os.getpid()67        suffix = f"|pid={pid}{suffix}"68 69    tp_info = None70 71    # Color formatter class72    class ColorFormatter(logging.Formatter):73        """Formatter with colors and pretty output"""74 75        # ANSI color codes76        COLORS = {77            "DEBUG": "\033[36m",  # Cyan78            "INFO": "\033[32m",  # Green79            "WARNING": "\033[33m",  # Yellow80            "ERROR": "\033[31m",  # Red81            "CRITICAL": "\033[35m",  # Magenta82        }83        RESET = "\033[0m"84        BOLD = "\033[1m"85 86        def format(self, record):87            from zonos2.distributed import try_get_tp_info88 89            # Format timestamp like SGLang: [YYYY-MM-DD|HH:MM:SS|pid=1234]90            timestamp = self.formatTime(record, "[%Y-%m-%d|%H:%M:%S{suffix}]")91            nonlocal tp_info92            tp_info = tp_info or try_get_tp_info()93            if tp_info is not None and use_tp_rank is not False:94                real_suffix = f"{suffix}|core|rank={tp_info.rank}"95            else:96                real_suffix = suffix97            timestamp = timestamp.format(suffix=real_suffix)98 99            # Get color for log level100            level_color = self.COLORS.get(record.levelname, "")101 102            # Format the message103            colored_level = f"{level_color}{record.levelname:<8}{self.RESET}"104            message = record.getMessage()105 106            # Pretty format: [timestamp] LEVEL message107            return f"{self.BOLD}{timestamp}{self.RESET} {colored_level} {message}"108 109    logger = logging.getLogger(name)110    logger.setLevel(_LOG_LEVEL)111 112    # Clear existing handlers to avoid duplicates113    logger.handlers.clear()114 115    handler = logging.StreamHandler(sys.stdout)116    formatter = ColorFormatter()117    handler.setFormatter(formatter)118    logger.addHandler(handler)119 120    # Prevent propagation to root logger121    logger.propagate = False122 123    def _call_rank0(msg, *args, _which, **kwargs):124        from zonos2.distributed import get_tp_info125 126        nonlocal tp_info127        tp_info = tp_info or get_tp_info()128        assert tp_info is not None, "TP info not set yet"129        if tp_info.is_primary():130            getattr(logger, _which)(msg, *args, **kwargs)131 132    if TYPE_CHECKING:133 134        class WrapperLogger(logging.Logger):135            """Custom logger to handle the color formatter."""136 137            def info_rank0(self, msg, *args, **kwargs): ...138            def warning_rank0(self, msg, *args, **kwargs): ...139            def debug_rank0(self, msg, *args, **kwargs): ...140            def critical_rank0(self, msg, *args, **kwargs): ...141 142        return WrapperLogger(name)143    else:144        logger.info_rank0 = partial(_call_rank0, _which="info")145        logger.debug_rank0 = partial(_call_rank0, _which="debug")146        logger.critical_rank0 = partial(_call_rank0, _which="critical")147        logger.warning_rank0 = partial(_call_rank0, _which="warning")148        return logger149