Add gcs server node info handler (#6595)

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.

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);
   }

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<rpc::ActorInfoHandler> GcsServer::InitActorInfoHandler() {
       new rpc::DefaultActorInfoHandler(*redis_gcs_client_));
 }
 
+std::unique_ptr<rpc::NodeInfoHandler> GcsServer::InitNodeInfoHandler() {
+  return std::unique_ptr<rpc::DefaultNodeInfoHandler>(
+      new rpc::DefaultNodeInfoHandler(*redis_gcs_client_));
+}
+
 }  // namespace gcs
 }  // namespace ray

src/ray/gcs/gcs_server/gcs_server.h

@@ -52,6 +52,9 @@ class GcsServer {
   /// The actor info handler
   virtual std::unique_ptr<rpc::ActorInfoHandler> InitActorInfoHandler();
 
+  /// The node info handler
+  virtual std::unique_ptr<rpc::NodeInfoHandler> InitNodeInfoHandler();
+
  private:
   /// Gcs server configuration
   GcsServerConfig config_;
@@ -65,6 +68,9 @@ class GcsServer {
   /// Actor info handler and service
   std::unique_ptr<rpc::ActorInfoHandler> actor_info_handler_;
   std::unique_ptr<rpc::ActorInfoGrpcService> actor_info_service_;
+  /// Node info handler and service
+  std::unique_ptr<rpc::NodeInfoHandler> node_info_handler_;
+  std::unique_ptr<rpc::NodeInfoGrpcService> node_info_service_;
   /// Backend client
   std::shared_ptr<RedisGcsClient> redis_gcs_client_;
 };

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);
   }

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<rpc::GcsNodeInfo> &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<rpc::GcsNodeInfo>());
+  }
+  RAY_LOG(DEBUG) << "Finished getting all node info.";
+}
+
+}  // namespace rpc
+}  // namespace ray

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

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<void(std::promise<bool> &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<bool> &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<bool> 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<bool> &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<bool> 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<bool> &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<bool> 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<bool> &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<bool> 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<bool> &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<bool> 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<bool> promise_;
-    auto future = promise_.get_future();
-    function(promise_);
+  bool RegisterNode(const rpc::RegisterNodeRequest &request) {
+    std::promise<bool> 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<bool> 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<rpc::GcsNodeInfo> GetAllNodeInfo() {
+    std::vector<rpc::GcsNodeInfo> node_info_list;
+    rpc::GetAllNodeInfoRequest request;
+    std::promise<bool> 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<bool> &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::GcsServer> 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<rpc::GcsNodeInfo> 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

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<GcsNodeInfo> &data) {
-    callback(Status::OK(), data);
+    std::vector<GcsNodeInfo> result;
+    std::set<std::string> 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);

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;
 

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<GcsNodeInfo>(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<GcsNodeInfo>();

src/ray/gcs/tables.h

@@ -876,6 +876,12 @@ class ClientTable : public Log<ClientID, GcsNodeInfo> {
   /// \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.
   ///

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);
+}

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<RegisterNodeReply> &callback) {
+    client_call_manager_
+        .CreateCall<NodeInfoGcsService, RegisterNodeRequest, RegisterNodeReply>(
+            *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<UnregisterNodeReply> &callback) {
+    client_call_manager_
+        .CreateCall<NodeInfoGcsService, UnregisterNodeRequest, UnregisterNodeReply>(
+            *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<GetAllNodeInfoReply> &callback) {
+    client_call_manager_
+        .CreateCall<NodeInfoGcsService, GetAllNodeInfoRequest, GetAllNodeInfoReply>(
+            *node_info_stub_, &NodeInfoGcsService::Stub::PrepareAsyncGetAllNodeInfo,
+            request, callback);
+  }
+
  private:
   /// The gRPC-generated stub.
   std::unique_ptr<JobInfoGcsService::Stub> job_info_stub_;
   std::unique_ptr<ActorInfoGcsService::Stub> actor_info_stub_;
+  std::unique_ptr<NodeInfoGcsService::Stub> node_info_stub_;
 
   /// The `ClientCallManager` used for managing requests.
   ClientCallManager &client_call_manager_;

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<ServerCallFactory> HANDLER##_call_factory(                           \
+      new ServerCallFactoryImpl<NodeInfoGcsService, NodeInfoHandler, HANDLER##Request, \
+                                HANDLER##Reply>(                                       \
+          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<grpc::ServerCompletionQueue> &cq,
+      std::vector<std::pair<std::unique_ptr<ServerCallFactory>, 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