| """A Ray compatible logging module with actor-scope logger support.""" |
| import contextvars |
| import logging |
| import os |
| import sys |
| from logging.handlers import RotatingFileHandler |
| from typing import Optional |
|
|
| import ray |
|
|
| from trinity.common.constants import ( |
| LOG_DIR_ENV_VAR, |
| LOG_LEVEL_ENV_VAR, |
| LOG_NODE_IP_ENV_VAR, |
| ) |
|
|
| _LOG_FORMAT = "%(levelname)s %(asctime)s [%(filename)s:%(lineno)d] %(message)s" |
| _LOG_DATE_FORMAT = "%m-%d %H:%M:%S" |
|
|
|
|
| class NewLineFormatter(logging.Formatter): |
| """ |
| Formatter that adds logging prefix to newlines to align multi-line messages. |
| """ |
|
|
| def __init__(self, fmt: str, datefmt: Optional[str] = None): |
| super().__init__(fmt, datefmt) |
|
|
| def format(self, record: logging.LogRecord) -> str: |
| msg = super().format(record) |
| if record.message: |
| prefix = msg.split(record.message)[0] |
| msg = msg.replace("\n", f"\r\n{prefix}") |
| return msg |
|
|
|
|
| _ray_logger_ctx: contextvars.ContextVar[Optional[logging.Logger]] = contextvars.ContextVar( |
| "ray_logger", default=None |
| ) |
|
|
|
|
| def get_logger( |
| name: Optional[str] = None, level: Optional[int] = None, in_ray_actor: bool = False |
| ) -> logging.Logger: |
| """ |
| Get a logger instance, compatible with Ray Actor and standard usage. |
| |
| In most cases, DO NOT USE this function in file-level scope, because the logger |
| instance will be created at import time, which will cause issues when used in Ray actors. |
| |
| Args: |
| name (Optional[str]): The name of the logger. If None, uses 'trinity'. |
| level (Optional[int]): The logging level. If None, uses LOG_LEVEL_ENV_VAR or INFO. |
| in_ray_actor (bool): Whether the logger is used within a Ray actor. |
| |
| Returns: |
| logging.Logger: Configured logger instance. |
| """ |
| |
| logger = _ray_logger_ctx.get() |
| if logger is not None: |
| return logger |
|
|
| resolved_level = ( |
| level |
| if level is not None |
| else getattr(logging, os.environ.get(LOG_LEVEL_ENV_VAR, "INFO").upper()) |
| ) |
| logger_name = f"trinity.{name}" if name else "trinity" |
| logger = logging.getLogger(logger_name) |
| logger.propagate = False |
| logger.setLevel(resolved_level) |
| logger.handlers.clear() |
|
|
| |
| stream_handler = logging.StreamHandler(sys.stdout) |
| stream_handler.setLevel(resolved_level) |
| formatter = NewLineFormatter(_LOG_FORMAT, datefmt=_LOG_DATE_FORMAT) |
| stream_handler.setFormatter(formatter) |
| logger.addHandler(stream_handler) |
|
|
| if in_ray_actor: |
| |
| log_dir = os.environ.get(LOG_DIR_ENV_VAR) |
| assert name is not None, "Logger name must be set when logging from a Ray actor" |
| if log_dir: |
| if os.environ.get(LOG_NODE_IP_ENV_VAR, "0") != "0": |
| |
| node_ip = ray.util.get_node_ip_address() |
| log_dir = os.path.join(log_dir, node_ip) |
| os.makedirs(log_dir, exist_ok=True) |
| |
| file_path = os.path.join(log_dir, f"{name}.log") |
| file_handler = RotatingFileHandler( |
| file_path, encoding="utf-8", maxBytes=64 * 1024 * 1024 |
| ) |
| file_handler.setLevel(resolved_level) |
| file_handler.setFormatter(formatter) |
| logger.addHandler(file_handler) |
| _ray_logger_ctx.set(logger) |
| |
|
|
| return logger |
|
|
