chore: import upstream snapshot with attribution

python/ray/_common/formatters.py

@@ -0,0 +1,120 @@
+import json
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List
+
+from ray._common.logging_constants import (
+    LOGGER_FLATTEN_KEYS,
+    LOGRECORD_STANDARD_ATTRS,
+    LogKey,
+)
+from ray._private.log import INTERNAL_TIMESTAMP_LOG_KEY
+from ray._private.ray_constants import LOGGER_FORMAT
+
+
+def _append_flatten_attributes(formatted_attrs: Dict[str, Any], key: str, value: Any):
+    """Flatten the dictionary values for special keys and append the values in place.
+
+    If the key is in `LOGGER_FLATTEN_KEYS`, the value will be flattened and appended
+    to the `formatted_attrs` dictionary. Otherwise, the key-value pair will be appended
+    directly.
+    """
+    if key in LOGGER_FLATTEN_KEYS:
+        if not isinstance(value, dict):
+            raise ValueError(
+                f"Expected a dictionary passing into {key}, but got {type(value)}"
+            )
+        for k, v in value.items():
+            if k in formatted_attrs:
+                raise KeyError(f"Found duplicated key in the log record: {k}")
+            formatted_attrs[k] = v
+    else:
+        formatted_attrs[key] = value
+
+
+class AbstractFormatter(logging.Formatter, ABC):
+    def __init__(self, fmt=None, datefmt=None, style="%", validate=True) -> None:
+        super().__init__(fmt, datefmt, style, validate)
+        self._additional_log_standard_attrs = []
+
+    def set_additional_log_standard_attrs(
+        self, additional_log_standard_attrs: List[str]
+    ) -> None:
+        self._additional_log_standard_attrs = additional_log_standard_attrs
+
+    @property
+    def additional_log_standard_attrs(self) -> List[str]:
+        return self._additional_log_standard_attrs
+
+    def generate_record_format_attrs(
+        self,
+        record: logging.LogRecord,
+        exclude_default_standard_attrs,
+    ) -> dict:
+        record_format_attrs = {}
+
+        # If `exclude_default_standard_attrs` is False, include the standard attributes.
+        # Otherwise, include only Ray and user-provided context.
+        if not exclude_default_standard_attrs:
+            record_format_attrs.update(
+                {
+                    LogKey.ASCTIME.value: self.formatTime(record),
+                    LogKey.LEVELNAME.value: record.levelname,
+                    LogKey.MESSAGE.value: record.getMessage(),
+                    LogKey.FILENAME.value: record.filename,
+                    LogKey.LINENO.value: record.lineno,
+                    LogKey.PROCESS.value: record.process,
+                }
+            )
+            if record.exc_info:
+                if not record.exc_text:
+                    record.exc_text = self.formatException(record.exc_info)
+                record_format_attrs[LogKey.EXC_TEXT.value] = record.exc_text
+
+        # Add the user specified additional standard attributes.
+        for key in self._additional_log_standard_attrs:
+            _append_flatten_attributes(
+                record_format_attrs, key, getattr(record, key, None)
+            )
+
+        for key, value in record.__dict__.items():
+            # Both Ray and user-provided context are stored in `record_format`.
+            if key not in LOGRECORD_STANDARD_ATTRS:
+                _append_flatten_attributes(record_format_attrs, key, value)
+
+        # Format the internal timestamp to the standardized `timestamp_ns` key.
+        if INTERNAL_TIMESTAMP_LOG_KEY in record_format_attrs:
+            record_format_attrs[LogKey.TIMESTAMP_NS.value] = record_format_attrs.pop(
+                INTERNAL_TIMESTAMP_LOG_KEY
+            )
+
+        return record_format_attrs
+
+    @abstractmethod
+    def format(self, record: logging.LogRecord) -> str:
+        pass
+
+
+class JSONFormatter(AbstractFormatter):
+    def format(self, record: logging.LogRecord) -> str:
+        record_format_attrs = self.generate_record_format_attrs(
+            record, exclude_default_standard_attrs=False
+        )
+        return json.dumps(record_format_attrs)
+
+
+class TextFormatter(AbstractFormatter):
+    def __init__(self, fmt=None, datefmt=None, style="%", validate=True) -> None:
+        super().__init__(fmt, datefmt, style, validate)
+        self._inner_formatter = logging.Formatter(LOGGER_FORMAT)
+
+    def format(self, record: logging.LogRecord) -> str:
+        s = self._inner_formatter.format(record)
+        record_format_attrs = self.generate_record_format_attrs(
+            record, exclude_default_standard_attrs=True
+        )
+
+        additional_attrs = " ".join(
+            [f"{key}={value}" for key, value in record_format_attrs.items()]
+        )
+        return f"{s} {additional_attrs}"