diff --git a/src/ray/gcs/gcs_server/node_info_handler_impl.cc b/src/ray/gcs/gcs_server/node_info_handler_impl.cc index 257ce40e4..7f8c1edc0 100644 --- a/src/ray/gcs/gcs_server/node_info_handler_impl.cc +++ b/src/ray/gcs/gcs_server/node_info_handler_impl.cc @@ -61,5 +61,54 @@ void DefaultNodeInfoHandler::HandleGetAllNodeInfo( RAY_LOG(DEBUG) << "Finished getting all node info."; } +void DefaultNodeInfoHandler::HandleReportHeartbeat( + const ReportHeartbeatRequest &request, ReportHeartbeatReply *reply, + SendReplyCallback send_reply_callback) { + ClientID node_id = ClientID::FromBinary(request.heartbeat().client_id()); + RAY_LOG(DEBUG) << "Reporting heartbeat, node id = " << node_id; + + auto on_done = [node_id, send_reply_callback](Status status) { + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to report heartbeat: " << status.ToString() + << ", node id = " << node_id; + } + send_reply_callback(status, nullptr, nullptr); + }; + + auto heartbeat_data = std::make_shared(); + heartbeat_data->CopyFrom(request.heartbeat()); + Status status = gcs_client_.Nodes().AsyncReportHeartbeat(heartbeat_data, on_done); + if (!status.ok()) { + on_done(status); + } + RAY_LOG(DEBUG) << "Finished reporting heartbeat, node id = " << node_id; +} + +void DefaultNodeInfoHandler::HandleReportBatchHeartbeat( + const ReportBatchHeartbeatRequest &request, ReportBatchHeartbeatReply *reply, + SendReplyCallback send_reply_callback) { + RAY_LOG(DEBUG) << "Reporting batch heartbeat, batch size = " + << request.heartbeat_batch().batch_size(); + + auto on_done = [&request, send_reply_callback](Status status) { + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to report batch heartbeat: " << status.ToString() + << ", batch size = " << request.heartbeat_batch().batch_size(); + } + send_reply_callback(status, nullptr, nullptr); + }; + + auto heartbeat_batch_data = std::make_shared(); + heartbeat_batch_data->CopyFrom(request.heartbeat_batch()); + Status status = + gcs_client_.Nodes().AsyncReportBatchHeartbeat(heartbeat_batch_data, on_done); + if (!status.ok()) { + on_done(status); + } + + RAY_LOG(DEBUG) << "Finished reporting batch heartbeat, batch size = " + << request.heartbeat_batch().batch_size(); +} + } // namespace rpc } // namespace ray diff --git a/src/ray/gcs/gcs_server/node_info_handler_impl.h b/src/ray/gcs/gcs_server/node_info_handler_impl.h index f1165185e..e779044c4 100644 --- a/src/ray/gcs/gcs_server/node_info_handler_impl.h +++ b/src/ray/gcs/gcs_server/node_info_handler_impl.h @@ -24,6 +24,14 @@ class DefaultNodeInfoHandler : public rpc::NodeInfoHandler { GetAllNodeInfoReply *reply, SendReplyCallback send_reply_callback) override; + void HandleReportHeartbeat(const ReportHeartbeatRequest &request, + ReportHeartbeatReply *reply, + SendReplyCallback send_reply_callback) override; + + void HandleReportBatchHeartbeat(const ReportBatchHeartbeatRequest &request, + ReportBatchHeartbeatReply *reply, + SendReplyCallback send_reply_callback) override; + private: gcs::RedisGcsClient &gcs_client_; }; diff --git a/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc b/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc index a0a4df26e..e82ba53aa 100644 --- a/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc +++ b/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc @@ -144,6 +144,27 @@ class GcsServerTest : public RedisServiceManagerForTest { return node_info_list; } + bool ReportHeartbeat(const rpc::ReportHeartbeatRequest &request) { + std::promise promise; + client_->ReportHeartbeat(request, [&promise](const Status &status, + const rpc::ReportHeartbeatReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + return WaitReady(promise.get_future(), timeout_ms_); + } + + bool ReportBatchHeartbeat(const rpc::ReportBatchHeartbeatRequest &request) { + std::promise promise; + client_->ReportBatchHeartbeat( + request, + [&promise](const Status &status, const rpc::ReportBatchHeartbeatReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + return WaitReady(promise.get_future(), timeout_ms_); + } + bool AddObjectLocation(const rpc::AddObjectLocationRequest &request) { std::promise promise; client_->AddObjectLocation( @@ -291,6 +312,15 @@ TEST_F(GcsServerTest, TestNodeInfo) { ASSERT_TRUE(node_info_list[0].state() == rpc::GcsNodeInfo_GcsNodeState::GcsNodeInfo_GcsNodeState_ALIVE); + // Report heartbeat + rpc::ReportHeartbeatRequest report_heartbeat_request; + report_heartbeat_request.mutable_heartbeat()->set_client_id(node_id.Binary()); + ASSERT_TRUE(ReportHeartbeat(report_heartbeat_request)); + rpc::ReportBatchHeartbeatRequest report_batch_heartbeat_request; + report_batch_heartbeat_request.mutable_heartbeat_batch()->add_batch()->set_client_id( + node_id.Binary()); + ASSERT_TRUE(ReportBatchHeartbeat(report_batch_heartbeat_request)); + // Unregister node info rpc::UnregisterNodeRequest unregister_node_info_request; unregister_node_info_request.set_node_id(node_id.Binary()); diff --git a/src/ray/protobuf/gcs_service.proto b/src/ray/protobuf/gcs_service.proto index c04d6bab0..77883f76d 100644 --- a/src/ray/protobuf/gcs_service.proto +++ b/src/ray/protobuf/gcs_service.proto @@ -89,6 +89,20 @@ message GetAllNodeInfoReply { repeated GcsNodeInfo node_info_list = 1; } +message ReportHeartbeatRequest { + HeartbeatTableData heartbeat = 1; +} + +message ReportHeartbeatReply { +} + +message ReportBatchHeartbeatRequest { + HeartbeatBatchTableData heartbeat_batch = 1; +} + +message ReportBatchHeartbeatReply { +} + // Service for node info access. service NodeInfoGcsService { // Register a node to GCS Service. @@ -97,6 +111,11 @@ service NodeInfoGcsService { rpc UnregisterNode(UnregisterNodeRequest) returns (UnregisterNodeReply); // Get information of all nodes from GCS Service. rpc GetAllNodeInfo(GetAllNodeInfoRequest) returns (GetAllNodeInfoReply); + // Report heartbeat of a node to GCS Service. + rpc ReportHeartbeat(ReportHeartbeatRequest) returns (ReportHeartbeatReply); + // Report batch heartbeat to GCS Service. + rpc ReportBatchHeartbeat(ReportBatchHeartbeatRequest) + returns (ReportBatchHeartbeatReply); } message GetObjectLocationsRequest { diff --git a/src/ray/rpc/gcs_server/gcs_rpc_client.h b/src/ray/rpc/gcs_server/gcs_rpc_client.h index 283770873..9e730c4da 100644 --- a/src/ray/rpc/gcs_server/gcs_rpc_client.h +++ b/src/ray/rpc/gcs_server/gcs_rpc_client.h @@ -76,6 +76,30 @@ class GcsRpcClient { VOID_RPC_CLIENT_METHOD(ObjectInfoGcsService, RemoveObjectLocation, request, callback, object_info_grpc_client_) + /// Report heartbeat of a node to GCS Service. + /// + /// \param request The request message. + /// \param callback The callback function that handles reply from server. + void ReportHeartbeat(const ReportHeartbeatRequest &request, + const ClientCallback &callback) { + client_call_manager_ + .CreateCall( + *node_info_stub_, &NodeInfoGcsService::Stub::PrepareAsyncReportHeartbeat, + request, callback); + } + + /// Report batch heartbeat to GCS Service. + /// + /// \param request The request message. + /// \param callback The callback function that handles reply from server. + void ReportBatchHeartbeat(const ReportBatchHeartbeatRequest &request, + const ClientCallback &callback) { + client_call_manager_.CreateCall( + *node_info_stub_, &NodeInfoGcsService::Stub::PrepareAsyncReportBatchHeartbeat, + request, callback); + } + private: /// The gRPC-generated stub. std::unique_ptr> job_info_grpc_client_; diff --git a/src/ray/rpc/gcs_server/gcs_rpc_server.h b/src/ray/rpc/gcs_server/gcs_rpc_server.h index 888965750..bd3b5958f 100644 --- a/src/ray/rpc/gcs_server/gcs_rpc_server.h +++ b/src/ray/rpc/gcs_server/gcs_rpc_server.h @@ -146,6 +146,14 @@ class NodeInfoHandler { virtual void HandleGetAllNodeInfo(const GetAllNodeInfoRequest &request, GetAllNodeInfoReply *reply, SendReplyCallback send_reply_callback) = 0; + + virtual void HandleReportHeartbeat(const ReportHeartbeatRequest &request, + ReportHeartbeatReply *reply, + SendReplyCallback send_reply_callback) = 0; + + virtual void HandleReportBatchHeartbeat(const ReportBatchHeartbeatRequest &request, + ReportBatchHeartbeatReply *reply, + SendReplyCallback send_reply_callback) = 0; }; /// The `GrpcService` for `NodeInfoGcsService`. @@ -168,6 +176,8 @@ class NodeInfoGrpcService : public GrpcService { NODE_INFO_SERVICE_RPC_HANDLER(RegisterNode, 1); NODE_INFO_SERVICE_RPC_HANDLER(UnregisterNode, 1); NODE_INFO_SERVICE_RPC_HANDLER(GetAllNodeInfo, 1); + NODE_INFO_SERVICE_RPC_HANDLER(ReportHeartbeat, 1); + NODE_INFO_SERVICE_RPC_HANDLER(ReportBatchHeartbeat, 1); } private: