Add node resources methods to gcs server node info handler (#6685)

src/ray/gcs/gcs_server/node_info_handler_impl.cc

@@ -110,5 +110,85 @@ void DefaultNodeInfoHandler::HandleReportBatchHeartbeat(
                  << request.heartbeat_batch().batch_size();
 }
 
+void DefaultNodeInfoHandler::HandleGetResources(const GetResourcesRequest &request,
+                                                GetResourcesReply *reply,
+                                                SendReplyCallback send_reply_callback) {
+  ClientID node_id = ClientID::FromBinary(request.node_id());
+  RAY_LOG(DEBUG) << "Getting node resources, node id = " << node_id;
+
+  auto on_done = [node_id, reply, send_reply_callback](
+                     Status status,
+                     const boost::optional<gcs::NodeInfoAccessor::ResourceMap> &result) {
+    if (status.ok()) {
+      if (result) {
+        for (auto &resource : *result) {
+          (*reply->mutable_resources())[resource.first] = *resource.second;
+        }
+      }
+    } else {
+      RAY_LOG(ERROR) << "Failed to get node resources: " << status.ToString()
+                     << ", node id = " << node_id;
+    }
+    send_reply_callback(status, nullptr, nullptr);
+  };
+
+  Status status = gcs_client_.Nodes().AsyncGetResources(node_id, on_done);
+  if (!status.ok()) {
+    on_done(status, boost::none);
+  }
+
+  RAY_LOG(DEBUG) << "Finished getting node resources, node id = " << node_id;
+}
+
+void DefaultNodeInfoHandler::HandleUpdateResources(
+    const UpdateResourcesRequest &request, UpdateResourcesReply *reply,
+    SendReplyCallback send_reply_callback) {
+  ClientID node_id = ClientID::FromBinary(request.node_id());
+  gcs::NodeInfoAccessor::ResourceMap resources;
+  for (auto resource : request.resources()) {
+    resources[resource.first] = std::make_shared<rpc::ResourceTableData>(resource.second);
+  }
+
+  RAY_LOG(DEBUG) << "Updating node resources, node id = " << node_id;
+  auto on_done = [node_id, send_reply_callback](Status status) {
+    if (!status.ok()) {
+      RAY_LOG(ERROR) << "Failed to update node resources: " << status.ToString()
+                     << ", node id = " << node_id;
+    }
+    send_reply_callback(status, nullptr, nullptr);
+  };
+
+  Status status = gcs_client_.Nodes().AsyncUpdateResources(node_id, resources, on_done);
+  if (!status.ok()) {
+    on_done(status);
+  }
+
+  RAY_LOG(DEBUG) << "Finished updating node resources, node id = " << node_id;
+}
+
+void DefaultNodeInfoHandler::HandleDeleteResources(
+    const DeleteResourcesRequest &request, DeleteResourcesReply *reply,
+    SendReplyCallback send_reply_callback) {
+  ClientID node_id = ClientID::FromBinary(request.node_id());
+  auto resource_names = VectorFromProtobuf(request.resource_name_list());
+  RAY_LOG(DEBUG) << "Deleting node resources, node id = " << node_id;
+
+  auto on_done = [node_id, send_reply_callback](Status status) {
+    if (!status.ok()) {
+      RAY_LOG(ERROR) << "Failed to delete node resources: " << status.ToString()
+                     << ", node id = " << node_id;
+    }
+    send_reply_callback(status, nullptr, nullptr);
+  };
+
+  Status status =
+      gcs_client_.Nodes().AsyncDeleteResources(node_id, resource_names, on_done);
+  if (!status.ok()) {
+    on_done(status);
+  }
+
+  RAY_LOG(DEBUG) << "Finished deleting node resources, node id = " << node_id;
+}
+
 }  // namespace rpc
 }  // namespace ray

src/ray/gcs/gcs_server/node_info_handler_impl.h

@@ -32,6 +32,17 @@ class DefaultNodeInfoHandler : public rpc::NodeInfoHandler {
                                   ReportBatchHeartbeatReply *reply,
                                   SendReplyCallback send_reply_callback) override;
 
+  void HandleGetResources(const GetResourcesRequest &request, GetResourcesReply *reply,
+                          SendReplyCallback send_reply_callback) override;
+
+  void HandleUpdateResources(const UpdateResourcesRequest &request,
+                             UpdateResourcesReply *reply,
+                             SendReplyCallback send_reply_callback) override;
+
+  void HandleDeleteResources(const DeleteResourcesRequest &request,
+                             DeleteResourcesReply *reply,
+                             SendReplyCallback send_reply_callback) override;
+
  private:
   gcs::RedisGcsClient &gcs_client_;
 };

src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc

@@ -208,6 +208,44 @@ class GcsServerTest : public RedisServiceManagerForTest {
     return WaitReady(promise.get_future(), timeout_ms_);
   }
 
+  bool UpdateResources(const rpc::UpdateResourcesRequest &request) {
+    std::promise<bool> promise;
+    client_->UpdateResources(request, [&promise](const Status &status,
+                                                 const rpc::UpdateResourcesReply &reply) {
+      RAY_CHECK_OK(status);
+      promise.set_value(true);
+    });
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
+  bool DeleteResources(const rpc::DeleteResourcesRequest &request) {
+    std::promise<bool> promise;
+    client_->DeleteResources(request, [&promise](const Status &status,
+                                                 const rpc::DeleteResourcesReply &reply) {
+      RAY_CHECK_OK(status);
+      promise.set_value(true);
+    });
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
+  std::map<std::string, gcs::ResourceTableData> GetResources(const std::string &node_id) {
+    rpc::GetResourcesRequest request;
+    request.set_node_id(node_id);
+    std::map<std::string, gcs::ResourceTableData> resources;
+    std::promise<bool> promise;
+    client_->GetResources(request,
+                          [&resources, &promise](const Status &status,
+                                                 const rpc::GetResourcesReply &reply) {
+                            RAY_CHECK_OK(status);
+                            for (auto resource : reply.resources()) {
+                              resources[resource.first] = resource.second;
+                            }
+                            promise.set_value(true);
+                          });
+    EXPECT_TRUE(WaitReady(promise.get_future(), timeout_ms_));
+    return resources;
+  }
+
   bool AddObjectLocation(const rpc::AddObjectLocationRequest &request) {
     std::promise<bool> promise;
     client_->AddObjectLocation(
@@ -390,6 +428,25 @@ TEST_F(GcsServerTest, TestNodeInfo) {
   ASSERT_TRUE(node_info_list.size() == 1);
   ASSERT_TRUE(node_info_list[0].state() ==
               rpc::GcsNodeInfo_GcsNodeState::GcsNodeInfo_GcsNodeState_DEAD);
+
+  // Update node resources
+  rpc::UpdateResourcesRequest update_resources_request;
+  update_resources_request.set_node_id(node_id.Binary());
+  rpc::ResourceTableData resource_table_data;
+  resource_table_data.set_resource_capacity(1.0);
+  std::string resource_name = "CPU";
+  (*update_resources_request.mutable_resources())[resource_name] = resource_table_data;
+  ASSERT_TRUE(UpdateResources(update_resources_request));
+  auto resources = GetResources(node_id.Binary());
+  ASSERT_TRUE(resources.size() == 1);
+
+  // Delete node resources
+  rpc::DeleteResourcesRequest delete_resources_request;
+  delete_resources_request.set_node_id(node_id.Binary());
+  delete_resources_request.add_resource_name_list(resource_name);
+  ASSERT_TRUE(DeleteResources(delete_resources_request));
+  resources = GetResources(node_id.Binary());
+  ASSERT_TRUE(resources.size() == 0);
 }
 
 TEST_F(GcsServerTest, TestObjectInfo) {

src/ray/protobuf/gcs_service.proto

@@ -133,6 +133,30 @@ message ReportBatchHeartbeatRequest {
 message ReportBatchHeartbeatReply {
 }
 
+message GetResourcesRequest {
+  bytes node_id = 1;
+}
+
+message GetResourcesReply {
+  map<string, ResourceTableData> resources = 1;
+}
+
+message UpdateResourcesRequest {
+  bytes node_id = 1;
+  map<string, ResourceTableData> resources = 2;
+}
+
+message UpdateResourcesReply {
+}
+
+message DeleteResourcesRequest {
+  bytes node_id = 1;
+  repeated string resource_name_list = 2;
+}
+
+message DeleteResourcesReply {
+}
+
 // Service for node info access.
 service NodeInfoGcsService {
   // Register a node to GCS Service.
@@ -146,6 +170,12 @@ service NodeInfoGcsService {
   // Report batch heartbeat to GCS Service.
   rpc ReportBatchHeartbeat(ReportBatchHeartbeatRequest)
       returns (ReportBatchHeartbeatReply);
+  // Get node's resources from GCS Service.
+  rpc GetResources(GetResourcesRequest) returns (GetResourcesReply);
+  // Update resources of a node in GCS Service.
+  rpc UpdateResources(UpdateResourcesRequest) returns (UpdateResourcesReply);
+  // Delete resources of a node in GCS Service.
+  rpc DeleteResources(DeleteResourcesRequest) returns (DeleteResourcesReply);
 }
 
 message GetObjectLocationsRequest {

src/ray/rpc/gcs_server/gcs_rpc_client.h

@@ -84,6 +84,18 @@ class GcsRpcClient {
   VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, ReportBatchHeartbeat, request, callback,
                          node_info_grpc_client_)
 
+  /// Get node's resources from GCS Service.
+  VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, GetResources, request, callback,
+                         node_info_grpc_client_)
+
+  /// Update resources of a node in GCS Service.
+  VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, UpdateResources, request, callback,
+                         node_info_grpc_client_)
+
+  /// Delete resources of a node in GCS Service.
+  VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, DeleteResources, request, callback,
+                         node_info_grpc_client_)
+
   /// Get object's locations from GCS Service.
   VOID_RPC_CLIENT_METHOD(ObjectInfoGcsService, GetObjectLocations, request, callback,
                          object_info_grpc_client_)

src/ray/rpc/gcs_server/gcs_rpc_server.h

@@ -169,6 +169,18 @@ class NodeInfoHandler {
   virtual void HandleReportBatchHeartbeat(const ReportBatchHeartbeatRequest &request,
                                           ReportBatchHeartbeatReply *reply,
                                           SendReplyCallback send_reply_callback) = 0;
+
+  virtual void HandleGetResources(const GetResourcesRequest &request,
+                                  GetResourcesReply *reply,
+                                  SendReplyCallback send_reply_callback) = 0;
+
+  virtual void HandleUpdateResources(const UpdateResourcesRequest &request,
+                                     UpdateResourcesReply *reply,
+                                     SendReplyCallback send_reply_callback) = 0;
+
+  virtual void HandleDeleteResources(const DeleteResourcesRequest &request,
+                                     DeleteResourcesReply *reply,
+                                     SendReplyCallback send_reply_callback) = 0;
 };
 
 /// The `GrpcService` for `NodeInfoGcsService`.
@@ -193,6 +205,9 @@ class NodeInfoGrpcService : public GrpcService {
     NODE_INFO_SERVICE_RPC_HANDLER(GetAllNodeInfo, 1);
     NODE_INFO_SERVICE_RPC_HANDLER(ReportHeartbeat, 1);
     NODE_INFO_SERVICE_RPC_HANDLER(ReportBatchHeartbeat, 1);
+    NODE_INFO_SERVICE_RPC_HANDLER(GetResources, 1);
+    NODE_INFO_SERVICE_RPC_HANDLER(UpdateResources, 1);
+    NODE_INFO_SERVICE_RPC_HANDLER(DeleteResources, 1);
   }
 
  private: