[core] Ensure client_call_manager_ outlives metrics_agent_client_ in core worker (ray-project#58315)

The `metrics_agent_client_` depends on `client_call_manager_`, but
previously it was pulling out a reference to it from the core worker,
which is not guaranteed to outlive the agent client.

Modifying it to keep the `client_call_manager_` as a field of the
`core_worker_process` instead.

I think we may also need to drain any ongoing RPCs from the
`metrics_agent_client_` on shutdown. Leaving that for a future PR.

---------

Signed-off-by: Edward Oakes <ed.nmi.oakes@gmail.com>

Author: edoakes

src/ray/core_worker/core_worker.cc

@@ -289,7 +289,6 @@ CoreWorker::CoreWorker(
     CoreWorkerOptions options,
     std::unique_ptr<WorkerContext> worker_context,
     instrumented_io_context &io_service,
-    std::unique_ptr<rpc::ClientCallManager> client_call_manager,
     std::shared_ptr<rpc::CoreWorkerClientPool> core_worker_client_pool,
     std::shared_ptr<rpc::RayletClientPool> raylet_client_pool,
     std::shared_ptr<PeriodicalRunnerInterface> periodical_runner,
@@ -325,7 +324,6 @@ CoreWorker::CoreWorker(
                          : nullptr),
       worker_context_(std::move(worker_context)),
       io_service_(io_service),
-      client_call_manager_(std::move(client_call_manager)),
       core_worker_client_pool_(std::move(core_worker_client_pool)),
       raylet_client_pool_(std::move(raylet_client_pool)),
       periodical_runner_(std::move(periodical_runner)),

src/ray/core_worker/core_worker.h

@@ -176,7 +176,6 @@ class CoreWorker {
   CoreWorker(CoreWorkerOptions options,
              std::unique_ptr<WorkerContext> worker_context,
              instrumented_io_context &io_service,
-             std::unique_ptr<rpc::ClientCallManager> client_call_manager,
              std::shared_ptr<rpc::CoreWorkerClientPool> core_worker_client_pool,
              std::shared_ptr<rpc::RayletClientPool> raylet_client_pool,
              std::shared_ptr<PeriodicalRunnerInterface> periodical_runner,
@@ -1737,9 +1736,6 @@ class CoreWorker {
   /// Event loop where the IO events are handled. e.g. async GCS operations.
   instrumented_io_context &io_service_;
 
-  /// Shared client call manager.
-  std::unique_ptr<rpc::ClientCallManager> client_call_manager_;
-
   /// Shared core worker client pool.
   std::shared_ptr<rpc::CoreWorkerClientPool> core_worker_client_pool_;
 

src/ray/core_worker/core_worker_process.cc

@@ -132,8 +132,6 @@ std::shared_ptr<CoreWorker> CoreWorkerProcess::TryGetWorker() {
 std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
     CoreWorkerOptions options, const WorkerID &worker_id) {
   /// Event loop where the IO events are handled. e.g. async GCS operations.
-  auto client_call_manager = std::make_unique<rpc::ClientCallManager>(
-      io_service_, /*record_stats=*/false, options.node_ip_address);
   auto periodical_runner = PeriodicalRunner::Create(io_service_);
   auto worker_context = std::make_unique<WorkerContext>(
       options.worker_type, worker_id, GetProcessJobID(options));
@@ -166,7 +164,7 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
   auto task_event_buffer = std::make_unique<worker::TaskEventBufferImpl>(
       std::make_unique<gcs::GcsClient>(options.gcs_options, options.node_ip_address),
       std::make_unique<rpc::EventAggregatorClientImpl>(options.metrics_agent_port,
-                                                       *client_call_manager),
+                                                       *client_call_manager_),
       options.session_name);
 
   // Start the IO thread first to make sure the checker is working.
@@ -237,7 +235,7 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
       local_node_id, options.node_ip_address, options.node_manager_port);
   auto local_raylet_rpc_client =
       std::make_shared<rpc::RayletClient>(std::move(raylet_address),
-                                          *client_call_manager,
+                                          *client_call_manager_,
                                           /*raylet_unavailable_timeout_callback=*/[] {});
   auto core_worker_server =
       std::make_unique<rpc::GrpcServer>(WorkerTypeString(options.worker_type),
@@ -277,7 +275,7 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
         auto core_worker = GetCoreWorker();
         return std::make_shared<ray::rpc::RayletClient>(
             addr,
-            *core_worker->client_call_manager_,
+            *client_call_manager_,
             rpc::RayletClientPool::GetDefaultUnavailableTimeoutCallback(
                 core_worker->gcs_client_.get(),
                 core_worker->raylet_client_pool_.get(),
@@ -289,7 +287,7 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
         auto core_worker = GetCoreWorker();
         return std::make_shared<rpc::CoreWorkerClient>(
             addr,
-            *core_worker->client_call_manager_,
+            *client_call_manager_,
             rpc::CoreWorkerClientPool::GetDefaultUnavailableTimeoutCallback(
                 core_worker->gcs_client_.get(),
                 core_worker->core_worker_client_pool_.get(),
@@ -658,7 +656,6 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
       std::make_shared<CoreWorker>(std::move(options),
                                    std::move(worker_context),
                                    io_service_,
-                                   std::move(client_call_manager),
                                    std::move(core_worker_client_pool),
                                    std::move(raylet_client_pool),
                                    std::move(periodical_runner),
@@ -696,6 +693,8 @@ CoreWorkerProcessImpl::CoreWorkerProcessImpl(const CoreWorkerOptions &options)
                      ? ComputeDriverIdFromJob(options_.job_id)
                      : WorkerID::FromRandom()),
       io_work_(io_service_.get_executor()),
+      client_call_manager_(std::make_unique<rpc::ClientCallManager>(
+          io_service_, /*record_stats=*/false, options.node_ip_address)),
       task_execution_service_work_(task_execution_service_.get_executor()),
       service_handler_(std::make_unique<CoreWorkerServiceHandlerProxy>()) {
   if (options_.enable_logging) {
@@ -824,10 +823,7 @@ CoreWorkerProcessImpl::CoreWorkerProcessImpl(const CoreWorkerOptions &options)
     write_locked.Get() = worker;
     // Initialize metrics agent client.
     metrics_agent_client_ = std::make_unique<ray::rpc::MetricsAgentClientImpl>(
-        "127.0.0.1",
-        options_.metrics_agent_port,
-        io_service_,
-        *write_locked.Get()->client_call_manager_);
+        "127.0.0.1", options_.metrics_agent_port, io_service_, *client_call_manager_);
     metrics_agent_client_->WaitForServerReady([this](const Status &server_status) {
       if (server_status.ok()) {
         stats::InitOpenTelemetryExporter(options_.metrics_agent_port);

src/ray/core_worker/core_worker_process.h

@@ -153,9 +153,6 @@ class CoreWorkerProcessImpl {
   /// The various options.
   const CoreWorkerOptions options_;
 
-  /// The core worker instance of this worker process.
-  MutexProtected<std::shared_ptr<CoreWorker>> core_worker_;
-
   /// The worker ID of this worker.
   const WorkerID worker_id_;
 
@@ -166,6 +163,10 @@ class CoreWorkerProcessImpl {
   /// Keeps the io_service_ alive.
   boost::asio::executor_work_guard<boost::asio::io_context::executor_type> io_work_;
 
+  /// Shared client call manager across all gRPC clients in the core worker process.
+  /// This is used by the CoreWorker and the MetricsAgentClient.
+  std::unique_ptr<rpc::ClientCallManager> client_call_manager_;
+
   /// Event loop where tasks are processed.
   /// task_execution_service_ should be destructed first to avoid
   /// issues like https://github.com/ray-project/ray/issues/18857
@@ -179,6 +180,9 @@ class CoreWorkerProcessImpl {
   // Thread that runs a boost::asio service to process IO events.
   boost::thread io_thread_;
 
+  /// The core worker instance of this worker process.
+  MutexProtected<std::shared_ptr<CoreWorker>> core_worker_;
+
   /// The proxy service handler that routes the RPC calls to the core worker.
   std::unique_ptr<CoreWorkerServiceHandlerProxy> service_handler_;
 

src/ray/core_worker/tests/core_worker_test.cc

@@ -96,7 +96,7 @@ class CoreWorkerTest : public ::testing::Test {
       return Status::OK();
     };
 
-    auto client_call_manager = std::make_unique<rpc::ClientCallManager>(
+    client_call_manager_ = std::make_unique<rpc::ClientCallManager>(
         io_service_, /*record_stats=*/false, /*local_address=*/"");
 
     auto core_worker_client_pool =
@@ -167,7 +167,7 @@ class CoreWorkerTest : public ::testing::Test {
 
     auto task_event_buffer = std::make_unique<worker::TaskEventBufferImpl>(
         std::make_unique<gcs::MockGcsClient>(),
-        std::make_unique<rpc::EventAggregatorClientImpl>(0, *client_call_manager),
+        std::make_unique<rpc::EventAggregatorClientImpl>(0, *client_call_manager_),
         "test_session");
 
     task_manager_ = std::make_shared<TaskManager>(
@@ -248,7 +248,6 @@ class CoreWorkerTest : public ::testing::Test {
     core_worker_ = std::make_shared<CoreWorker>(std::move(options),
                                                 std::move(worker_context),
                                                 io_service_,
-                                                std::move(client_call_manager),
                                                 std::move(core_worker_client_pool),
                                                 std::move(raylet_client_pool),
                                                 std::move(periodical_runner),
@@ -289,6 +288,7 @@ class CoreWorkerTest : public ::testing::Test {
   boost::thread io_thread_;
 
   rpc::Address rpc_address_;
+  std::unique_ptr<rpc::ClientCallManager> client_call_manager_;
   std::shared_ptr<ReferenceCounterInterface> reference_counter_;
   std::shared_ptr<CoreWorkerMemoryStore> memory_store_;
   ActorTaskSubmitter *actor_task_submitter_;