wehub-resource-sync
439 lines
15 KiB
Python
439 lines
15 KiB
Python
"""Metadata exporter API for Ray Data datasets."""
|
|
|
|
import logging
|
|
import os
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import asdict, dataclass, field, is_dataclass
|
|
from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Sequence
|
|
|
|
import ray
|
|
from ray._private.event.export_event_logger import (
|
|
EventLogType,
|
|
check_export_api_enabled,
|
|
get_export_event_logger,
|
|
)
|
|
from ray.core.generated.export_dataset_metadata_pb2 import (
|
|
ExportDatasetMetadata as ProtoDatasetMetadata,
|
|
)
|
|
from ray.dashboard.modules.metrics.dashboards.common import Panel
|
|
from ray.dashboard.modules.metrics.dashboards.data_dashboard_panels import (
|
|
OPERATOR_PANELS,
|
|
)
|
|
from ray.data._internal.execution.dataset_state import DatasetState
|
|
from ray.data.context import DataContext
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.data._internal.execution.interfaces.physical_operator import (
|
|
PhysicalOperator,
|
|
)
|
|
from ray.data.context import DataContext
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
UNKNOWN = "unknown"
|
|
|
|
# Number of characters to truncate to when
|
|
# exporting dataset operator arguments
|
|
DEFAULT_TRUNCATION_LENGTH = 100
|
|
|
|
# NOTE: These dataclasses need to be updated in sync with the protobuf definitions in
|
|
# src/ray/protobuf/export_api/export_dataset_metadata.proto
|
|
@dataclass
|
|
class SubStage:
|
|
"""Represents a sub-stage within an operator in the DAG.
|
|
|
|
Attributes:
|
|
name: The name of the sub-stage.
|
|
id: The unique identifier of the sub-stage.
|
|
"""
|
|
|
|
name: str
|
|
id: str
|
|
|
|
|
|
@dataclass
|
|
class Operator:
|
|
"""Represents a data processing operator in the DAG.
|
|
|
|
Attributes:
|
|
name: The name of the operator.
|
|
id: The unique identifier of the operator within the DAG structure, typically
|
|
incorporating a position or index (e.g., "ReadParquet_0"). This is used for
|
|
referencing operators within the DAG topology.
|
|
uuid: The system-generated UUID of the physical operator instance. This is the
|
|
internal unique identifier created when the operator instance is initialized
|
|
and remains consistent throughout its lifetime.
|
|
input_dependencies: List of operator IDs that this operator depends on for input.
|
|
sub_stages: List of sub-stages contained within this operator.
|
|
args: User-specified arguments associated with the operator, which may
|
|
include configuration settings, options, or other relevant data for the operator.
|
|
execution_start_time: The timestamp when the operator execution begins.
|
|
execution_end_time: The timestamp when the operator execution ends.
|
|
state: The state of the operator.
|
|
"""
|
|
|
|
name: str
|
|
id: str
|
|
uuid: str
|
|
execution_start_time: Optional[float]
|
|
execution_end_time: Optional[float]
|
|
state: str
|
|
input_dependencies: List[str] = field(default_factory=list)
|
|
sub_stages: List[SubStage] = field(default_factory=list)
|
|
args: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class DataContextMetadata:
|
|
"""Represents sanitized DataContext metadata for export.
|
|
|
|
This class wraps the sanitized dictionary representation of DataContext to provide
|
|
type safety and consistency with other metadata structures. The actual data is stored
|
|
as a dictionary that has been processed by sanitize_for_struct() to ensure it can be
|
|
safely serialized without module dependency issues.
|
|
|
|
Attributes:
|
|
config: Dictionary containing sanitized DataContext configuration values.
|
|
All complex objects have been converted to basic types (strings, numbers,
|
|
lists, dicts) suitable for protobuf Struct format.
|
|
"""
|
|
|
|
config: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
@staticmethod
|
|
def from_data_context(data_context: "DataContext") -> "DataContextMetadata":
|
|
"""Create DataContextMetadata from a DataContext object.
|
|
|
|
Args:
|
|
data_context: The DataContext object to convert.
|
|
|
|
Returns:
|
|
A DataContextMetadata instance with sanitized configuration.
|
|
"""
|
|
return DataContextMetadata(config=sanitize_for_struct(data_context))
|
|
|
|
|
|
@dataclass
|
|
class Topology:
|
|
"""Represents the complete structure of the operator DAG.
|
|
|
|
Attributes:
|
|
operators: List of all operators in the DAG.
|
|
"""
|
|
|
|
operators: List[Operator] = field(default_factory=list)
|
|
|
|
@staticmethod
|
|
def create_topology_metadata(
|
|
dag: "PhysicalOperator", op_to_id: Dict["PhysicalOperator", str]
|
|
) -> "Topology":
|
|
"""Create a Topology structure from the physical operator DAG.
|
|
|
|
Args:
|
|
dag: The operator DAG to analyze.
|
|
op_to_id: Mapping from each physical operator to its stable string ID.
|
|
|
|
Returns:
|
|
A Topology object representing the operator DAG structure.
|
|
"""
|
|
# Create the result structure
|
|
result = Topology()
|
|
|
|
# Add detailed operator information with dependencies
|
|
for op in dag.post_order_iter():
|
|
op_id = op_to_id[op]
|
|
|
|
# Create operator object
|
|
operator = Operator(
|
|
name=op.name,
|
|
id=op_id,
|
|
uuid=op.id,
|
|
input_dependencies=[
|
|
op_to_id[dep] for dep in op.input_dependencies if dep in op_to_id
|
|
],
|
|
args=sanitize_for_struct(op._get_logical_args()),
|
|
execution_start_time=None,
|
|
execution_end_time=None,
|
|
state=DatasetState.PENDING.name,
|
|
)
|
|
|
|
# Add sub-stages if they exist
|
|
if hasattr(op, "_sub_progress_bar_names") and op._sub_progress_bar_names:
|
|
for j, sub_name in enumerate(op._sub_progress_bar_names):
|
|
sub_stage_id = f"{op_id}_sub_{j}"
|
|
operator.sub_stages.append(SubStage(name=sub_name, id=sub_stage_id))
|
|
|
|
result.operators.append(operator)
|
|
return result
|
|
|
|
|
|
@dataclass
|
|
class DatasetMetadata:
|
|
"""Metadata about a Ray Data dataset.
|
|
|
|
This class represents the metadata associated with a dataset, including its provenance
|
|
information and execution details.
|
|
|
|
Attributes:
|
|
job_id: The ID of the job running this dataset.
|
|
topology: The structure of the dataset's operator DAG.
|
|
dataset_id: The unique ID of the dataset.
|
|
start_time: The timestamp when the dataset is registered.
|
|
data_context: DataContextMetadata containing sanitized DataContext configuration.
|
|
This is pre-processed using sanitize_for_struct() to avoid serialization
|
|
issues with module dependencies.
|
|
execution_start_time: The timestamp when the dataset execution starts.
|
|
execution_end_time: The timestamp when the dataset execution ends.
|
|
state: The state of the dataset.
|
|
"""
|
|
|
|
job_id: str
|
|
topology: Topology
|
|
dataset_id: str
|
|
start_time: float
|
|
data_context: DataContextMetadata
|
|
execution_start_time: Optional[float]
|
|
execution_end_time: Optional[float]
|
|
state: str
|
|
|
|
|
|
def _add_ellipsis_for_string(s: str, truncate_length: int) -> str:
|
|
if len(s) > truncate_length:
|
|
return s[:truncate_length] + "..."
|
|
return s
|
|
|
|
|
|
def sanitize_for_struct(obj, truncate_length=DEFAULT_TRUNCATION_LENGTH):
|
|
"""Prepares the obj for Struct Protobuf format by recursively
|
|
going through dictionaries, lists, etc...
|
|
|
|
- Dataclasses will be converted to dicts
|
|
- Dictionary keys will be converted to strings
|
|
- Lists, tuples, sets, bytes, bytearrays will be converted to lists
|
|
"""
|
|
if isinstance(obj, Mapping):
|
|
# protobuf Struct key names must be strings.
|
|
return {str(k): sanitize_for_struct(v, truncate_length) for k, v in obj.items()}
|
|
elif isinstance(obj, str):
|
|
return _add_ellipsis_for_string(obj, truncate_length)
|
|
elif isinstance(obj, (Sequence, set)):
|
|
# Convert all sequence-like types (lists, tuples, sets, bytes, other sequences) to lists
|
|
res = []
|
|
for i, v in enumerate(obj):
|
|
if i >= truncate_length:
|
|
res.append("...")
|
|
break
|
|
res.append(sanitize_for_struct(v, truncate_length))
|
|
return res
|
|
else:
|
|
try:
|
|
if is_dataclass(obj):
|
|
return sanitize_for_struct(asdict(obj), truncate_length)
|
|
return _add_ellipsis_for_string(str(obj), truncate_length)
|
|
except Exception:
|
|
unk_name = f"{UNKNOWN}: {type(obj).__name__}"
|
|
return _add_ellipsis_for_string(unk_name, truncate_length)
|
|
|
|
|
|
def dataset_metadata_to_proto(
|
|
dataset_metadata: DatasetMetadata,
|
|
include_data_context: bool = True,
|
|
include_op_args: bool = True,
|
|
) -> Any:
|
|
"""Convert the dataset metadata to a protobuf message.
|
|
|
|
Args:
|
|
dataset_metadata: DatasetMetadata object containing the dataset's
|
|
information and DAG structure.
|
|
include_data_context: If DataContext will be exported
|
|
include_op_args: If operator args will be exported
|
|
|
|
Returns:
|
|
The protobuf message representing the dataset metadata.
|
|
"""
|
|
|
|
from google.protobuf.struct_pb2 import Struct
|
|
|
|
from ray.core.generated.export_dataset_metadata_pb2 import (
|
|
ExportDatasetMetadata as ProtoDatasetMetadata,
|
|
Operator as ProtoOperator,
|
|
SubStage as ProtoSubStage,
|
|
Topology as ProtoTopology,
|
|
)
|
|
|
|
# Create the protobuf message
|
|
proto_dataset_metadata = ProtoDatasetMetadata()
|
|
proto_topology = ProtoTopology()
|
|
|
|
# Add operators to the DAG
|
|
for op in dataset_metadata.topology.operators:
|
|
args = Struct()
|
|
if include_op_args:
|
|
args.update(op.args)
|
|
proto_operator = ProtoOperator(
|
|
name=op.name,
|
|
id=op.id,
|
|
uuid=op.uuid,
|
|
args=args,
|
|
execution_start_time=op.execution_start_time,
|
|
execution_end_time=op.execution_end_time,
|
|
state=ProtoOperator.OperatorState.Value(op.state),
|
|
)
|
|
|
|
# Add input dependencies
|
|
for dep_id in op.input_dependencies:
|
|
proto_operator.input_dependencies.append(dep_id)
|
|
|
|
# Add sub-stages
|
|
for sub_stage in op.sub_stages:
|
|
proto_sub_stage = ProtoSubStage(
|
|
name=sub_stage.name,
|
|
id=sub_stage.id,
|
|
)
|
|
proto_operator.sub_stages.append(proto_sub_stage)
|
|
|
|
# Add the operator to the DAG
|
|
proto_topology.operators.append(proto_operator)
|
|
|
|
# Populate the data metadata proto
|
|
data_context = Struct()
|
|
if include_data_context:
|
|
data_context.update(dataset_metadata.data_context.config)
|
|
proto_dataset_metadata = ProtoDatasetMetadata(
|
|
dataset_id=dataset_metadata.dataset_id,
|
|
job_id=dataset_metadata.job_id,
|
|
start_time=dataset_metadata.start_time,
|
|
data_context=data_context,
|
|
execution_start_time=dataset_metadata.execution_start_time,
|
|
execution_end_time=dataset_metadata.execution_end_time,
|
|
state=ProtoDatasetMetadata.DatasetState.Value(dataset_metadata.state),
|
|
operator_panels=[_to_proto_dashboard_panel(p) for p in OPERATOR_PANELS],
|
|
)
|
|
proto_dataset_metadata.topology.CopyFrom(proto_topology)
|
|
|
|
return proto_dataset_metadata
|
|
|
|
|
|
def _to_proto_dashboard_panel(
|
|
panel: Panel,
|
|
) -> ProtoDatasetMetadata.DashboardPanelMetadata:
|
|
"""Convert Dashboard Panel to protobuf format."""
|
|
proto_panel = ProtoDatasetMetadata.DashboardPanelMetadata(
|
|
id=str(panel.id),
|
|
title=panel.title,
|
|
)
|
|
|
|
return proto_panel
|
|
|
|
|
|
def get_dataset_metadata_exporter() -> "DatasetMetadataExporter":
|
|
"""Get the dataset metadata exporter instance.
|
|
|
|
Returns:
|
|
The dataset metadata exporter instance.
|
|
"""
|
|
return LoggerDatasetMetadataExporter.create_if_enabled()
|
|
|
|
|
|
class DatasetMetadataExporter(ABC):
|
|
"""Abstract base class for dataset metadata exporters.
|
|
|
|
Implementations of this interface can export Ray Data metadata to various destinations
|
|
like log files, databases, or monitoring systems.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def export_dataset_metadata(
|
|
self,
|
|
dataset_metadata: DatasetMetadata,
|
|
include_data_context: bool = True,
|
|
include_op_args: bool = True,
|
|
) -> None:
|
|
"""Export dataset metadata to the destination.
|
|
|
|
Args:
|
|
dataset_metadata: DatasetMetadata object containing dataset information.
|
|
include_data_context: If DataContext will be exported
|
|
include_op_args: If operator args will be exported
|
|
"""
|
|
pass
|
|
|
|
@classmethod
|
|
@abstractmethod
|
|
def create_if_enabled(cls) -> Optional["DatasetMetadataExporter"]:
|
|
"""Create an exporter instance if the export functionality is enabled.
|
|
|
|
Returns:
|
|
An exporter instance if enabled, None otherwise.
|
|
"""
|
|
pass
|
|
|
|
|
|
class LoggerDatasetMetadataExporter(DatasetMetadataExporter):
|
|
"""Dataset metadata exporter implementation that uses the Ray export event logger.
|
|
|
|
This exporter writes dataset metadata to log files using Ray's export event system.
|
|
"""
|
|
|
|
def __init__(self, logger: logging.Logger):
|
|
"""Initialize with a configured export event logger.
|
|
|
|
Args:
|
|
logger: The export event logger to use for writing events.
|
|
"""
|
|
self._export_logger = logger
|
|
|
|
def export_dataset_metadata(
|
|
self,
|
|
dataset_metadata: DatasetMetadata,
|
|
include_data_context: bool = True,
|
|
include_op_args: bool = True,
|
|
) -> None:
|
|
"""Export dataset metadata using the export event logger.
|
|
|
|
Args:
|
|
dataset_metadata: DatasetMetadata object containing dataset information.
|
|
include_data_context: If DataContext will be exported
|
|
include_op_args: If operator args will be exported
|
|
"""
|
|
data_metadata_proto = dataset_metadata_to_proto(
|
|
dataset_metadata,
|
|
include_data_context,
|
|
include_op_args,
|
|
)
|
|
self._export_logger.send_event(data_metadata_proto)
|
|
|
|
@classmethod
|
|
def create_if_enabled(cls) -> Optional["LoggerDatasetMetadataExporter"]:
|
|
"""Create a logger-based exporter if the export API is enabled.
|
|
|
|
Returns:
|
|
A LoggerDatasetMetadataExporter instance if enabled, None otherwise.
|
|
"""
|
|
# Proto schemas should be imported here to prevent serialization errors
|
|
from ray.core.generated.export_event_pb2 import ExportEvent
|
|
|
|
is_dataset_metadata_export_api_enabled = check_export_api_enabled(
|
|
ExportEvent.SourceType.EXPORT_DATASET_METADATA
|
|
)
|
|
if not is_dataset_metadata_export_api_enabled:
|
|
# The export API is not enabled, so we shouldn't create an exporter
|
|
return None
|
|
|
|
log_directory = os.path.join(
|
|
ray._private.worker._global_node.get_session_dir_path(), "logs"
|
|
)
|
|
|
|
try:
|
|
logger = get_export_event_logger(
|
|
EventLogType.DATASET_METADATA,
|
|
log_directory,
|
|
)
|
|
return LoggerDatasetMetadataExporter(logger)
|
|
except Exception:
|
|
logger.exception(
|
|
"Unable to initialize the export event logger, so no Dataset Metadata export "
|
|
"events will be written."
|
|
)
|
|
return None
|