mlflow--mlflow/mlflow/utils/logging_utils.py

import contextlib
import logging
import logging.config
import re
import sys

from mlflow.environment_variables import MLFLOW_LOGGING_LEVEL
from mlflow.utils.thread_utils import ThreadLocalVariable


def get_mlflow_log_level() -> str:
    """Returns the log level from MLFLOW_LOGGING_LEVEL env var, defaulting to INFO."""
    return (MLFLOW_LOGGING_LEVEL.get() or "INFO").upper()


# Logging format example:
# 2018/11/20 12:36:37 INFO mlflow.sagemaker: Creating new SageMaker endpoint
LOGGING_LINE_FORMAT = "%(asctime)s %(levelname)s %(name)s: %(message)s"
LOGGING_DATETIME_FORMAT = "%Y/%m/%d %H:%M:%S"


class MlflowLoggingStream:
    """
    A Python stream for use with event logging APIs throughout MLflow (`eprint()`,
    `logger.info()`, etc.). This stream wraps `sys.stderr`, forwarding `write()` and
    `flush()` calls to the stream referred to by `sys.stderr` at the time of the call.
    It also provides capabilities for disabling the stream to silence event logs.
    """

    def __init__(self):
        self._enabled = True

    def write(self, text):
        if self._enabled:
            sys.stderr.write(text)

    def flush(self):
        if self._enabled:
            sys.stderr.flush()

    @property
    def enabled(self):
        return self._enabled

    @enabled.setter
    def enabled(self, value):
        self._enabled = value


MLFLOW_LOGGING_STREAM = MlflowLoggingStream()


def disable_logging():
    """
    Disables the `MlflowLoggingStream` used by event logging APIs throughout MLflow
    (`eprint()`, `logger.info()`, etc), silencing all subsequent event logs.
    """
    MLFLOW_LOGGING_STREAM.enabled = False


def enable_logging():
    """
    Enables the `MlflowLoggingStream` used by event logging APIs throughout MLflow
    (`eprint()`, `logger.info()`, etc), emitting all subsequent event logs. This
    reverses the effects of `disable_logging()`.
    """
    MLFLOW_LOGGING_STREAM.enabled = True


class MlflowFormatter(logging.Formatter):
    """
    Custom Formatter Class to support colored log
    ANSI characters might not work natively on older Windows, so disabling the feature for win32.
    See https://github.com/borntyping/python-colorlog/blob/dfa10f59186d3d716aec4165ee79e58f2265c0eb/colorlog/escape_codes.py#L16C8-L16C31
    """

    # Copied from color log package https://github.com/borntyping/python-colorlog/blob/dfa10f59186d3d716aec4165ee79e58f2265c0eb/colorlog/escape_codes.py#L33-L50
    COLORS = {
        "black": 30,
        "red": 31,
        "green": 32,
        "yellow": 33,
        "blue": 34,
        "purple": 35,
        "cyan": 36,
        "white": 37,
        "light_black": 90,
        "light_red": 91,
        "light_green": 92,
        "light_yellow": 93,
        "light_blue": 94,
        "light_purple": 95,
        "light_cyan": 96,
        "light_white": 97,
    }
    RESET = "\033[0m"

    def format(self, record):
        if color := getattr(record, "color", None):
            if color in self.COLORS and sys.platform != "win32":
                color_code = self._escape(self.COLORS[color])
                return f"{color_code}{super().format(record)}{self.RESET}"
        return super().format(record)

    def _escape(self, code: int) -> str:
        return f"\033[{code}m"


# Thread-local variable to suppress logs in the certain thread, used
# in telemetry client to suppress logs in the consumer thread
should_suppress_logs_in_thread = ThreadLocalVariable(default_factory=lambda: False)


class SuppressLogFilter(logging.Filter):
    def filter(self, record):
        if should_suppress_logs_in_thread.get():
            return False
        return super().filter(record)


# Cloud storage presigned URLs (AWS S3, GCS, Azure SAS) embed credentials in
# query parameters. urllib3 logs the full target URL on connection retries
# (e.g. when DNS resolution fails), which leaks those credentials to stderr.
_SENSITIVE_QUERY_PARAM_NAMES = (
    "X-Amz-Signature",
    "X-Amz-Security-Token",
    "X-Amz-Credential",
    "X-Goog-Signature",
    "X-Goog-Credential",
    "Signature",
    "sig",
)

_SENSITIVE_QUERY_PARAM_PATTERN = re.compile(
    r"\b(" + "|".join(re.escape(p) for p in _SENSITIVE_QUERY_PARAM_NAMES) + r")=[^&\s'\"]+",
    re.IGNORECASE,
)


def _redact_sensitive_query_params(message: str) -> str:
    return _SENSITIVE_QUERY_PARAM_PATTERN.sub(lambda m: f"{m.group(1)}=[REDACTED]", message)


class SensitiveQueryParamFilter(logging.Filter):
    """
    Logging filter that masks cloud storage credentials embedded in URL query
    strings. Attached to urllib3 to avoid leaking presigned URL credentials
    (X-Amz-Signature, X-Amz-Security-Token, X-Amz-Credential, ...) when
    urllib3 logs full URLs on connection failures or retries.
    """

    def filter(self, record: logging.LogRecord) -> bool:
        try:
            message = record.getMessage()
        except Exception:
            return True
        redacted = _redact_sensitive_query_params(message)
        if redacted != message:
            record.msg = redacted
            record.args = None
        return True


def _configure_mlflow_loggers(root_module_name):
    log_level = (MLFLOW_LOGGING_LEVEL.get() or "INFO").upper()
    # For alembic, use WARNING minimum to reduce noise, but respect higher levels
    alembic_level = log_level if log_level in ("WARNING", "ERROR", "CRITICAL") else "WARNING"

    logging.config.dictConfig({
        "version": 1,
        "disable_existing_loggers": False,
        "formatters": {
            "mlflow_formatter": {
                "()": MlflowFormatter,
                "format": LOGGING_LINE_FORMAT,
                "datefmt": LOGGING_DATETIME_FORMAT,
            },
        },
        "handlers": {
            "mlflow_handler": {
                "formatter": "mlflow_formatter",
                "class": "logging.StreamHandler",
                "stream": MLFLOW_LOGGING_STREAM,
                "filters": ["suppress_in_thread"],
            },
        },
        "loggers": {
            root_module_name: {
                "handlers": ["mlflow_handler"],
                "level": get_mlflow_log_level(),
                "propagate": False,
            },
            "sqlalchemy.engine": {
                "handlers": ["mlflow_handler"],
                "level": "WARN",
                "propagate": False,
            },
            "alembic": {
                "handlers": ["mlflow_handler"],
                "level": alembic_level,
                "propagate": False,
            },
            "huey": {
                "handlers": ["mlflow_handler"],
                "level": alembic_level,
                "propagate": False,
            },
        },
        "filters": {
            "suppress_in_thread": {
                "()": SuppressLogFilter,
            }
        },
    })


def _install_sensitive_query_param_filter() -> None:
    """
    Attach SensitiveQueryParamFilter to urllib3.connectionpool. urllib3 logs
    full target URLs (with presigned-URL credentials in query params) on
    connection retries, which can leak credentials to stderr. Idempotent.
    """
    urllib3_logger = logging.getLogger("urllib3.connectionpool")
    if not any(isinstance(f, SensitiveQueryParamFilter) for f in urllib3_logger.filters):
        urllib3_logger.addFilter(SensitiveQueryParamFilter())


def eprint(*args, **kwargs):
    print(*args, file=MLFLOW_LOGGING_STREAM, **kwargs)


class LoggerMessageFilter(logging.Filter):
    def __init__(self, module: str, filter_regex: re.Pattern):
        super().__init__()
        self._pattern = filter_regex
        self._module = module

    def filter(self, record):
        if record.name == self._module and self._pattern.search(record.msg):
            return False
        return True


@contextlib.contextmanager
def suppress_logs(module: str, filter_regex: re.Pattern):
    """
    Context manager that suppresses log messages from the specified module that match the specified
    regular expression. This is useful for suppressing expected log messages from third-party
    libraries that are not relevant to the current test.
    """
    logger = logging.getLogger(module)
    filter = LoggerMessageFilter(module=module, filter_regex=filter_regex)
    logger.addFilter(filter)
    try:
        yield
    finally:
        logger.removeFilter(filter)


def _debug(s: str) -> None:
    """
    Debug function to test logging level.
    """
    logging.getLogger(__name__).debug(s)


@contextlib.contextmanager
def suppress_logs_in_thread():
    """
    Context manager to suppress logs in the current thread.
    """
    original_value = should_suppress_logs_in_thread.get()
    try:
        should_suppress_logs_in_thread.set(True)
        yield
    finally:
        should_suppress_logs_in_thread.set(original_value)