[Core] Add Logic to Emit Task Events to Event Aggregator (ray-project#53402)

This PR adds the logic to emit task events to event aggregator, mainly:
* Update the task event buffer logic to convert and send the ray event
to the event aggregator at the same time when the task events are
converted and sent to GCS
* Added 2 configs to control turning on/off the path to send to GCS and
the path to send to event aggregator
* Added test task event buffer tests for the ray event to aggregator
path

---------

Signed-off-by: Mengjin Yan <mengjinyan3@gmail.com>
Signed-off-by: myan <myan@anyscale.com>

Author: MengjinYan

BUILD.bazel

@@ -196,7 +196,9 @@ ray_cc_library(
     name = "event_aggregator_client",
     hdrs = ["src/ray/rpc/event_aggregator_client.h"],
     deps = [
+        ":event_aggregator_cc_rpc",
         ":grpc_client",
+        "//src/ray/common:status",
         "//src/ray/protobuf:events_event_aggregator_service_cc_proto",
         "//src/ray/util:logging",
         "@com_github_grpc_grpc//:grpc++",
@@ -389,6 +391,13 @@ ray_cc_library(
     ],
 )
 
+cc_grpc_library(
+    name = "event_aggregator_cc_rpc",
+    srcs = ["//src/ray/protobuf:events_event_aggregator_service_proto"],
+    grpc_only = True,
+    deps = ["//src/ray/protobuf:events_event_aggregator_service_cc_proto"],
+)
+
 # pubsub.
 cc_grpc_library(
     name = "pubsub_cc_grpc",

python/ray/dashboard/modules/aggregator/aggregator_agent.py

@@ -213,10 +213,10 @@ def _receive_events(self, request):
                 status_message += (
                     f", and {len(error_messages) - truncate_num} more events dropped"
                 )
-        status = events_event_aggregator_service_pb2.AddEventStatus(
+        status = events_event_aggregator_service_pb2.AddEventsStatus(
             code=status_code, message=status_message
         )
-        return events_event_aggregator_service_pb2.AddEventReply(status=status)
+        return events_event_aggregator_service_pb2.AddEventsReply(status=status)
 
     def _send_events_to_external_service(self, event_batch):
         """

python/ray/dashboard/modules/aggregator/tests/test_aggregator_agent.py

@@ -22,7 +22,7 @@
     EventAggregatorServiceStub,
 )
 from ray.core.generated.events_event_aggregator_service_pb2 import (
-    AddEventRequest,
+    AddEventsRequest,
     RayEventsData,
     TaskEventsMetadata,
 )
@@ -92,7 +92,7 @@ def test_aggregator_agent_receive_publish_events_normally(
     seconds, nanos = divmod(test_time, 10**9)
     timestamp = Timestamp(seconds=seconds, nanos=nanos)
 
-    request = AddEventRequest(
+    request = AddEventsRequest(
         events_data=RayEventsData(
             events=[
                 RayEvent(
@@ -154,7 +154,7 @@ def test_aggregator_agent_receive_event_full(
     seconds, nanos = divmod(now, 10**9)
     timestamp = Timestamp(seconds=seconds, nanos=nanos)
 
-    request = AddEventRequest(
+    request = AddEventsRequest(
         events_data=RayEventsData(
             events=[
                 RayEvent(
@@ -196,7 +196,7 @@ def test_aggregator_agent_receive_dropped_at_core_worker(
     seconds, nanos = divmod(now, 10**9)
     timestamp = Timestamp(seconds=seconds, nanos=nanos)
 
-    request = AddEventRequest(
+    request = AddEventsRequest(
         events_data=RayEventsData(
             events=[
                 RayEvent(
@@ -247,7 +247,7 @@ def test_aggregator_agent_receive_multiple_events(
     now = time.time_ns()
     seconds, nanos = divmod(now, 10**9)
     timestamp = Timestamp(seconds=seconds, nanos=nanos)
-    request = AddEventRequest(
+    request = AddEventsRequest(
         events_data=RayEventsData(
             events=[
                 RayEvent(
@@ -306,7 +306,7 @@ def test_aggregator_agent_receive_multiple_events_failures(
     now = time.time_ns()
     seconds, nanos = divmod(now, 10**9)
     timestamp = Timestamp(seconds=seconds, nanos=nanos)
-    request = AddEventRequest(
+    request = AddEventsRequest(
         events_data=RayEventsData(
             events=[
                 RayEvent(
@@ -353,7 +353,7 @@ def test_aggregator_agent_receive_empty_events(
         cluster.webui_url, cluster.gcs_address, cluster.head_node.node_id
     )
     httpserver.expect_request("/", method="POST").respond_with_data("", status=200)
-    request = AddEventRequest(
+    request = AddEventsRequest(
         events_data=RayEventsData(
             events=[],
             task_events_metadata=TaskEventsMetadata(

python/ray/tests/test_metrics_agent.py

@@ -22,7 +22,7 @@
 from ray.core.generated.common_pb2 import TaskAttempt
 from ray.core.generated.events_base_event_pb2 import RayEvent
 from ray.core.generated.events_event_aggregator_service_pb2 import (
-    AddEventRequest,
+    AddEventsRequest,
     RayEventsData,
     TaskEventsMetadata,
 )
@@ -528,7 +528,7 @@ def test_case_value_correct():
     now = time.time_ns()
     seconds, nanos = divmod(now, 10**9)
     timestamp = Timestamp(seconds=seconds, nanos=nanos)
-    request = AddEventRequest(
+    request = AddEventsRequest(
         events_data=RayEventsData(
             events=[
                 RayEvent(

src/ray/common/grpc_util.h

@@ -16,6 +16,7 @@
 
 #include <google/protobuf/map.h>
 #include <google/protobuf/repeated_field.h>
+#include <google/protobuf/timestamp.pb.h>
 #include <google/protobuf/util/message_differencer.h>
 #include <grpcpp/grpcpp.h>
 
@@ -231,4 +232,15 @@ inline grpc::ChannelArguments CreateDefaultChannelArguments() {
   return arguments;
 }
 
+// Convert an epoch time in nanoseconds to a protobuf timestamp
+// Ref: https://protobuf.dev/reference/php/api-docs/Google/Protobuf/Timestamp.html
+inline google::protobuf::Timestamp AbslTimeNanosToProtoTimestamp(int64_t nanos) {
+  google::protobuf::Timestamp timestamp;
+
+  // Extract the seconds and the fractional nanoseconds from the epoch time
+  timestamp.set_seconds(nanos / 1000000000);
+  timestamp.set_nanos(nanos % 1000000000);
+  return timestamp;
+}
+
 }  // namespace ray

src/ray/common/ray_config_def.h

@@ -921,6 +921,17 @@ RAY_CONFIG(bool, enable_export_api_write, false)
 // Example config: `export RAY_enable_export_api_write_config='EXPORT_ACTOR,EXPORT_TASK'`
 RAY_CONFIG(std::vector<std::string>, enable_export_api_write_config, {})
 
+// Whether the task events from the core worker are sent to GCS directly.
+// TODO(myan): #54515 Remove this flag after the task events to GCS path is fully
+// migrated to the event aggregator.
+RAY_CONFIG(bool, enable_core_worker_task_event_to_gcs, true)
+
+// Whether to enable the ray event to send to the event aggregator.
+// Currently, only task events are supported.
+// TODO(myan): #54515 Remove this flag after the task events are fully migrated to the
+// event aggregator.
+RAY_CONFIG(bool, enable_core_worker_ray_event_to_aggregator, false)
+
 // Configuration for pipe logger buffer size.
 RAY_CONFIG(uint64_t, pipe_logger_read_buf_size, 1024)
 

src/ray/core_worker/BUILD.bazel

@@ -207,6 +207,7 @@ ray_cc_library(
     srcs = ["task_event_buffer.cc"],
     hdrs = ["task_event_buffer.h"],
     deps = [
+        "//:event_aggregator_client",
         "//src/ray/common:asio",
         "//src/ray/common:id",
         "//src/ray/common:task_common",

src/ray/core_worker/core_worker.cc

@@ -41,6 +41,7 @@
 #include "ray/common/task/task_util.h"
 #include "ray/gcs/gcs_client/gcs_client.h"
 #include "ray/gcs/pb_util.h"
+#include "ray/rpc/event_aggregator_client.h"
 #include "ray/util/container_util.h"
 #include "ray/util/event.h"
 #include "ray/util/subreaper.h"
@@ -399,7 +400,7 @@ CoreWorker::CoreWorker(
                               options_.language,
                               worker_context_->GetCurrentJobID(),
                               // Driver has no parent task
-                              /* parent_task_id */ TaskID::Nil(),
+                              /*parent_task_id=*/TaskID::Nil(),
                               GetCallerId(),
                               rpc_address_,
                               TaskID::Nil());
@@ -417,6 +418,7 @@ CoreWorker::CoreWorker(
           /*attempt_number=*/0,
           rpc::TaskStatus::RUNNING,
           /*timestamp=*/absl::GetCurrentTimeNanos(),
+          /*is_actor_task=*/false,
           std::make_shared<const TaskSpecification>(std::move(spec)));
       task_event_buffer_->AddTaskEvent(std::move(task_event));
     }
@@ -588,7 +590,8 @@ void CoreWorker::Disconnect(
         worker_context_->GetCurrentJobID(),
         /* attempt_number */ 0,
         rpc::TaskStatus::FINISHED,
-        /* timestamp */ absl::GetCurrentTimeNanos());
+        /*timestamp=*/absl::GetCurrentTimeNanos(),
+        /*is_actor_task_event=*/worker_context_->GetCurrentActorID().IsNil());
     task_event_buffer_->AddTaskEvent(std::move(task_event));
   }
 
@@ -1041,7 +1044,7 @@ Status CoreWorker::PutInLocalPlasmaStore(const RayObject &object,
                                          bool pin_object) {
   bool object_exists = false;
   RAY_RETURN_NOT_OK(plasma_store_provider_->Put(
-      object, object_id, /* owner_address = */ rpc_address_, &object_exists));
+      object, object_id, /*owner_address=*/rpc_address_, &object_exists));
   if (!object_exists) {
     if (pin_object) {
       // Tell the raylet to pin the object **after** it is created.
@@ -1158,7 +1161,7 @@ Status CoreWorker::CreateOwnedAndIncrementLocalRef(
       status = plasma_store_provider_->Create(metadata,
                                               data_size,
                                               *object_id,
-                                              /* owner_address = */ real_owner_address,
+                                              /*owner_address=*/real_owner_address,
                                               data,
                                               /*created_by_worker=*/true,
                                               is_experimental_mutable_object);
@@ -1200,7 +1203,7 @@ Status CoreWorker::ExperimentalChannelWriteAcquire(
     int64_t timeout_ms,
     std::shared_ptr<Buffer> *data) {
   Status status = experimental_mutable_object_provider_->GetChannelStatus(
-      object_id, /*is_reader*/ false);
+      object_id, /*is_reader=*/false);
   if (!status.ok()) {
     return status;
   }
@@ -1350,7 +1353,7 @@ Status CoreWorker::Get(const std::vector<ObjectID> &ids,
   bool is_experimental_channel = false;
   for (const ObjectID &id : ids) {
     Status status =
-        experimental_mutable_object_provider_->GetChannelStatus(id, /*is_reader*/ true);
+        experimental_mutable_object_provider_->GetChannelStatus(id, /*is_reader=*/true);
     if (status.ok()) {
       is_experimental_channel = true;
       // We continue rather than break because we want to check that *all* of the
@@ -2052,7 +2055,7 @@ std::vector<rpc::ObjectReference> CoreWorker::SubmitTask(
                       /*include_job_config*/ true,
                       /*generator_backpressure_num_objects*/
                       task_options.generator_backpressure_num_objects,
-                      /*enable_task_event*/ task_options.enable_task_events,
+                      /*enable_task_event=*/task_options.enable_task_events,
                       task_options.labels,
                       task_options.label_selector);
   ActorID root_detached_actor_id;
@@ -2136,10 +2139,10 @@ Status CoreWorker::CreateActor(const RayFunction &function,
                       rpc_address_,
                       function,
                       args,
-                      /*num_returns*/ 0,
+                      /*num_returns=*/0,
                       new_resource,
                       new_placement_resources,
-                      "" /* debugger_breakpoint */,
+                      /*debugger_breakpoint=*/"",
                       depth,
                       actor_creation_options.serialized_runtime_env_info,
                       call_site,
@@ -2242,7 +2245,7 @@ Status CoreWorker::CreateActor(const RayFunction &function,
       task_spec,
       CurrentCallSite(),
       // Actor creation task retry happens on GCS not on core worker.
-      /*max_retries*/ 0);
+      /*max_retries=*/0);
 
   if (actor_name.empty()) {
     io_service_.post(
@@ -2416,19 +2419,19 @@ Status CoreWorker::SubmitActorTask(
                       task_options.num_returns,
                       task_options.resources,
                       required_resources,
-                      "",    /* debugger_breakpoint */
-                      depth, /*depth*/
-                      "{}",  /* serialized_runtime_env_info */
+                      /*debugger_breakpoint=*/"",
+                      depth,
+                      /*serialized_runtime_env_info=*/"{}",
                       call_site,
                       worker_context_->GetMainThreadOrActorCreationTaskID(),
                       task_options.concurrency_group_name,
-                      /*include_job_config*/ false,
-                      /*generator_backpressure_num_objects*/
+                      /*include_job_config=*/false,
+                      /*generator_backpressure_num_objects=*/
                       task_options.generator_backpressure_num_objects,
-                      /*enable_task_events*/ task_options.enable_task_events,
-                      /*labels*/ {},
-                      /*label_selector*/ {},
-                      /*tensor_transport*/ task_options.tensor_transport);
+                      /*enable_task_events=*/task_options.enable_task_events,
+                      /*labels=*/{},
+                      /*label_selector=*/{},
+                      /*tensor_transport=*/task_options.tensor_transport);
   // NOTE: placement_group_capture_child_tasks and runtime_env will
   // be ignored in the actor because we should always follow the actor's option.
 
@@ -2824,7 +2827,7 @@ Status CoreWorker::ExecuteTask(
                                                           task_spec.AttemptNumber(),
                                                           task_spec,
                                                           rpc::TaskStatus::RUNNING,
-                                                          /* include_task_info */ false,
+                                                          /*include_task_info=*/false,
                                                           update));
 
     worker_context_->SetCurrentTask(task_spec);
@@ -3212,7 +3215,7 @@ void CoreWorker::HandleReportGeneratorItemReturns(
   auto worker_id = WorkerID::FromBinary(request.worker_addr().worker_id());
   task_manager_->HandleReportGeneratorItemReturns(
       request,
-      /*execution_signal_callback*/
+      /*execution_signal_callback=*/
       [reply,
        worker_id = std::move(worker_id),
        generator_id = std::move(generator_id),
@@ -3736,8 +3739,8 @@ void CoreWorker::HandleUpdateObjectLocationBatch(
   }
 
   send_reply_callback(Status::OK(),
-                      /*success_callback_on_reply*/ nullptr,
-                      /*failure_callback_on_reply*/ nullptr);
+                      /*success_callback_on_reply=*/nullptr,
+                      /*failure_callback_on_reply=*/nullptr);
 }
 
 void CoreWorker::AddSpilledObjectLocationOwner(
@@ -3975,7 +3978,7 @@ void CoreWorker::CancelTaskOnExecutor(TaskID task_id,
     }
   }
 
-  on_canceled(/*success*/ success, /*requested_task_running*/ requested_task_running);
+  on_canceled(/*success=*/success, /*requested_task_running=*/requested_task_running);
 }
 
 void CoreWorker::CancelActorTaskOnExecutor(WorkerID caller_worker_id,
@@ -4567,7 +4570,7 @@ void CoreWorker::RecordTaskLogStart(const TaskID &task_id,
       attempt_number,
       *current_task,
       rpc::TaskStatus::NIL,
-      /* include_task_info */ false,
+      /*include_task_info=*/false,
       worker::TaskStatusEvent::TaskStateUpdate(task_log_info)));
 }
 
@@ -4591,7 +4594,7 @@ void CoreWorker::RecordTaskLogEnd(const TaskID &task_id,
       attempt_number,
       *current_task,
       rpc::TaskStatus::NIL,
-      /* include_task_info */ false,
+      /*include_task_info=*/false,
       worker::TaskStatusEvent::TaskStateUpdate(task_log_info)));
 }
 
@@ -4609,7 +4612,7 @@ void CoreWorker::UpdateTaskIsDebuggerPaused(const TaskID &task_id,
       running_task_it->second.AttemptNumber(),
       running_task_it->second,
       rpc::TaskStatus::NIL,
-      /* include_task_info */ false,
+      /*include_task_info=*/false,
       worker::TaskStatusEvent::TaskStateUpdate(is_debugger_paused)));
 }
 

src/ray/core_worker/core_worker_process.cc

@@ -171,7 +171,9 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
   }
 
   auto task_event_buffer = std::make_unique<worker::TaskEventBufferImpl>(
-      std::make_shared<gcs::GcsClient>(options.gcs_options));
+      std::make_unique<gcs::GcsClient>(options.gcs_options),
+      std::make_unique<rpc::EventAggregatorClientImpl>(options.metrics_agent_port,
+                                                       *client_call_manager));
 
   // Start the IO thread first to make sure the checker is working.
   boost::thread::attributes io_thread_attrs;