chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,230 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
import ray
|
||||
import ray.dashboard.utils as dashboard_utils
|
||||
from ray._private import ray_constants
|
||||
from ray._private.telemetry.open_telemetry_metric_recorder import (
|
||||
OpenTelemetryMetricRecorder,
|
||||
)
|
||||
from ray.core.generated import (
|
||||
events_event_aggregator_service_pb2,
|
||||
events_event_aggregator_service_pb2_grpc,
|
||||
)
|
||||
from ray.dashboard.modules.aggregator.constants import AGGREGATOR_AGENT_METRIC_PREFIX
|
||||
from ray.dashboard.modules.aggregator.multi_consumer_event_buffer import (
|
||||
MultiConsumerEventBuffer,
|
||||
)
|
||||
from ray.dashboard.modules.aggregator.publisher.async_publisher_client import (
|
||||
AsyncGCSTaskEventsPublisherClient,
|
||||
AsyncHttpPublisherClient,
|
||||
)
|
||||
from ray.dashboard.modules.aggregator.publisher.ray_event_publisher import (
|
||||
NoopPublisher,
|
||||
RayEventPublisher,
|
||||
)
|
||||
from ray.dashboard.modules.aggregator.task_events_metadata_buffer import (
|
||||
TaskEventsMetadataBuffer,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Max number of threads for the thread pool executor handling CPU intensive tasks
|
||||
THREAD_POOL_EXECUTOR_MAX_WORKERS = ray_constants.env_integer(
|
||||
"RAY_DASHBOARD_AGGREGATOR_AGENT_THREAD_POOL_EXECUTOR_MAX_WORKERS", 1
|
||||
)
|
||||
# Interval to check the main thread liveness
|
||||
CHECK_MAIN_THREAD_LIVENESS_INTERVAL_SECONDS = ray_constants.env_float(
|
||||
"RAY_DASHBOARD_AGGREGATOR_AGENT_CHECK_MAIN_THREAD_LIVENESS_INTERVAL_SECONDS", 0.1
|
||||
)
|
||||
# Maximum size of the event buffer in the aggregator agent
|
||||
# The default value was 1,000,000 but was reduced to 100,000 now to avoid being OOM Killed.
|
||||
# We observed that the previous 1,000,000 could take up to 20 GB of memory.
|
||||
# TODO (rueian): Find a better way for the event buffer to store events while avoiding being OOM Killed. For example:
|
||||
# 1. Store bytes instead of python objects and count the size in bytes.
|
||||
# 2. Compress the bytes before storing them in the buffer? (This will increase the CPU usage)
|
||||
# 3. Don't be fixed at 10,0000 but adjust the buffer size based on the available memory on startup.
|
||||
MAX_EVENT_BUFFER_SIZE = ray_constants.env_integer(
|
||||
"RAY_DASHBOARD_AGGREGATOR_AGENT_MAX_EVENT_BUFFER_SIZE", 100000
|
||||
)
|
||||
# Maximum number of events to send in a single batch to the destination
|
||||
MAX_EVENT_SEND_BATCH_SIZE = ray_constants.env_integer(
|
||||
"RAY_DASHBOARD_AGGREGATOR_AGENT_MAX_EVENT_SEND_BATCH_SIZE", 1000
|
||||
)
|
||||
# Address of the external service to send events with format of "http://<ip>:<port>"
|
||||
EVENTS_EXPORT_ADDR = os.environ.get(
|
||||
"RAY_DASHBOARD_AGGREGATOR_AGENT_EVENTS_EXPORT_ADDR", ""
|
||||
)
|
||||
# flag to enable publishing events to the external HTTP service
|
||||
PUBLISH_EVENTS_TO_EXTERNAL_HTTP_SERVICE = ray_constants.env_bool(
|
||||
"RAY_DASHBOARD_AGGREGATOR_AGENT_PUBLISH_EVENTS_TO_EXTERNAL_HTTP_SERVICE", True
|
||||
)
|
||||
# flag to enable publishing events to GCS
|
||||
PUBLISH_EVENTS_TO_GCS = ray_constants.env_bool(
|
||||
"RAY_DASHBOARD_AGGREGATOR_AGENT_PUBLISH_EVENTS_TO_GCS", False
|
||||
)
|
||||
# flag to control whether preserve the proto field name when converting the events to
|
||||
# JSON. If True, the proto field name will be preserved. If False, the proto field name
|
||||
# will be converted to camel case.
|
||||
PRESERVE_PROTO_FIELD_NAME = ray_constants.env_bool(
|
||||
"RAY_DASHBOARD_AGGREGATOR_AGENT_PRESERVE_PROTO_FIELD_NAME", False
|
||||
)
|
||||
|
||||
|
||||
class AggregatorAgent(
|
||||
dashboard_utils.DashboardAgentModule,
|
||||
events_event_aggregator_service_pb2_grpc.EventAggregatorServiceServicer,
|
||||
):
|
||||
"""
|
||||
AggregatorAgent is a dashboard agent module that collects events sent with
|
||||
gRPC from other components, buffers them, and periodically sends them to GCS and
|
||||
an external service with HTTP POST requests for further processing or storage
|
||||
"""
|
||||
|
||||
def __init__(self, dashboard_agent) -> None:
|
||||
super().__init__(dashboard_agent)
|
||||
self._ip = dashboard_agent.ip
|
||||
self._pid = os.getpid()
|
||||
|
||||
# common prometheus labels for aggregator-owned metrics
|
||||
self._common_tags = {
|
||||
"ip": self._ip,
|
||||
"pid": str(self._pid),
|
||||
"Version": ray.__version__,
|
||||
"Component": "aggregator_agent",
|
||||
"SessionName": self.session_name,
|
||||
}
|
||||
|
||||
self._event_buffer = MultiConsumerEventBuffer(
|
||||
max_size=MAX_EVENT_BUFFER_SIZE,
|
||||
max_batch_size=MAX_EVENT_SEND_BATCH_SIZE,
|
||||
common_metric_tags=self._common_tags,
|
||||
)
|
||||
self._executor = ThreadPoolExecutor(
|
||||
max_workers=THREAD_POOL_EXECUTOR_MAX_WORKERS,
|
||||
thread_name_prefix="aggregator_agent_executor",
|
||||
)
|
||||
|
||||
# Task metadata buffer accumulates dropped task attempts for GCS publishing
|
||||
self._task_metadata_buffer = TaskEventsMetadataBuffer(
|
||||
common_metric_tags=self._common_tags
|
||||
)
|
||||
|
||||
self._events_export_addr = (
|
||||
dashboard_agent.events_export_addr or EVENTS_EXPORT_ADDR
|
||||
)
|
||||
|
||||
self._event_processing_enabled = False
|
||||
if PUBLISH_EVENTS_TO_EXTERNAL_HTTP_SERVICE and self._events_export_addr:
|
||||
logger.info(
|
||||
f"Publishing events to external HTTP service is enabled. events_export_addr: {self._events_export_addr}"
|
||||
)
|
||||
self._event_processing_enabled = True
|
||||
self._http_endpoint_publisher = RayEventPublisher(
|
||||
name="http_service",
|
||||
publish_client=AsyncHttpPublisherClient(
|
||||
endpoint=self._events_export_addr,
|
||||
executor=self._executor,
|
||||
preserve_proto_field_name=PRESERVE_PROTO_FIELD_NAME,
|
||||
),
|
||||
event_buffer=self._event_buffer,
|
||||
common_metric_tags=self._common_tags,
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
f"Event HTTP target is not enabled or publishing events to external HTTP service is disabled. Skipping sending events to external HTTP service. events_export_addr: {self._events_export_addr}"
|
||||
)
|
||||
self._http_endpoint_publisher = NoopPublisher()
|
||||
|
||||
if PUBLISH_EVENTS_TO_GCS:
|
||||
logger.info("Publishing events to GCS is enabled")
|
||||
self._event_processing_enabled = True
|
||||
self._gcs_publisher = RayEventPublisher(
|
||||
name="ray_gcs",
|
||||
publish_client=AsyncGCSTaskEventsPublisherClient(
|
||||
gcs_client=self._dashboard_agent.gcs_client,
|
||||
executor=self._executor,
|
||||
),
|
||||
event_buffer=self._event_buffer,
|
||||
common_metric_tags=self._common_tags,
|
||||
task_metadata_buffer=self._task_metadata_buffer,
|
||||
)
|
||||
else:
|
||||
logger.info("Publishing events to GCS is disabled")
|
||||
self._gcs_publisher = NoopPublisher()
|
||||
|
||||
# Metrics
|
||||
self._open_telemetry_metric_recorder = OpenTelemetryMetricRecorder()
|
||||
|
||||
# Register counter metrics
|
||||
self._events_received_metric_name = (
|
||||
f"{AGGREGATOR_AGENT_METRIC_PREFIX}_events_received_total"
|
||||
)
|
||||
self._open_telemetry_metric_recorder.register_counter_metric(
|
||||
self._events_received_metric_name,
|
||||
"Total number of events received via AddEvents gRPC.",
|
||||
)
|
||||
|
||||
self._events_failed_to_add_metric_name = (
|
||||
f"{AGGREGATOR_AGENT_METRIC_PREFIX}_events_buffer_add_failures_total"
|
||||
)
|
||||
self._open_telemetry_metric_recorder.register_counter_metric(
|
||||
self._events_failed_to_add_metric_name,
|
||||
"Total number of events that failed to be added to the event buffer.",
|
||||
)
|
||||
|
||||
async def AddEvents(self, request, context) -> None:
|
||||
"""
|
||||
gRPC handler for adding events to the event aggregator. Receives events from the
|
||||
request and adds them to the event buffer.
|
||||
"""
|
||||
if not self._event_processing_enabled:
|
||||
return events_event_aggregator_service_pb2.AddEventsReply()
|
||||
|
||||
received_count = len(request.events_data.events)
|
||||
failed_count = 0
|
||||
events_data = request.events_data
|
||||
|
||||
if PUBLISH_EVENTS_TO_GCS:
|
||||
self._task_metadata_buffer.merge(events_data.task_events_metadata)
|
||||
|
||||
for event in events_data.events:
|
||||
try:
|
||||
await self._event_buffer.add_event(event)
|
||||
except Exception as e:
|
||||
failed_count += 1
|
||||
logger.error(
|
||||
f"Failed to add event with id={event.event_id.decode()} to buffer. "
|
||||
"Error: %s",
|
||||
e,
|
||||
)
|
||||
|
||||
if received_count > 0:
|
||||
self._open_telemetry_metric_recorder.set_metric_value(
|
||||
self._events_received_metric_name, self._common_tags, received_count
|
||||
)
|
||||
if failed_count > 0:
|
||||
self._open_telemetry_metric_recorder.set_metric_value(
|
||||
self._events_failed_to_add_metric_name, self._common_tags, failed_count
|
||||
)
|
||||
|
||||
return events_event_aggregator_service_pb2.AddEventsReply()
|
||||
|
||||
async def run(self, server) -> None:
|
||||
if server:
|
||||
events_event_aggregator_service_pb2_grpc.add_EventAggregatorServiceServicer_to_server(
|
||||
self, server
|
||||
)
|
||||
try:
|
||||
await asyncio.gather(
|
||||
self._http_endpoint_publisher.run_forever(),
|
||||
self._gcs_publisher.run_forever(),
|
||||
)
|
||||
finally:
|
||||
self._executor.shutdown()
|
||||
|
||||
@staticmethod
|
||||
def is_minimal_module() -> bool:
|
||||
return False
|
||||
Reference in New Issue
Block a user