chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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
|
||||
@@ -0,0 +1,261 @@
|
||||
import json
|
||||
import logging
|
||||
import pathlib
|
||||
import random
|
||||
import string
|
||||
import threading
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import Union
|
||||
|
||||
from ray._private import ray_constants
|
||||
from ray._private.protobuf_compat import message_to_dict
|
||||
from ray.core.generated.export_dataset_metadata_pb2 import (
|
||||
ExportDatasetMetadata,
|
||||
)
|
||||
from ray.core.generated.export_dataset_operator_event_pb2 import (
|
||||
ExportDatasetOperatorEventData,
|
||||
)
|
||||
from ray.core.generated.export_dataset_operator_schema_pb2 import (
|
||||
ExportDatasetOperatorSchema,
|
||||
)
|
||||
from ray.core.generated.export_event_pb2 import ExportEvent
|
||||
from ray.core.generated.export_submission_job_event_pb2 import (
|
||||
ExportSubmissionJobEventData,
|
||||
)
|
||||
from ray.core.generated.export_train_state_pb2 import (
|
||||
ExportTrainRunAttemptEventData,
|
||||
ExportTrainRunEventData,
|
||||
)
|
||||
|
||||
global_logger = logging.getLogger(__name__)
|
||||
|
||||
# This contains the union of export event data types which emit events
|
||||
# using the python ExportEventLoggerAdapter
|
||||
ExportEventDataType = Union[
|
||||
ExportSubmissionJobEventData,
|
||||
ExportTrainRunEventData,
|
||||
ExportTrainRunAttemptEventData,
|
||||
ExportDatasetMetadata,
|
||||
ExportDatasetOperatorEventData,
|
||||
ExportDatasetOperatorSchema,
|
||||
]
|
||||
|
||||
|
||||
class EventLogType(Enum):
|
||||
"""Enum class representing different types of export event logs.
|
||||
|
||||
Each enum value contains a log type name and a set of supported event data types.
|
||||
|
||||
Attributes:
|
||||
TRAIN_STATE: Export events related to training state, supporting train run and attempt events.
|
||||
SUBMISSION_JOB: Export events related to job submissions.
|
||||
DATASET_METADATA: Export events related to dataset metadata.
|
||||
DATASET_OPERATOR_EVENT: Export events related to Ray Data operator.
|
||||
DATASET_OPERATOR_SCHEMA: Export schema related to Ray Data operator.
|
||||
"""
|
||||
|
||||
TRAIN_STATE = (
|
||||
"EXPORT_TRAIN_STATE",
|
||||
{ExportTrainRunEventData, ExportTrainRunAttemptEventData},
|
||||
)
|
||||
SUBMISSION_JOB = ("EXPORT_SUBMISSION_JOB", {ExportSubmissionJobEventData})
|
||||
DATASET_METADATA = ("EXPORT_DATASET_METADATA", {ExportDatasetMetadata})
|
||||
DATASET_OPERATOR_EVENT = (
|
||||
"EXPORT_DATASET_OPERATOR_EVENT",
|
||||
{ExportDatasetOperatorEventData},
|
||||
)
|
||||
DATASET_OPERATOR_SCHEMA = (
|
||||
"EXPORT_DATASET_OPERATOR_SCHEMA",
|
||||
{ExportDatasetOperatorSchema},
|
||||
)
|
||||
|
||||
def __init__(self, log_type_name: str, event_types: set[ExportEventDataType]):
|
||||
"""Initialize an EventLogType enum value.
|
||||
|
||||
Args:
|
||||
log_type_name: String identifier for the log type. This name is used to construct the log file name.
|
||||
See `_build_export_event_file_logger` for more details.
|
||||
event_types: Set of event data types that this log type supports.
|
||||
"""
|
||||
self.log_type_name = log_type_name
|
||||
self.event_types = event_types
|
||||
|
||||
def supports_event_type(self, event_type: ExportEventDataType) -> bool:
|
||||
"""Check if this log type supports the given event data type.
|
||||
|
||||
Args:
|
||||
event_type: The event data type to check for support.
|
||||
|
||||
Returns:
|
||||
bool: True if the event type is supported, False otherwise.
|
||||
"""
|
||||
return type(event_type) in self.event_types
|
||||
|
||||
|
||||
def generate_event_id():
|
||||
return "".join([random.choice(string.hexdigits) for _ in range(18)])
|
||||
|
||||
|
||||
class ExportEventLoggerAdapter:
|
||||
def __init__(self, log_type: EventLogType, logger: logging.Logger):
|
||||
"""Adapter for the Python logger that's used to emit export events."""
|
||||
self.logger = logger
|
||||
self.log_type = log_type
|
||||
|
||||
def send_event(self, event_data: ExportEventDataType):
|
||||
# NOTE: Python logger is thread-safe,
|
||||
# so we don't need to protect it using locks.
|
||||
try:
|
||||
event = self._create_export_event(event_data)
|
||||
except TypeError:
|
||||
global_logger.exception(
|
||||
"Failed to create ExportEvent from event_data so no "
|
||||
"event will be written to file."
|
||||
)
|
||||
return
|
||||
|
||||
event_as_str = self._export_event_to_string(event)
|
||||
|
||||
self.logger.info(event_as_str)
|
||||
# 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 export event logger handler.")
|
||||
|
||||
def _create_export_event(self, event_data: ExportEventDataType) -> ExportEvent:
|
||||
event = ExportEvent()
|
||||
event.event_id = generate_event_id()
|
||||
event.timestamp = int(datetime.now().timestamp())
|
||||
if isinstance(event_data, ExportSubmissionJobEventData):
|
||||
event.submission_job_event_data.CopyFrom(event_data)
|
||||
event.source_type = ExportEvent.SourceType.EXPORT_SUBMISSION_JOB
|
||||
elif isinstance(event_data, ExportTrainRunEventData):
|
||||
event.train_run_event_data.CopyFrom(event_data)
|
||||
event.source_type = ExportEvent.SourceType.EXPORT_TRAIN_RUN
|
||||
elif isinstance(event_data, ExportTrainRunAttemptEventData):
|
||||
event.train_run_attempt_event_data.CopyFrom(event_data)
|
||||
event.source_type = ExportEvent.SourceType.EXPORT_TRAIN_RUN_ATTEMPT
|
||||
elif isinstance(event_data, ExportDatasetMetadata):
|
||||
event.dataset_metadata.CopyFrom(event_data)
|
||||
event.source_type = ExportEvent.SourceType.EXPORT_DATASET_METADATA
|
||||
elif isinstance(event_data, ExportDatasetOperatorEventData):
|
||||
event.dataset_operator_event_data.CopyFrom(event_data)
|
||||
event.source_type = ExportEvent.SourceType.EXPORT_DATASET_OPERATOR_EVENT
|
||||
elif isinstance(event_data, ExportDatasetOperatorSchema):
|
||||
event.dataset_operator_schema.CopyFrom(event_data)
|
||||
event.source_type = ExportEvent.SourceType.EXPORT_DATASET_OPERATOR_SCHEMA
|
||||
else:
|
||||
raise TypeError(f"Invalid event_data type: {type(event_data)}")
|
||||
if not self.log_type.supports_event_type(event_data):
|
||||
global_logger.error(
|
||||
f"event_data has source type {event.source_type}, however "
|
||||
f"the event was sent to a logger with log type {self.log_type.log_type_name}. "
|
||||
f"The event will still be written to the file of {self.log_type.log_type_name} "
|
||||
"but this indicates a bug in the code."
|
||||
)
|
||||
pass
|
||||
return event
|
||||
|
||||
def _export_event_to_string(self, event: ExportEvent) -> str:
|
||||
event_data_json = {}
|
||||
proto_to_dict_options = {
|
||||
"always_print_fields_with_no_presence": True,
|
||||
"preserving_proto_field_name": True,
|
||||
"use_integers_for_enums": False,
|
||||
}
|
||||
event_data_field_set = event.WhichOneof("event_data")
|
||||
if event_data_field_set:
|
||||
event_data_json = message_to_dict(
|
||||
getattr(event, event_data_field_set),
|
||||
**proto_to_dict_options,
|
||||
)
|
||||
else:
|
||||
global_logger.error(
|
||||
f"event_data missing from export event with id {event.event_id} "
|
||||
f"and type {event.source_type}. An empty event will be written, "
|
||||
"but this indicates a bug in the code. "
|
||||
)
|
||||
pass
|
||||
event_json = {
|
||||
"event_id": event.event_id,
|
||||
"timestamp": event.timestamp,
|
||||
"source_type": ExportEvent.SourceType.Name(event.source_type),
|
||||
"event_data": event_data_json,
|
||||
}
|
||||
return json.dumps(event_json)
|
||||
|
||||
|
||||
def _build_export_event_file_logger(
|
||||
log_type_name: str, sink_dir: str
|
||||
) -> logging.Logger:
|
||||
logger = logging.getLogger("_ray_export_event_logger_" + log_type_name)
|
||||
logger.setLevel(logging.INFO)
|
||||
dir_path = pathlib.Path(sink_dir) / "export_events"
|
||||
filepath = dir_path / f"event_{log_type_name}.log"
|
||||
dir_path.mkdir(exist_ok=True)
|
||||
filepath.touch(exist_ok=True)
|
||||
# Configure the logger.
|
||||
# Default is 100 MB max file size
|
||||
handler = logging.handlers.RotatingFileHandler(
|
||||
filepath,
|
||||
maxBytes=(ray_constants.RAY_EXPORT_EVENT_MAX_FILE_SIZE_BYTES),
|
||||
backupCount=ray_constants.RAY_EXPORT_EVENT_MAX_BACKUP_COUNT,
|
||||
)
|
||||
logger.addHandler(handler)
|
||||
logger.propagate = False
|
||||
return logger
|
||||
|
||||
|
||||
# This lock must be used when accessing or updating global event logger dict.
|
||||
_export_event_logger_lock = threading.Lock()
|
||||
_export_event_logger = {}
|
||||
|
||||
|
||||
def get_export_event_logger(log_type: EventLogType, sink_dir: str) -> logging.Logger:
|
||||
"""Get the export event logger of the current process.
|
||||
|
||||
There's only one logger per export event source.
|
||||
|
||||
Args:
|
||||
log_type: The type of the export event.
|
||||
sink_dir: The directory to sink event logs.
|
||||
|
||||
Returns:
|
||||
The export event logger adapter for the given log type.
|
||||
"""
|
||||
with _export_event_logger_lock:
|
||||
global _export_event_logger
|
||||
log_type_name = log_type.log_type_name
|
||||
if log_type_name not in _export_event_logger:
|
||||
logger = _build_export_event_file_logger(log_type.log_type_name, sink_dir)
|
||||
_export_event_logger[log_type_name] = ExportEventLoggerAdapter(
|
||||
log_type, logger
|
||||
)
|
||||
|
||||
return _export_event_logger[log_type_name]
|
||||
|
||||
|
||||
def check_export_api_enabled(
|
||||
source: ExportEvent.SourceType,
|
||||
) -> bool:
|
||||
"""
|
||||
Check RAY_ENABLE_EXPORT_API_WRITE and RAY_ENABLE_EXPORT_API_WRITE_CONFIG environment
|
||||
variables to verify if export events should be written for the given source type.
|
||||
|
||||
Args:
|
||||
source: The source of the export event.
|
||||
|
||||
Returns:
|
||||
True if the export API is enabled for the given source, else False.
|
||||
"""
|
||||
if ray_constants.RAY_ENABLE_EXPORT_API_WRITE:
|
||||
return True
|
||||
source_name = ExportEvent.SourceType.Name(source)
|
||||
return (
|
||||
source_name in ray_constants.RAY_ENABLE_EXPORT_API_WRITE_CONFIG
|
||||
if ray_constants.RAY_ENABLE_EXPORT_API_WRITE_CONFIG
|
||||
else False
|
||||
)
|
||||
Reference in New Issue
Block a user