wehub-resource-sync
203 lines
6.3 KiB
Python
203 lines
6.3 KiB
Python
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
|