[GCS]Use new getting all available resources interface instead of pub-sub … (#10914)

* Use new all available resources getting interface instead of pub-sub in state.py

* add missing server handler and test cases, fix comments

* add fine grained test assert

* per comments

* involve new added function _available_resources_per_node

* change  ClientID to NodeID

* fix compile

* fix client id and lint

* robust tests check

* robust tests

src/ray/gcs/accessor.h

@@ -538,6 +538,13 @@ class NodeInfoAccessor {
   virtual Status AsyncGetResources(const NodeID &node_id,
                                    const OptionalItemCallback<ResourceMap> &callback) = 0;
 
+  /// Get available resources of all nodes from GCS asynchronously.
+  ///
+  /// \param callback Callback that will be called after lookup finishes.
+  /// \return Status
+  virtual Status AsyncGetAllAvailableResources(
+      const MultiItemCallback<rpc::AvailableResources> &callback) = 0;
+
   /// Update resources of node in GCS asynchronously.
   ///
   /// \param node_id The ID of node to update dynamic resources.

src/ray/gcs/gcs_client/global_state_accessor.cc

@@ -148,6 +148,16 @@ std::string GlobalStateAccessor::GetNodeResourceInfo(const NodeID &node_id) {
   return node_resource_map.SerializeAsString();
 }
 
+std::vector<std::string> GlobalStateAccessor::GetAllAvailableResources() {
+  std::vector<std::string> available_resources;
+  std::promise<bool> promise;
+  RAY_CHECK_OK(gcs_client_->Nodes().AsyncGetAllAvailableResources(
+      TransformForMultiItemCallback<rpc::AvailableResources>(available_resources,
+                                                             promise)));
+  promise.get_future().get();
+  return available_resources;
+}
+
 std::string GlobalStateAccessor::GetInternalConfig() {
   rpc::StoredConfig config_proto;
   std::promise<void> promise;

src/ray/gcs/gcs_client/global_state_accessor.h

@@ -86,6 +86,13 @@ class GlobalStateAccessor {
   /// deserialized with protobuf function.
   std::string GetNodeResourceInfo(const NodeID &node_id);
 
+  /// Get available resources of all nodes.
+  ///
+  /// \return available resources of all nodes. To support multi-language, we serialize
+  /// each AvailableResources and return the serialized string. Where used, it needs to be
+  /// deserialized with protobuf function.
+  std::vector<std::string> GetAllAvailableResources();
+
   /// Get internal config from GCS Service.
   ///
   /// \return map of internal config keys and values. It is stored as a StoredConfig proto

src/ray/gcs/gcs_client/service_based_accessor.cc

@@ -638,6 +638,21 @@ Status ServiceBasedNodeInfoAccessor::AsyncGetResources(
   return Status::OK();
 }
 
+Status ServiceBasedNodeInfoAccessor::AsyncGetAllAvailableResources(
+    const MultiItemCallback<rpc::AvailableResources> &callback) {
+  rpc::GetAllAvailableResourcesRequest request;
+  client_impl_->GetGcsRpcClient().GetAllAvailableResources(
+      request,
+      [callback](const Status &status, const rpc::GetAllAvailableResourcesReply &reply) {
+        std::vector<rpc::AvailableResources> result =
+            VectorFromProtobuf(reply.resources_list());
+        callback(status, result);
+        RAY_LOG(DEBUG) << "Finished getting available resources of all nodes, status = "
+                       << status;
+      });
+  return Status::OK();
+}
+
 Status ServiceBasedNodeInfoAccessor::AsyncUpdateResources(
     const NodeID &node_id, const ResourceMap &resources, const StatusCallback &callback) {
   RAY_LOG(DEBUG) << "Updating node resources, node id = " << node_id;

src/ray/gcs/gcs_client/service_based_accessor.h

@@ -182,6 +182,9 @@ class ServiceBasedNodeInfoAccessor : public NodeInfoAccessor {
   Status AsyncGetResources(const NodeID &node_id,
                            const OptionalItemCallback<ResourceMap> &callback) override;
 
+  Status AsyncGetAllAvailableResources(
+      const MultiItemCallback<rpc::AvailableResources> &callback) override;
+
   Status AsyncUpdateResources(const NodeID &node_id, const ResourceMap &resources,
                               const StatusCallback &callback) override;
 

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

@@ -331,6 +331,20 @@ class ServiceBasedGcsClientTest : public ::testing::Test {
     return WaitReady(promise.get_future(), timeout_ms_);
   }
 
+  std::vector<rpc::AvailableResources> GetAllAvailableResources() {
+    std::promise<bool> promise;
+    std::vector<rpc::AvailableResources> resources;
+    RAY_CHECK_OK(gcs_client_->Nodes().AsyncGetAllAvailableResources(
+        [&resources, &promise](Status status,
+                               const std::vector<rpc::AvailableResources> &result) {
+          EXPECT_TRUE(!result.empty());
+          resources.assign(result.begin(), result.end());
+          promise.set_value(status.ok());
+        }));
+    EXPECT_TRUE(WaitReady(promise.get_future(), timeout_ms_));
+    return resources;
+  }
+
   bool SubscribeTask(
       const TaskID &task_id,
       const gcs::SubscribeCallback<TaskID, rpc::TaskTableData> &subscribe) {
@@ -732,6 +746,38 @@ TEST_F(ServiceBasedGcsClientTest, TestNodeHeartbeat) {
   WaitForExpectedCount(heartbeat_batch_count, 1);
 }
 
+TEST_F(ServiceBasedGcsClientTest, TestGetAllAvailableResources) {
+  // Subscribe batched state of all nodes from GCS.
+  std::atomic<int> heartbeat_batch_count(0);
+  auto on_subscribe =
+      [&heartbeat_batch_count](const gcs::HeartbeatBatchTableData &result) {
+        ++heartbeat_batch_count;
+      };
+  ASSERT_TRUE(SubscribeBatchHeartbeat(on_subscribe));
+
+  // Register node.
+  auto node_info = Mocker::GenNodeInfo();
+  RAY_CHECK(RegisterNode(*node_info));
+
+  // Report heartbeat of a node to GCS.
+  NodeID node_id = NodeID::FromBinary(node_info->node_id());
+  auto heartbeat = std::make_shared<rpc::HeartbeatTableData>();
+  heartbeat->set_client_id(node_id.Binary());
+  // Set this flag because GCS won't publish unchanged heartbeat.
+  heartbeat->set_should_global_gc(true);
+  (*heartbeat->mutable_resources_available())["CPU"] = 1.0;
+  (*heartbeat->mutable_resources_available())["GPU"] = 10.0;
+  ASSERT_TRUE(ReportHeartbeat(heartbeat));
+  WaitForExpectedCount(heartbeat_batch_count, 1);
+
+  // Assert get all available resources right.
+  std::vector<rpc::AvailableResources> resources = GetAllAvailableResources();
+  EXPECT_EQ(resources.size(), 1);
+  EXPECT_EQ(resources[0].resources_available_size(), 2);
+  EXPECT_EQ((*resources[0].mutable_resources_available())["CPU"], 1.0);
+  EXPECT_EQ((*resources[0].mutable_resources_available())["GPU"], 10.0);
+}
+
 TEST_F(ServiceBasedGcsClientTest, TestTaskInfo) {
   JobID job_id = JobID::FromInt(1);
   TaskID task_id = TaskID::ForDriverTask(job_id);

src/ray/gcs/gcs_server/gcs_node_manager.cc

@@ -354,6 +354,21 @@ void GcsNodeManager::HandleGetInternalConfig(const rpc::GetInternalConfigRequest
       gcs_table_storage_->InternalConfigTable().Get(UniqueID::Nil(), get_system_config));
 }
 
+void GcsNodeManager::HandleGetAllAvailableResources(
+    const rpc::GetAllAvailableResourcesRequest &request,
+    rpc::GetAllAvailableResourcesReply *reply,
+    rpc::SendReplyCallback send_reply_callback) {
+  for (const auto &iter : GetClusterRealtimeResources()) {
+    rpc::AvailableResources resource;
+    resource.set_node_id(iter.first.Binary());
+    for (auto res : iter.second->GetResourceAmountMap()) {
+      (*resource.mutable_resources_available())[res.first] = res.second.ToDouble();
+    }
+    reply->add_resources_list()->CopyFrom(resource);
+  }
+  GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
+}
+
 std::shared_ptr<rpc::GcsNodeInfo> GcsNodeManager::GetNode(
     const ray::NodeID &node_id) const {
   auto iter = alive_nodes_.find(node_id);

src/ray/gcs/gcs_server/gcs_node_manager.h

@@ -79,16 +79,22 @@ class GcsNodeManager : public rpc::NodeInfoHandler {
                              rpc::DeleteResourcesReply *reply,
                              rpc::SendReplyCallback send_reply_callback) override;
 
-  /// Handle setting internal config.
+  /// Handle set internal config.
   void HandleSetInternalConfig(const rpc::SetInternalConfigRequest &request,
                                rpc::SetInternalConfigReply *reply,
                                rpc::SendReplyCallback send_reply_callback) override;
 
-  /// Handle getting internal config.
+  /// Handle get internal config.
   void HandleGetInternalConfig(const rpc::GetInternalConfigRequest &request,
                                rpc::GetInternalConfigReply *reply,
                                rpc::SendReplyCallback send_reply_callback) override;
 
+  /// Handle get available resources of all nodes.
+  void HandleGetAllAvailableResources(
+      const rpc::GetAllAvailableResourcesRequest &request,
+      rpc::GetAllAvailableResourcesReply *reply,
+      rpc::SendReplyCallback send_reply_callback) override;
+
   /// Add an alive node.
   ///
   /// \param node The info of the node to be added.

src/ray/gcs/redis_accessor.h

@@ -356,6 +356,11 @@ class RedisNodeInfoAccessor : public NodeInfoAccessor {
   Status AsyncGetResources(const NodeID &node_id,
                            const OptionalItemCallback<ResourceMap> &callback) override;
 
+  Status AsyncGetAllAvailableResources(
+      const MultiItemCallback<rpc::AvailableResources> &callback) override {
+    return Status::NotImplemented("AsyncGetAllAvailableResources not implemented");
+  }
+
   Status AsyncUpdateResources(const NodeID &node_id, const ResourceMap &resources,
                               const StatusCallback &callback) override;
 

src/ray/protobuf/gcs.proto

@@ -222,6 +222,13 @@ message ResourceTableData {
   double resource_capacity = 1;
 }
 
+message AvailableResources {
+  // Node manager client id.
+  bytes node_id = 1;
+  // Resource capacity currently available on this node manager.
+  map<string, double> resources_available = 2;
+}
+
 message GcsNodeInfo {
   // State of a node.
   enum GcsNodeState {

src/ray/protobuf/gcs_service.proto

@@ -263,6 +263,14 @@ message GetInternalConfigReply {
   StoredConfig config = 2;
 }
 
+message GetAllAvailableResourcesRequest {
+}
+
+message GetAllAvailableResourcesReply {
+  GcsStatus status = 1;
+  repeated AvailableResources resources_list = 2;
+}
+
 // Service for node info access.
 service NodeInfoGcsService {
   // Register a node to GCS Service.
@@ -283,6 +291,9 @@ service NodeInfoGcsService {
   rpc SetInternalConfig(SetInternalConfigRequest) returns (SetInternalConfigReply);
   // Get cluster internal config.
   rpc GetInternalConfig(GetInternalConfigRequest) returns (GetInternalConfigReply);
+  // Get available resources of all nodes.
+  rpc GetAllAvailableResources(GetAllAvailableResourcesRequest)
+      returns (GetAllAvailableResourcesReply);
 }
 
 message GetObjectLocationsRequest {

src/ray/rpc/gcs_server/gcs_rpc_client.h

@@ -190,6 +190,10 @@ class GcsRpcClient {
   VOID_GCS_RPC_CLIENT_METHOD(NodeInfoGcsService, GetInternalConfig,
                              node_info_grpc_client_, )
 
+  /// Get available resources of all nodes from the GCS Service.
+  VOID_GCS_RPC_CLIENT_METHOD(NodeInfoGcsService, GetAllAvailableResources,
+                             node_info_grpc_client_, )
+
   /// Get object's locations from GCS Service.
   VOID_GCS_RPC_CLIENT_METHOD(ObjectInfoGcsService, GetObjectLocations,
                              object_info_grpc_client_, )

src/ray/rpc/gcs_server/gcs_rpc_server.h

@@ -215,6 +215,11 @@ class NodeInfoGcsServiceHandler {
   virtual void HandleGetInternalConfig(const GetInternalConfigRequest &request,
                                        GetInternalConfigReply *reply,
                                        SendReplyCallback send_reply_callback) = 0;
+
+  virtual void HandleGetAllAvailableResources(
+      const rpc::GetAllAvailableResourcesRequest &request,
+      rpc::GetAllAvailableResourcesReply *reply,
+      rpc::SendReplyCallback send_reply_callback) = 0;
 };
 
 /// The `GrpcService` for `NodeInfoGcsService`.
@@ -242,6 +247,7 @@ class NodeInfoGrpcService : public GrpcService {
     NODE_INFO_SERVICE_RPC_HANDLER(DeleteResources);
     NODE_INFO_SERVICE_RPC_HANDLER(SetInternalConfig);
     NODE_INFO_SERVICE_RPC_HANDLER(GetInternalConfig);
+    NODE_INFO_SERVICE_RPC_HANDLER(GetAllAvailableResources);
   }
 
  private: