[GCS] Global state accessor from node resource table (#8658)

src/ray/core_worker/lib/java/io_ray_runtime_gcs_GlobalStateAccessor.cc

@@ -76,6 +76,16 @@ Java_io_ray_runtime_gcs_GlobalStateAccessor_nativeGetAllNodeInfo(JNIEnv *env, jo
       });
 }
 
+JNIEXPORT jbyteArray JNICALL
+Java_io_ray_runtime_gcs_GlobalStateAccessor_nativeGetNodeResourceInfo(
+    JNIEnv *env, jobject o, jlong gcs_accessor_ptr, jbyteArray node_id_bytes) {
+  auto *gcs_accessor =
+      reinterpret_cast<ray::gcs::GlobalStateAccessor *>(gcs_accessor_ptr);
+  auto node_id = JavaByteArrayToId<ray::ClientID>(env, node_id_bytes);
+  auto node_resource_info = gcs_accessor->GetNodeResourceInfo(node_id);
+  return static_cast<jbyteArray>(NativeStringToJavaByteArray(env, node_resource_info));
+}
+
 JNIEXPORT jobject JNICALL
 Java_io_ray_runtime_gcs_GlobalStateAccessor_nativeGetAllActorInfo(
     JNIEnv *env, jobject o, jlong gcs_accessor_ptr) {

src/ray/core_worker/lib/java/io_ray_runtime_gcs_GlobalStateAccessor.h

@@ -75,6 +75,14 @@ JNIEXPORT jobject JNICALL
 Java_io_ray_runtime_gcs_GlobalStateAccessor_nativeGetAllNodeInfo(JNIEnv *, jobject,
                                                                  jlong);
 
+/*
+ * Class:     io_ray_runtime_gcs_GlobalStateAccessor
+ * Method:    nativeGetNodeResourceInfo
+ * Signature: (J[B)[B
+ */
+JNIEXPORT jbyteArray JNICALL
+Java_io_ray_runtime_gcs_GlobalStateAccessor_nativeGetNodeResourceInfo(JNIEnv *, jobject,
+                                                                      jlong, jbyteArray);
 /*
  * Class:     io_ray_runtime_gcs_GlobalStateAccessor
  * Method:    nativeGetAllActorInfo

src/ray/gcs/gcs_client/global_state_accessor.cc

@@ -129,6 +129,27 @@ std::unique_ptr<std::string> GlobalStateAccessor::GetObjectInfo(
   return object_info;
 }
 
+std::string GlobalStateAccessor::GetNodeResourceInfo(const ClientID &node_id) {
+  rpc::ResourceMap node_resource_map;
+  std::promise<void> promise;
+  auto on_done =
+      [&node_resource_map, &promise](
+          const Status &status,
+          const boost::optional<ray::gcs::NodeInfoAccessor::ResourceMap> &result) {
+        RAY_CHECK_OK(status);
+        if (result) {
+          auto result_value = result.get();
+          for (auto &data : result_value) {
+            (*node_resource_map.mutable_items())[data.first] = *data.second;
+          }
+        }
+        promise.set_value();
+      };
+  RAY_CHECK_OK(gcs_client_->Nodes().AsyncGetResources(node_id, on_done));
+  promise.get_future().get();
+  return node_resource_map.SerializeAsString();
+}
+
 std::vector<std::string> GlobalStateAccessor::GetAllActorInfo() {
   std::vector<std::string> actor_table_data;
   std::promise<bool> promise;

src/ray/gcs/gcs_client/global_state_accessor.h

@@ -79,6 +79,14 @@ class GlobalStateAccessor {
   /// protobuf function.
   std::unique_ptr<std::string> GetObjectInfo(const ObjectID &object_id);
 
+  /// Get information of a node resource from GCS Service.
+  ///
+  /// \param node_id The ID of node to look up in the GCS Service.
+  /// \return node resource map info. To support multi-language, we serialize each
+  /// ResourceTableData and return the serialized string. Where used, it needs to be
+  /// deserialized with protobuf function.
+  std::string GetNodeResourceInfo(const ClientID &node_id);
+
   /// Get information of all actors from GCS Service.
   ///
   /// \return All actor info. To support multi-language, we serialize each ActorTableData

src/ray/gcs/gcs_client/test/global_state_accessor_test.cc

@@ -132,6 +132,42 @@ TEST_F(GlobalStateAccessorTest, TestNodeTable) {
   }
 }
 
+TEST_F(GlobalStateAccessorTest, TestNodeResourceTable) {
+  int node_count = 100;
+  ASSERT_EQ(global_state_->GetAllNodeInfo().size(), 0);
+  for (int index = 0; index < node_count; ++index) {
+    auto node_table_data =
+        Mocker::GenNodeInfo(index, std::string("127.0.0.") + std::to_string(index));
+    auto node_id = ClientID::FromBinary(node_table_data->node_id());
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Nodes().AsyncRegister(
+        *node_table_data, [&promise](Status status) { promise.set_value(status.ok()); }));
+    WaitReady(promise.get_future(), timeout_ms_);
+    ray::gcs::NodeInfoAccessor::ResourceMap resources;
+    rpc::ResourceTableData resource_table_data;
+    resource_table_data.set_resource_capacity(static_cast<double>(index + 1) + 0.1);
+    resources[std::to_string(index)] =
+        std::make_shared<rpc::ResourceTableData>(resource_table_data);
+    RAY_IGNORE_EXPR(gcs_client_->Nodes().AsyncUpdateResources(
+        node_id, resources, [](Status status) { RAY_CHECK(status.ok()); }));
+  }
+  auto node_table = global_state_->GetAllNodeInfo();
+  ASSERT_EQ(node_table.size(), node_count);
+  for (int index = 0; index < node_count; ++index) {
+    rpc::GcsNodeInfo node_data;
+    node_data.ParseFromString(node_table[index]);
+    auto resource_map_str =
+        global_state_->GetNodeResourceInfo(ClientID::FromBinary(node_data.node_id()));
+    rpc::ResourceMap resource_map;
+    resource_map.ParseFromString(resource_map_str);
+    ASSERT_EQ(
+        static_cast<uint32_t>(
+            (*resource_map.mutable_items())[std::to_string(node_data.node_manager_port())]
+                .resource_capacity()),
+        node_data.node_manager_port() + 1);
+  }
+}
+
 TEST_F(GlobalStateAccessorTest, TestProfileTable) {
   int profile_count = 100;
   ASSERT_EQ(global_state_->GetAllProfileInfo().size(), 0);

src/ray/gcs/gcs_server/gcs_node_manager.cc

@@ -202,8 +202,8 @@ void GcsNodeManager::HandleGetResources(const rpc::GetResourcesRequest &request,
   RAY_LOG(DEBUG) << "Getting node resources, node id = " << node_id;
   auto iter = cluster_resources_.find(node_id);
   if (iter != cluster_resources_.end()) {
-    for (auto &resource : iter->second) {
-      (*reply->mutable_resources())[resource.first] = *resource.second;
+    for (auto &resource : iter->second.items()) {
+      (*reply->mutable_resources())[resource.first] = resource.second;
     }
   }
   GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
@@ -216,23 +216,19 @@ void GcsNodeManager::HandleUpdateResources(const rpc::UpdateResourcesRequest &re
   ClientID node_id = ClientID::FromBinary(request.node_id());
   RAY_LOG(DEBUG) << "Updating resources, node id = " << node_id;
   auto iter = cluster_resources_.find(node_id);
+  auto to_be_updated_resources = request.resources();
   if (iter != cluster_resources_.end()) {
-    auto to_be_updated_resources = std::make_shared<gcs::NodeInfoAccessor::ResourceMap>();
-    for (auto resource : request.resources()) {
-      (*to_be_updated_resources)[resource.first] =
-          std::make_shared<rpc::ResourceTableData>(resource.second);
-    }
-    for (auto &entry : *to_be_updated_resources) {
-      iter->second[entry.first] = entry.second;
+    for (auto &entry : to_be_updated_resources) {
+      (*iter->second.mutable_items())[entry.first] = entry.second;
     }
     auto on_done = [this, node_id, to_be_updated_resources, reply,
                     send_reply_callback](const Status &status) {
       RAY_CHECK_OK(status);
       rpc::NodeResourceChange node_resource_change;
       node_resource_change.set_node_id(node_id.Binary());
-      for (auto &it : *to_be_updated_resources) {
+      for (auto &it : to_be_updated_resources) {
         (*node_resource_change.mutable_updated_resources())[it.first] =
-            it.second->resource_capacity();
+            it.second.resource_capacity();
       }
       RAY_CHECK_OK(gcs_pub_sub_->Publish(NODE_RESOURCE_CHANNEL, node_id.Hex(),
                                          node_resource_change.SerializeAsString(),
@@ -242,8 +238,8 @@ void GcsNodeManager::HandleUpdateResources(const rpc::UpdateResourcesRequest &re
       RAY_LOG(DEBUG) << "Finished updating resources, node id = " << node_id;
     };
 
-    RAY_CHECK_OK(node_info_accessor_.AsyncUpdateResources(
-        node_id, *to_be_updated_resources, on_done));
+    RAY_CHECK_OK(
+        gcs_table_storage_->NodeResourceTable().Put(node_id, iter->second, on_done));
   } else {
     GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::Invalid("Node is not exist."));
     RAY_LOG(ERROR) << "Failed to update resources as node " << node_id
@@ -260,7 +256,7 @@ void GcsNodeManager::HandleDeleteResources(const rpc::DeleteResourcesRequest &re
   auto iter = cluster_resources_.find(node_id);
   if (iter != cluster_resources_.end()) {
     for (auto &resource_name : resource_names) {
-      iter->second.erase(resource_name);
+      RAY_IGNORE_EXPR(iter->second.mutable_items()->erase(resource_name));
     }
     auto on_done = [this, node_id, resource_names, reply,
                     send_reply_callback](const Status &status) {
@@ -277,7 +273,7 @@ void GcsNodeManager::HandleDeleteResources(const rpc::DeleteResourcesRequest &re
       GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
     };
     RAY_CHECK_OK(
-        node_info_accessor_.AsyncDeleteResources(node_id, resource_names, on_done));
+        gcs_table_storage_->NodeResourceTable().Put(node_id, iter->second, on_done));
   } else {
     GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
     RAY_LOG(DEBUG) << "Finished deleting node resources, node id = " << node_id;
@@ -300,8 +296,7 @@ void GcsNodeManager::AddNode(std::shared_ptr<rpc::GcsNodeInfo> node) {
   if (iter == alive_nodes_.end()) {
     alive_nodes_.emplace(node_id, node);
     // Add an empty resources for this node.
-    RAY_CHECK(
-        cluster_resources_.emplace(node_id, gcs::NodeInfoAccessor::ResourceMap()).second);
+    RAY_CHECK(cluster_resources_.emplace(node_id, rpc::ResourceMap()).second);
     // Register this node to the `node_failure_detector_` which will start monitoring it.
     node_failure_detector_->AddNode(node_id);
     // Notify all listeners.

src/ray/gcs/gcs_server/gcs_node_manager.h

@@ -198,7 +198,7 @@ class GcsNodeManager : public rpc::NodeInfoHandler {
   /// Dead nodes.
   absl::flat_hash_map<ClientID, std::shared_ptr<rpc::GcsNodeInfo>> dead_nodes_;
   /// Cluster resources.
-  absl::flat_hash_map<ClientID, gcs::NodeInfoAccessor::ResourceMap> cluster_resources_;
+  absl::flat_hash_map<ClientID, rpc::ResourceMap> cluster_resources_;
   /// Listeners which monitors the addition of nodes.
   std::vector<std::function<void(std::shared_ptr<rpc::GcsNodeInfo>)>>
       node_added_listeners_;