wehub-resource-sync
82 lines
3.2 KiB
Python
82 lines
3.2 KiB
Python
"""Utility functions for operator naming and logical-op usage recording for Ray Data telemetry.
|
|
"""
|
|
|
|
import json
|
|
import threading
|
|
from typing import Dict
|
|
|
|
from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag
|
|
from ray.data._internal.logical.interfaces import LogicalOperator
|
|
from ray.data._internal.logical.operators import Read, ReadFiles, Write
|
|
|
|
# The dictionary for the operator name and count.
|
|
_recorded_operators = dict()
|
|
_recorded_operators_lock = threading.Lock()
|
|
|
|
|
|
def _is_builtin_cls(cls: type) -> bool:
|
|
"""Return True if ``cls`` is defined under the ``ray`` package.
|
|
|
|
Used to gate which operator / datasource / datasink class names are safe
|
|
to surface in telemetry. Anything outside ``ray.*`` is treated as
|
|
user-defined and anonymized.
|
|
"""
|
|
return (cls.__module__ or "").startswith("ray.")
|
|
|
|
|
|
def record_operators_usage(op: LogicalOperator):
|
|
"""Record logical operator usage with Ray telemetry."""
|
|
ops_dict = dict()
|
|
_collect_operators_to_dict(op, ops_dict)
|
|
ops_json_str = ""
|
|
with _recorded_operators_lock:
|
|
for op_name, count in ops_dict.items():
|
|
_recorded_operators.setdefault(op_name, 0)
|
|
_recorded_operators[op_name] += count
|
|
ops_json_str = json.dumps(_recorded_operators)
|
|
|
|
record_extra_usage_tag(TagKey.DATA_LOGICAL_OPS, ops_json_str)
|
|
|
|
|
|
def anonymize_op_name(op: LogicalOperator) -> str:
|
|
"""Return an op name suitable for usage collection.
|
|
|
|
Read/Write surface their datasource/datasink suffix (``ReadParquet``,
|
|
``WriteIceberg``) when the underlying class ships under ``ray.data.*``;
|
|
user-defined datasources/datasinks collapse to ``ReadCustom`` /
|
|
``WriteCustom``. ``ReadFiles`` (the V2 file-read op) surfaces its
|
|
format via ``datasource_name`` (e.g. ``ReadFilesParquetV2``) when the
|
|
scanner class is built-in; user-defined scanners collapse to
|
|
``ReadFilesCustom``. All other built-in operators emit their class
|
|
name (``Sort``, ``MapBatches``, ``Limit``, …); user-defined
|
|
``LogicalOperator`` subclasses collapse to ``Unknown``.
|
|
"""
|
|
if isinstance(op, Read):
|
|
if _is_builtin_cls(type(op.datasource)):
|
|
return f"Read{op.datasource.get_name()}"
|
|
return "ReadCustom"
|
|
if isinstance(op, Write):
|
|
sink = op.datasink_or_legacy_datasource
|
|
if _is_builtin_cls(type(sink)):
|
|
return f"Write{sink.get_name()}"
|
|
return "WriteCustom"
|
|
if isinstance(op, ReadFiles):
|
|
# Gate on the scanner class — the string ``datasource_name`` field
|
|
# could be set to anything by a user-defined V2 datasource, so it's
|
|
# not safe to surface on its own.
|
|
if _is_builtin_cls(type(op.scanner)):
|
|
return f"ReadFiles{op.datasource_name}"
|
|
return "ReadFilesCustom"
|
|
cls = type(op)
|
|
return cls.__name__ if _is_builtin_cls(cls) else "Unknown"
|
|
|
|
|
|
def _collect_operators_to_dict(op: LogicalOperator, ops_dict: Dict[str, int]):
|
|
"""Collect the logical operator name and count into `ops_dict`."""
|
|
for child in op.input_dependencies:
|
|
_collect_operators_to_dict(child, ops_dict)
|
|
|
|
op_name = anonymize_op_name(op)
|
|
ops_dict.setdefault(op_name, 0)
|
|
ops_dict[op_name] += 1
|