diff --git a/src/ray/gcs/accessor.h b/src/ray/gcs/accessor.h index e4d01d21e..e21770b71 100644 --- a/src/ray/gcs/accessor.h +++ b/src/ray/gcs/accessor.h @@ -205,6 +205,12 @@ class NodeInfoAccessor { /// \return GcsNodeInfo virtual const rpc::GcsNodeInfo &GetSelfInfo() const = 0; + /// Register node to GCS synchronously. + /// + /// \param node_info The information of node to register to GCS. + /// \return Status + virtual Status Register(const rpc::GcsNodeInfo &node_info) = 0; + /// Cancel registration of a node to GCS asynchronously. /// /// \param node_id The ID of node that to be unregistered. diff --git a/src/ray/gcs/gcs_server/actor_info_handler_impl.cc b/src/ray/gcs/gcs_server/actor_info_handler_impl.cc index bcaaaa14d..628332e3d 100644 --- a/src/ray/gcs/gcs_server/actor_info_handler_impl.cc +++ b/src/ray/gcs/gcs_server/actor_info_handler_impl.cc @@ -45,10 +45,7 @@ void DefaultActorInfoHandler::HandleRegisterActorInfo( send_reply_callback(status, nullptr, nullptr); }; - Status status = gcs_client_.Actors().AsyncRegister( - actor_table_data, [send_reply_callback](Status status) { - send_reply_callback(status, nullptr, nullptr); - }); + Status status = gcs_client_.Actors().AsyncRegister(actor_table_data, on_done); if (!status.ok()) { on_done(status); } @@ -70,10 +67,7 @@ void DefaultActorInfoHandler::HandleUpdateActorInfo( send_reply_callback(status, nullptr, nullptr); }; - Status status = gcs_client_.Actors().AsyncUpdate( - actor_id, actor_table_data, [send_reply_callback](Status status) { - send_reply_callback(status, nullptr, nullptr); - }); + Status status = gcs_client_.Actors().AsyncUpdate(actor_id, actor_table_data, on_done); if (!status.ok()) { on_done(status); } diff --git a/src/ray/gcs/gcs_server/gcs_server.cc b/src/ray/gcs/gcs_server/gcs_server.cc index f4a647584..4041445a3 100644 --- a/src/ray/gcs/gcs_server/gcs_server.cc +++ b/src/ray/gcs/gcs_server/gcs_server.cc @@ -1,6 +1,7 @@ #include "gcs_server.h" #include "actor_info_handler_impl.h" #include "job_info_handler_impl.h" +#include "node_info_handler_impl.h" namespace ray { namespace gcs { @@ -26,6 +27,11 @@ void GcsServer::Start() { new rpc::ActorInfoGrpcService(main_service_, *actor_info_handler_)); rpc_server_.RegisterService(*actor_info_service_); + node_info_handler_ = InitNodeInfoHandler(); + node_info_service_.reset( + new rpc::NodeInfoGrpcService(main_service_, *node_info_handler_)); + rpc_server_.RegisterService(*node_info_service_); + // Run rpc server. rpc_server_.Run(); @@ -62,5 +68,10 @@ std::unique_ptr GcsServer::InitActorInfoHandler() { new rpc::DefaultActorInfoHandler(*redis_gcs_client_)); } +std::unique_ptr GcsServer::InitNodeInfoHandler() { + return std::unique_ptr( + new rpc::DefaultNodeInfoHandler(*redis_gcs_client_)); +} + } // namespace gcs } // namespace ray diff --git a/src/ray/gcs/gcs_server/gcs_server.h b/src/ray/gcs/gcs_server/gcs_server.h index 5bc7b408e..46fea49b7 100644 --- a/src/ray/gcs/gcs_server/gcs_server.h +++ b/src/ray/gcs/gcs_server/gcs_server.h @@ -52,6 +52,9 @@ class GcsServer { /// The actor info handler virtual std::unique_ptr InitActorInfoHandler(); + /// The node info handler + virtual std::unique_ptr InitNodeInfoHandler(); + private: /// Gcs server configuration GcsServerConfig config_; @@ -65,6 +68,9 @@ class GcsServer { /// Actor info handler and service std::unique_ptr actor_info_handler_; std::unique_ptr actor_info_service_; + /// Node info handler and service + std::unique_ptr node_info_handler_; + std::unique_ptr node_info_service_; /// Backend client std::shared_ptr redis_gcs_client_; }; diff --git a/src/ray/gcs/gcs_server/job_info_handler_impl.cc b/src/ray/gcs/gcs_server/job_info_handler_impl.cc index 85da8e6b6..b64720831 100644 --- a/src/ray/gcs/gcs_server/job_info_handler_impl.cc +++ b/src/ray/gcs/gcs_server/job_info_handler_impl.cc @@ -19,11 +19,7 @@ void DefaultJobInfoHandler::HandleAddJob(const rpc::AddJobRequest &request, send_reply_callback(status, nullptr, nullptr); }; - Status status = gcs_client_.Jobs().AsyncAdd( - job_table_data, [reply, send_reply_callback](Status status) { - reply->set_success(status.ok()); - send_reply_callback(status, nullptr, nullptr); - }); + Status status = gcs_client_.Jobs().AsyncAdd(job_table_data, on_done); if (!status.ok()) { on_done(status); } @@ -44,11 +40,7 @@ void DefaultJobInfoHandler::HandleMarkJobFinished( send_reply_callback(status, nullptr, nullptr); }; - Status status = gcs_client_.Jobs().AsyncMarkFinished( - job_id, [reply, send_reply_callback](Status status) { - reply->set_success(status.ok()); - send_reply_callback(status, nullptr, nullptr); - }); + Status status = gcs_client_.Jobs().AsyncMarkFinished(job_id, on_done); if (!status.ok()) { on_done(status); } diff --git a/src/ray/gcs/gcs_server/node_info_handler_impl.cc b/src/ray/gcs/gcs_server/node_info_handler_impl.cc new file mode 100644 index 000000000..257ce40e4 --- /dev/null +++ b/src/ray/gcs/gcs_server/node_info_handler_impl.cc @@ -0,0 +1,65 @@ +#include "node_info_handler_impl.h" +#include "ray/util/logging.h" + +namespace ray { +namespace rpc { + +void DefaultNodeInfoHandler::HandleRegisterNode( + const rpc::RegisterNodeRequest &request, rpc::RegisterNodeReply *reply, + rpc::SendReplyCallback send_reply_callback) { + ClientID node_id = ClientID::FromBinary(request.node_info().node_id()); + RAY_LOG(DEBUG) << "Registering node info, node id = " << node_id; + Status status = gcs_client_.Nodes().Register(request.node_info()); + if (!status.ok()) { + RAY_LOG(DEBUG) << "Failed to register node info, node id = " << node_id; + } + send_reply_callback(status, nullptr, nullptr); + RAY_LOG(DEBUG) << "Finished registering node info, node id = " << node_id; +} + +void DefaultNodeInfoHandler::HandleUnregisterNode( + const rpc::UnregisterNodeRequest &request, rpc::UnregisterNodeReply *reply, + rpc::SendReplyCallback send_reply_callback) { + ClientID node_id = ClientID::FromBinary(request.node_id()); + RAY_LOG(DEBUG) << "Unregistering node info, node id = " << node_id; + + auto on_done = [node_id, send_reply_callback](Status status) { + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to unregister node info: " << status.ToString() + << ", node id = " << node_id; + } + send_reply_callback(status, nullptr, nullptr); + }; + + Status status = gcs_client_.Nodes().AsyncUnregister(node_id, on_done); + if (!status.ok()) { + on_done(status); + } + RAY_LOG(DEBUG) << "Finished unregistering node info, node id = " << node_id; +} + +void DefaultNodeInfoHandler::HandleGetAllNodeInfo( + const rpc::GetAllNodeInfoRequest &request, rpc::GetAllNodeInfoReply *reply, + rpc::SendReplyCallback send_reply_callback) { + RAY_LOG(DEBUG) << "Getting all nodes info."; + auto on_done = [reply, send_reply_callback]( + Status status, const std::vector &result) { + if (status.ok()) { + for (const rpc::GcsNodeInfo &node_info : result) { + reply->add_node_info_list()->CopyFrom(node_info); + } + } else { + RAY_LOG(ERROR) << "Failed to get all nodes info: " << status.ToString(); + } + send_reply_callback(status, nullptr, nullptr); + }; + + Status status = gcs_client_.Nodes().AsyncGetAll(on_done); + if (!status.ok()) { + on_done(status, std::vector()); + } + RAY_LOG(DEBUG) << "Finished getting all node info."; +} + +} // 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 new file mode 100644 index 000000000..f1165185e --- /dev/null +++ b/src/ray/gcs/gcs_server/node_info_handler_impl.h @@ -0,0 +1,34 @@ +#ifndef RAY_GCS_NODE_INFO_HANDLER_IMPL_H +#define RAY_GCS_NODE_INFO_HANDLER_IMPL_H + +#include "ray/gcs/redis_gcs_client.h" +#include "ray/rpc/gcs_server/gcs_rpc_server.h" + +namespace ray { +namespace rpc { + +/// This implementation class of `NodeInfoHandler`. +class DefaultNodeInfoHandler : public rpc::NodeInfoHandler { + public: + explicit DefaultNodeInfoHandler(gcs::RedisGcsClient &gcs_client) + : gcs_client_(gcs_client) {} + + void HandleRegisterNode(const RegisterNodeRequest &request, RegisterNodeReply *reply, + SendReplyCallback send_reply_callback) override; + + void HandleUnregisterNode(const UnregisterNodeRequest &request, + UnregisterNodeReply *reply, + SendReplyCallback send_reply_callback) override; + + void HandleGetAllNodeInfo(const GetAllNodeInfoRequest &request, + GetAllNodeInfoReply *reply, + SendReplyCallback send_reply_callback) override; + + private: + gcs::RedisGcsClient &gcs_client_; +}; + +} // namespace rpc +} // namespace ray + +#endif // RAY_GCS_NODE_INFO_HANDLER_IMPL_H 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 253bf1cb4..e07ec475b 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 @@ -13,8 +13,6 @@ static std::string libray_redis_module_path; class GcsServerTest : public RedisServiceManagerForTest { public: - using CallFunction = std::function &promise)>; - void SetUp() override { gcs::GcsServerConfig config; config.grpc_server_port = 0; @@ -51,74 +49,104 @@ class GcsServerTest : public RedisServiceManagerForTest { thread_gcs_server_->join(); } - void TestAddJob(const rpc::AddJobRequest &request) { - auto call_function = [this, request](std::promise &promise) { - client_->AddJob(request, - [&promise](const Status &status, const rpc::AddJobReply &reply) { - RAY_CHECK_OK(status); - promise.set_value(true); - }); - }; - AsyncCall(call_function, timeout_ms_); + bool AddJob(const rpc::AddJobRequest &request) { + std::promise promise; + client_->AddJob(request, + [&promise](const Status &status, const rpc::AddJobReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + return WaitReady(promise.get_future(), timeout_ms_); } - void TestMarkJobFinished(const rpc::MarkJobFinishedRequest &request) { - auto call_function = [this, request](std::promise &promise) { - client_->MarkJobFinished( - request, - [&promise](const Status &status, const rpc::MarkJobFinishedReply &reply) { - RAY_CHECK_OK(status); - promise.set_value(true); - }); - }; - AsyncCall(call_function, timeout_ms_); + bool MarkJobFinished(const rpc::MarkJobFinishedRequest &request) { + std::promise promise; + client_->MarkJobFinished(request, [&promise](const Status &status, + const rpc::MarkJobFinishedReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + return WaitReady(promise.get_future(), timeout_ms_); } - void TestRegisterActorInfo(const rpc::RegisterActorInfoRequest &request) { - auto call_function = [this, request](std::promise &promise) { - client_->RegisterActorInfo( - request, - [&promise](const Status &status, const rpc::RegisterActorInfoReply &reply) { - RAY_CHECK_OK(status); - promise.set_value(true); - }); - }; - AsyncCall(call_function, timeout_ms_); + bool RegisterActorInfo(const rpc::RegisterActorInfoRequest &request) { + std::promise promise; + client_->RegisterActorInfo( + request, + [&promise](const Status &status, const rpc::RegisterActorInfoReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + return WaitReady(promise.get_future(), timeout_ms_); } - void TestUpdateActorInfo(const rpc::UpdateActorInfoRequest &request) { - auto call_function = [this, request](std::promise &promise) { - client_->UpdateActorInfo( - request, - [&promise](const Status &status, const rpc::UpdateActorInfoReply &reply) { - RAY_CHECK_OK(status); - promise.set_value(true); - }); - }; - AsyncCall(call_function, timeout_ms_); + bool UpdateActorInfo(const rpc::UpdateActorInfoRequest &request) { + std::promise promise; + client_->UpdateActorInfo(request, [&promise](const Status &status, + const rpc::UpdateActorInfoReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + return WaitReady(promise.get_future(), timeout_ms_); } - void TestGetActorInfo(const rpc::ActorTableData &expected) { + rpc::ActorTableData GetActorInfo(const std::string &actor_id) { rpc::GetActorInfoRequest request; - request.set_actor_id(expected.actor_id()); - auto call_function = [this, request, expected](std::promise &promise) { - client_->GetActorInfo( - request, [&promise, &expected](const Status &status, - const rpc::GetActorInfoReply &reply) { - RAY_CHECK_OK(status); - promise.set_value(true); - ASSERT_TRUE(reply.actor_table_data().state() == expected.state()); - }); - }; - AsyncCall(call_function, timeout_ms_); + request.set_actor_id(actor_id); + rpc::ActorTableData actor_table_data; + std::promise promise; + client_->GetActorInfo( + request, [&actor_table_data, &promise](const Status &status, + const rpc::GetActorInfoReply &reply) { + RAY_CHECK_OK(status); + actor_table_data.CopyFrom(reply.actor_table_data()); + promise.set_value(true); + }); + EXPECT_TRUE(WaitReady(promise.get_future(), timeout_ms_)); + return actor_table_data; } - void AsyncCall(const CallFunction &function, uint64_t timeout_ms) { - std::promise promise_; - auto future = promise_.get_future(); - function(promise_); + bool RegisterNode(const rpc::RegisterNodeRequest &request) { + std::promise promise; + client_->RegisterNode( + request, [&promise](const Status &status, const rpc::RegisterNodeReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + + return WaitReady(promise.get_future(), timeout_ms_); + } + + bool UnregisterNode(const rpc::UnregisterNodeRequest &request) { + std::promise promise; + client_->UnregisterNode( + request, [&promise](const Status &status, const rpc::UnregisterNodeReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + return WaitReady(promise.get_future(), timeout_ms_); + } + + std::vector GetAllNodeInfo() { + std::vector node_info_list; + rpc::GetAllNodeInfoRequest request; + std::promise promise; + client_->GetAllNodeInfo( + request, [&node_info_list, &promise](const Status &status, + const rpc::GetAllNodeInfoReply &reply) { + RAY_CHECK_OK(status); + for (int index = 0; index < reply.node_info_list_size(); ++index) { + node_info_list.push_back(reply.node_info_list(index)); + } + promise.set_value(true); + }); + EXPECT_TRUE(WaitReady(promise.get_future(), timeout_ms_)); + return node_info_list; + } + + bool WaitReady(const std::future &future, uint64_t timeout_ms) { auto status = future.wait_for(std::chrono::milliseconds(timeout_ms)); - ASSERT_EQ(status, std::future_status::ready); + return status == std::future_status::ready; } rpc::JobTableData GenJobTableData(JobID job_id) { @@ -143,6 +171,13 @@ class GcsServerTest : public RedisServiceManagerForTest { return actor_table_data; } + rpc::GcsNodeInfo GenGcsNodeInfo(const std::string &node_id) { + rpc::GcsNodeInfo gcs_node_info; + gcs_node_info.set_node_id(node_id); + gcs_node_info.set_state(rpc::GcsNodeInfo_GcsNodeState_ALIVE); + return gcs_node_info; + } + protected: // Gcs server std::unique_ptr gcs_server_; @@ -166,8 +201,10 @@ TEST_F(GcsServerTest, TestActorInfo) { // Register actor rpc::RegisterActorInfoRequest register_actor_info_request; register_actor_info_request.mutable_actor_table_data()->CopyFrom(actor_table_data); - TestRegisterActorInfo(register_actor_info_request); - TestGetActorInfo(actor_table_data); + ASSERT_TRUE(RegisterActorInfo(register_actor_info_request)); + rpc::ActorTableData result = GetActorInfo(actor_table_data.actor_id()); + ASSERT_TRUE(result.state() == + rpc::ActorTableData_ActorState::ActorTableData_ActorState_ALIVE); // Update actor state rpc::UpdateActorInfoRequest update_actor_info_request; @@ -175,8 +212,10 @@ TEST_F(GcsServerTest, TestActorInfo) { rpc::ActorTableData_ActorState::ActorTableData_ActorState_DEAD); update_actor_info_request.set_actor_id(actor_table_data.actor_id()); update_actor_info_request.mutable_actor_table_data()->CopyFrom(actor_table_data); - TestUpdateActorInfo(update_actor_info_request); - TestGetActorInfo(actor_table_data); + ASSERT_TRUE(UpdateActorInfo(update_actor_info_request)); + result = GetActorInfo(actor_table_data.actor_id()); + ASSERT_TRUE(result.state() == + rpc::ActorTableData_ActorState::ActorTableData_ActorState_DEAD); } TEST_F(GcsServerTest, TestJobInfo) { @@ -187,12 +226,36 @@ TEST_F(GcsServerTest, TestJobInfo) { // Add job rpc::AddJobRequest add_job_request; add_job_request.mutable_data()->CopyFrom(job_table_data); - TestAddJob(add_job_request); + ASSERT_TRUE(AddJob(add_job_request)); // Mark job finished rpc::MarkJobFinishedRequest mark_job_finished_request; mark_job_finished_request.set_job_id(job_table_data.job_id()); - TestMarkJobFinished(mark_job_finished_request); + ASSERT_TRUE(MarkJobFinished(mark_job_finished_request)); +} + +TEST_F(GcsServerTest, TestNodeInfo) { + // Create gcs node info + ClientID node_id = ClientID::FromRandom(); + rpc::GcsNodeInfo gcs_node_info = GenGcsNodeInfo(node_id.Binary()); + + // Register node info + rpc::RegisterNodeRequest register_node_info_request; + register_node_info_request.mutable_node_info()->CopyFrom(gcs_node_info); + ASSERT_TRUE(RegisterNode(register_node_info_request)); + std::vector node_info_list = GetAllNodeInfo(); + ASSERT_TRUE(node_info_list.size() == 1); + ASSERT_TRUE(node_info_list[0].state() == + rpc::GcsNodeInfo_GcsNodeState::GcsNodeInfo_GcsNodeState_ALIVE); + + // Unregister node info + rpc::UnregisterNodeRequest unregister_node_info_request; + unregister_node_info_request.set_node_id(node_id.Binary()); + ASSERT_TRUE(UnregisterNode(unregister_node_info_request)); + node_info_list = GetAllNodeInfo(); + ASSERT_TRUE(node_info_list.size() == 1); + ASSERT_TRUE(node_info_list[0].state() == + rpc::GcsNodeInfo_GcsNodeState::GcsNodeInfo_GcsNodeState_DEAD); } } // namespace ray diff --git a/src/ray/gcs/redis_accessor.cc b/src/ray/gcs/redis_accessor.cc index 249885bcb..40ac2b0d5 100644 --- a/src/ray/gcs/redis_accessor.cc +++ b/src/ray/gcs/redis_accessor.cc @@ -251,6 +251,11 @@ const GcsNodeInfo &RedisNodeInfoAccessor::GetSelfInfo() const { return client_table.GetLocalClient(); } +Status RedisNodeInfoAccessor::Register(const GcsNodeInfo &node_info) { + ClientTable &client_table = client_impl_->client_table(); + return client_table.Register(node_info); +} + Status RedisNodeInfoAccessor::AsyncUnregister(const ClientID &node_id, const StatusCallback &callback) { ClientTable::WriteCallback on_done = nullptr; @@ -275,7 +280,14 @@ Status RedisNodeInfoAccessor::AsyncGetAll( RAY_CHECK(callback != nullptr); auto on_done = [callback](RedisGcsClient *client, const ClientID &id, const std::vector &data) { - callback(Status::OK(), data); + std::vector result; + std::set node_ids; + for (int index = data.size() - 1; index >= 0; --index) { + if (node_ids.insert(data[index].node_id()).second) { + result.emplace_back(data[index]); + } + } + callback(Status::OK(), result); }; ClientTable &client_table = client_impl_->client_table(); return client_table.Lookup(on_done); diff --git a/src/ray/gcs/redis_accessor.h b/src/ray/gcs/redis_accessor.h index 2b637b3a7..071add72a 100644 --- a/src/ray/gcs/redis_accessor.h +++ b/src/ray/gcs/redis_accessor.h @@ -149,6 +149,8 @@ class RedisNodeInfoAccessor : public NodeInfoAccessor { const GcsNodeInfo &GetSelfInfo() const override; + Status Register(const GcsNodeInfo &node_info) override; + Status AsyncUnregister(const ClientID &node_id, const StatusCallback &callback) override; diff --git a/src/ray/gcs/tables.cc b/src/ray/gcs/tables.cc index 7dedc55f7..d3ef086bd 100644 --- a/src/ray/gcs/tables.cc +++ b/src/ray/gcs/tables.cc @@ -621,6 +621,12 @@ Status ClientTable::Disconnect() { return status; } +ray::Status ClientTable::Register(const GcsNodeInfo &node_info) { + RAY_CHECK(node_info.state() == GcsNodeInfo::ALIVE); + auto node_info_ptr = std::make_shared(node_info); + return SyncAppend(JobID::Nil(), client_log_key_, node_info_ptr); +} + ray::Status ClientTable::MarkDisconnected(const ClientID &dead_node_id, const WriteCallback &done) { auto node_info = std::make_shared(); diff --git a/src/ray/gcs/tables.h b/src/ray/gcs/tables.h index 3ca5a4352..f1feb11b1 100644 --- a/src/ray/gcs/tables.h +++ b/src/ray/gcs/tables.h @@ -876,6 +876,12 @@ class ClientTable : public Log { /// \return Status ray::Status Disconnect(); + /// Register a new client to the GCS. + /// + /// \param node_info Information about the client. + /// \return Status + ray::Status Register(const GcsNodeInfo &node_info); + /// Mark a different client as disconnected. The client ID should never be /// reused for a new client. /// diff --git a/src/ray/protobuf/gcs_service.proto b/src/ray/protobuf/gcs_service.proto index 3c4936cb7..9fba90b40 100644 --- a/src/ray/protobuf/gcs_service.proto +++ b/src/ray/protobuf/gcs_service.proto @@ -65,3 +65,36 @@ service ActorInfoGcsService { // Update actor info in GCS Service. rpc UpdateActorInfo(UpdateActorInfoRequest) returns (UpdateActorInfoReply); } + +message RegisterNodeRequest { + // Info of node. + GcsNodeInfo node_info = 1; +} + +message RegisterNodeReply { +} + +message UnregisterNodeRequest { + // The ID of node. + bytes node_id = 1; +} + +message UnregisterNodeReply { +} + +message GetAllNodeInfoRequest { +} + +message GetAllNodeInfoReply { + repeated GcsNodeInfo node_info_list = 1; +} + +// Service for node info access. +service NodeInfoGcsService { + // Register a node to GCS Service. + rpc RegisterNode(RegisterNodeRequest) returns (RegisterNodeReply); + // Unregister a node from GCS Service. + rpc UnregisterNode(UnregisterNodeRequest) returns (UnregisterNodeReply); + // Get information of all nodes from GCS Service. + rpc GetAllNodeInfo(GetAllNodeInfoRequest) returns (GetAllNodeInfoReply); +} diff --git a/src/ray/rpc/gcs_server/gcs_rpc_client.h b/src/ray/rpc/gcs_server/gcs_rpc_client.h index cd825ee96..4bbe6cb11 100644 --- a/src/ray/rpc/gcs_server/gcs_rpc_client.h +++ b/src/ray/rpc/gcs_server/gcs_rpc_client.h @@ -26,6 +26,7 @@ class GcsRpcClient { address + ":" + std::to_string(port), grpc::InsecureChannelCredentials()); job_info_stub_ = JobInfoGcsService::NewStub(channel); actor_info_stub_ = ActorInfoGcsService::NewStub(channel); + node_info_stub_ = NodeInfoGcsService::NewStub(channel); }; /// Add job info to gcs server. @@ -85,10 +86,47 @@ class GcsRpcClient { request, callback); } + /// Register a node to GCS Service. + /// + /// \param request The request message. + /// \param callback The callback function that handles reply from server. + void RegisterNode(const RegisterNodeRequest &request, + const ClientCallback &callback) { + client_call_manager_ + .CreateCall( + *node_info_stub_, &NodeInfoGcsService::Stub::PrepareAsyncRegisterNode, + request, callback); + } + + /// Unregister a node from GCS Service. + /// + /// \param request The request message. + /// \param callback The callback function that handles reply from server. + void UnregisterNode(const UnregisterNodeRequest &request, + const ClientCallback &callback) { + client_call_manager_ + .CreateCall( + *node_info_stub_, &NodeInfoGcsService::Stub::PrepareAsyncUnregisterNode, + request, callback); + } + + /// Get information of all nodes from GCS Service. + /// + /// \param request The request message. + /// \param callback The callback function that handles reply from server. + void GetAllNodeInfo(const GetAllNodeInfoRequest &request, + const ClientCallback &callback) { + client_call_manager_ + .CreateCall( + *node_info_stub_, &NodeInfoGcsService::Stub::PrepareAsyncGetAllNodeInfo, + request, callback); + } + private: /// The gRPC-generated stub. std::unique_ptr job_info_stub_; std::unique_ptr actor_info_stub_; + std::unique_ptr node_info_stub_; /// The `ClientCallManager` used for managing requests. ClientCallManager &client_call_manager_; diff --git a/src/ray/rpc/gcs_server/gcs_rpc_server.h b/src/ray/rpc/gcs_server/gcs_rpc_server.h index cff996a91..25c02b2b0 100644 --- a/src/ray/rpc/gcs_server/gcs_rpc_server.h +++ b/src/ray/rpc/gcs_server/gcs_rpc_server.h @@ -27,6 +27,15 @@ namespace rpc { server_call_factories_and_concurrencies->emplace_back( \ std::move(HANDLER##_call_factory), CONCURRENCY); +#define NODE_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \ + std::unique_ptr HANDLER##_call_factory( \ + new ServerCallFactoryImpl( \ + service_, &NodeInfoGcsService::AsyncService::Request##HANDLER, \ + service_handler_, &NodeInfoHandler::Handle##HANDLER, cq, main_service_)); \ + server_call_factories_and_concurrencies->emplace_back( \ + std::move(HANDLER##_call_factory), CONCURRENCY); + class JobInfoHandler { public: virtual ~JobInfoHandler() = default; @@ -113,6 +122,52 @@ class ActorInfoGrpcService : public GrpcService { ActorInfoHandler &service_handler_; }; +class NodeInfoHandler { + public: + virtual ~NodeInfoHandler() = default; + + virtual void HandleRegisterNode(const RegisterNodeRequest &request, + RegisterNodeReply *reply, + SendReplyCallback send_reply_callback) = 0; + + virtual void HandleUnregisterNode(const UnregisterNodeRequest &request, + UnregisterNodeReply *reply, + SendReplyCallback send_reply_callback) = 0; + + virtual void HandleGetAllNodeInfo(const GetAllNodeInfoRequest &request, + GetAllNodeInfoReply *reply, + SendReplyCallback send_reply_callback) = 0; +}; + +/// The `GrpcService` for `NodeInfoGcsService`. +class NodeInfoGrpcService : public GrpcService { + public: + /// Constructor. + /// + /// \param[in] handler The service handler that actually handle the requests. + explicit NodeInfoGrpcService(boost::asio::io_service &io_service, + NodeInfoHandler &handler) + : GrpcService(io_service), service_handler_(handler){}; + + protected: + grpc::Service &GetGrpcService() override { return service_; } + + void InitServerCallFactories( + const std::unique_ptr &cq, + std::vector, int>> + *server_call_factories_and_concurrencies) override { + NODE_INFO_SERVICE_RPC_HANDLER(RegisterNode, 1); + NODE_INFO_SERVICE_RPC_HANDLER(UnregisterNode, 1); + NODE_INFO_SERVICE_RPC_HANDLER(GetAllNodeInfo, 1); + } + + private: + /// The grpc async service object. + NodeInfoGcsService::AsyncService service_; + /// The service handler that actually handle the requests. + NodeInfoHandler &service_handler_; +}; + } // namespace rpc } // namespace ray