chore: import upstream snapshot with attribution

python/ray/_private/event/event_logger.py

@@ -0,0 +1,202 @@
+import json
+import logging
+import os
+import pathlib
+import random
+import socket
+import string
+import threading
+from datetime import datetime
+from typing import Dict, Optional
+
+from google.protobuf.json_format import Parse
+
+from ray._private.protobuf_compat import message_to_dict
+from ray.core.generated.event_pb2 import Event
+
+global_logger = logging.getLogger(__name__)
+
+
+def get_event_id():
+    return "".join([random.choice(string.hexdigits) for _ in range(36)])
+
+
+class EventLoggerAdapter:
+    def __init__(self, source: Event.SourceType, logger: logging.Logger):
+        """Adapter for the Python logger that's used to emit events.
+
+        When events are emitted, they are aggregated and available via
+        state API and dashboard.
+
+        This class is thread-safe.
+        """
+        self.logger = logger
+        # Aligned with `event.proto`'s `message Event``
+        self.source = source
+        self.source_hostname = socket.gethostname()
+        self.source_pid = os.getpid()
+
+        # The below fields must be protected by this lock.
+        self.lock = threading.Lock()
+        # {str -> str} typed dict
+        self.global_context = {}
+
+    def set_global_context(self, global_context: Dict[str, str] = None):
+        """Set the global metadata.
+
+        This method overwrites the global metadata if it is called more than once.
+        """
+        with self.lock:
+            self.global_context = {} if not global_context else global_context
+
+    def trace(self, message: str, **kwargs):
+        self._emit(Event.Severity.TRACE, message, **kwargs)
+
+    def debug(self, message: str, **kwargs):
+        self._emit(Event.Severity.DEBUG, message, **kwargs)
+
+    def info(self, message: str, **kwargs):
+        self._emit(Event.Severity.INFO, message, **kwargs)
+
+    def warning(self, message: str, **kwargs):
+        self._emit(Event.Severity.WARNING, message, **kwargs)
+
+    def error(self, message: str, **kwargs):
+        self._emit(Event.Severity.ERROR, message, **kwargs)
+
+    def fatal(self, message: str, **kwargs):
+        self._emit(Event.Severity.FATAL, message, **kwargs)
+
+    def _emit(self, severity: Event.Severity, message: str, **kwargs):
+        # NOTE: Python logger is thread-safe,
+        # so we don't need to protect it using locks.
+        event = Event()
+        event.event_id = get_event_id()
+        event.timestamp = int(datetime.now().timestamp())
+        event.message = message
+        event.severity = severity
+        # TODO(sang): Support event type & schema.
+        event.label = ""
+        event.source_type = self.source
+        event.source_hostname = self.source_hostname
+        event.source_pid = self.source_pid
+        custom_fields = event.custom_fields
+        with self.lock:
+            for k, v in self.global_context.items():
+                if v is not None and k is not None:
+                    custom_fields[k] = v
+        for k, v in kwargs.items():
+            if v is not None and k is not None:
+                custom_fields[k] = v
+
+        self.logger.info(
+            json.dumps(
+                message_to_dict(
+                    event,
+                    always_print_fields_with_no_presence=True,
+                    preserving_proto_field_name=True,
+                    use_integers_for_enums=False,
+                )
+            )
+        )
+
+        # Force flush all handlers so that we won't lose events.
+        for handler in self.logger.handlers[:]:
+            try:
+                handler.flush()
+            except Exception:
+                global_logger.exception("Failed to flush event logger handler.")
+
+
+def _build_event_file_logger(source: Event.SourceType, sink_dir: str):
+    logger = logging.getLogger("_ray_event_logger")
+    logger.setLevel(logging.INFO)
+    dir_path = pathlib.Path(sink_dir) / "events"
+    filepath = dir_path / f"event_{source}.log"
+    dir_path.mkdir(exist_ok=True)
+    filepath.touch(exist_ok=True)
+    # Configure the logger.
+    handler = logging.FileHandler(filepath)
+    formatter = logging.Formatter("%(message)s")
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    logger.propagate = False
+    return logger
+
+
+# This lock must be used when accessing or updating global event logger dict.
+_event_logger_lock = threading.Lock()
+_event_logger = {}
+
+
+def get_event_logger(source: Event.SourceType, sink_dir: str):
+    """Get the event logger of the current process.
+
+    There's only 1 event logger per (process, source).
+
+    TODO(sang): Support more impl than file-based logging.
+                Currently, the interface also ties to the
+                file-based logging impl.
+
+    Args:
+        source: The source of the event.
+        sink_dir: The directory to sink event logs.
+
+    Returns:
+        The event logger adapter for the given source.
+    """
+    with _event_logger_lock:
+        global _event_logger
+        source_name = Event.SourceType.Name(source)
+        if source_name not in _event_logger:
+            logger = _build_event_file_logger(source_name, sink_dir)
+            _event_logger[source_name] = EventLoggerAdapter(source, logger)
+
+        return _event_logger[source_name]
+
+
+def parse_event(event_str: str) -> Optional[Event]:
+    """Parse an event from a string.
+
+    Args:
+        event_str: The string to parse. Expect to be a JSON serialized
+            Event protobuf.
+
+    Returns:
+        The parsed event if parsable, else None
+    """
+    try:
+        return Parse(event_str, Event())
+    except Exception:
+        global_logger.exception(f"Failed to parse event: {event_str}")
+        return None
+
+
+def filter_event_by_level(event: Event, filter_event_level: str) -> bool:
+    """Filter an event based on event level.
+
+    Args:
+        event: The event to filter.
+        filter_event_level: The event level string to filter by. Any events
+            that are lower than this level will be filtered.
+
+    Returns:
+        True if the event should be filtered, else False.
+    """
+
+    event_levels = {
+        Event.Severity.TRACE: 0,
+        Event.Severity.DEBUG: 1,
+        Event.Severity.INFO: 2,
+        Event.Severity.WARNING: 3,
+        Event.Severity.ERROR: 4,
+        Event.Severity.FATAL: 5,
+    }
+
+    filter_event_level = filter_event_level.upper()
+    filter_event_level = Event.Severity.Value(filter_event_level)
+
+    if event_levels[event.severity] < event_levels[filter_event_level]:
+        return True
+
+    return False