[GCS]Use direct getting instead of pub-sub to update load metrics in monitor.py (#11339)

src/ray/gcs/accessor.h

@@ -576,6 +576,14 @@ class NodeInfoAccessor {
   /// Resend heartbeat when GCS restarts from a failure.
   virtual void AsyncReReportHeartbeat() = 0;
 
+  /// Get newest heartbeat of all nodes from GCS asynchronously. Only used when light
+  /// heartbeat enabled.
+  ///
+  /// \param callback Callback that will be called after lookup finishes.
+  /// \return Status
+  virtual Status AsyncGetAllHeartbeat(
+      const ItemCallback<rpc::HeartbeatBatchTableData> &callback) = 0;
+
   /// Subscribe batched state of all nodes from GCS.
   ///
   /// \param subscribe Callback that will be called each time when batch heartbeat is

src/ray/gcs/gcs_client/global_state_accessor.cc

@@ -174,6 +174,16 @@ std::string GlobalStateAccessor::GetInternalConfig() {
   return config_proto.SerializeAsString();
 }
 
+std::unique_ptr<std::string> GlobalStateAccessor::GetAllHeartbeat() {
+  std::unique_ptr<std::string> heartbeat_batch_data;
+  std::promise<bool> promise;
+  RAY_CHECK_OK(gcs_client_->Nodes().AsyncGetAllHeartbeat(
+      TransformForItemCallback<rpc::HeartbeatBatchTableData>(heartbeat_batch_data,
+                                                             promise)));
+  promise.get_future().get();
+  return heartbeat_batch_data;
+}
+
 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

@@ -99,6 +99,14 @@ class GlobalStateAccessor {
   /// and serialized as a string to allow multi-language support.
   std::string GetInternalConfig();
 
+  /// Get newest heartbeat of all nodes from GCS Service. Only used when light
+  /// heartbeat enabled.
+  ///
+  /// \return node heartbeat info. To support multi-language, we serialize each
+  /// HeartbeatTableData and return the serialized string. Where used, it needs to be
+  /// deserialized with protobuf function.
+  std::unique_ptr<std::string> GetAllHeartbeat();
+
   /// Get information of all actors from GCS Service.
   ///
   /// \return All actor info. To support multi-language, we serialize each ActorTableData
@@ -190,6 +198,18 @@ class GlobalStateAccessor {
     };
   }
 
+  /// Item transformation helper in template style.
+  ///
+  /// \return ItemCallback within in rpc type DATA.
+  template <class DATA>
+  ItemCallback<DATA> TransformForItemCallback(std::unique_ptr<std::string> &data,
+                                              std::promise<bool> &promise) {
+    return [&data, &promise](const DATA &result) {
+      data.reset(new std::string(result.SerializeAsString()));
+      promise.set_value(true);
+    };
+  }
+
   /// Whether this client is connected to gcs server.
   bool is_connected_{false};
 

src/ray/gcs/gcs_client/service_based_accessor.cc

@@ -709,6 +709,17 @@ void ServiceBasedNodeInfoAccessor::AsyncReReportHeartbeat() {
   }
 }
 
+Status ServiceBasedNodeInfoAccessor::AsyncGetAllHeartbeat(
+    const ItemCallback<rpc::HeartbeatBatchTableData> &callback) {
+  rpc::GetAllHeartbeatRequest request;
+  client_impl_->GetGcsRpcClient().GetAllHeartbeat(
+      request, [callback](const Status &status, const rpc::GetAllHeartbeatReply &reply) {
+        callback(reply.heartbeat_data());
+        RAY_LOG(DEBUG) << "Finished getting heartbeat of all nodes, status = " << status;
+      });
+  return Status::OK();
+}
+
 Status ServiceBasedNodeInfoAccessor::AsyncSubscribeBatchHeartbeat(
     const ItemCallback<rpc::HeartbeatBatchTableData> &subscribe,
     const StatusCallback &done) {

src/ray/gcs/gcs_client/service_based_accessor.h

@@ -194,6 +194,9 @@ class ServiceBasedNodeInfoAccessor : public NodeInfoAccessor {
 
   void AsyncReReportHeartbeat() override;
 
+  Status AsyncGetAllHeartbeat(
+      const ItemCallback<rpc::HeartbeatBatchTableData> &callback) override;
+
   Status AsyncSubscribeBatchHeartbeat(
       const ItemCallback<rpc::HeartbeatBatchTableData> &subscribe,
       const StatusCallback &done) override;

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

@@ -191,6 +191,34 @@ TEST_F(GlobalStateAccessorTest, TestInternalConfig) {
   }
 }
 
+TEST_F(GlobalStateAccessorTest, TestGetAllHeartbeat) {
+  std::unique_ptr<std::string> heartbeats = global_state_->GetAllHeartbeat();
+  rpc::HeartbeatBatchTableData heartbeat_batch_data;
+  heartbeat_batch_data.ParseFromString(*heartbeats.get());
+
+  ASSERT_EQ(heartbeat_batch_data.batch_size(), 0);
+
+  auto node_table_data = Mocker::GenNodeInfo();
+  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_);
+  auto node_table = global_state_->GetAllNodeInfo();
+  ASSERT_EQ(node_table.size(), 1);
+
+  // Report heartbeat
+  std::promise<bool> promise1;
+  auto heartbeat = std::make_shared<rpc::HeartbeatTableData>();
+  heartbeat->set_client_id(node_table_data->node_id());
+  RAY_CHECK_OK(gcs_client_->Nodes().AsyncReportHeartbeat(
+      heartbeat, [&promise1](Status status) { promise1.set_value(status.ok()); }));
+  WaitReady(promise1.get_future(), timeout_ms_);
+
+  heartbeats = global_state_->GetAllHeartbeat();
+  heartbeat_batch_data.ParseFromString(*heartbeats.get());
+  ASSERT_EQ(heartbeat_batch_data.batch_size(), 1);
+}
+
 TEST_F(GlobalStateAccessorTest, TestProfileTable) {
   int profile_count = RayConfig::instance().maximum_profile_table_rows_count() + 1;
   ASSERT_EQ(global_state_->GetAllProfileInfo().size(), 0);

src/ray/gcs/gcs_server/gcs_node_manager.cc

@@ -193,6 +193,8 @@ void GcsNodeManager::HandleReportHeartbeat(const rpc::ReportHeartbeatRequest &re
   auto heartbeat_data = std::make_shared<rpc::HeartbeatTableData>();
   heartbeat_data->CopyFrom(request.heartbeat());
 
+  UpdateNodeHeartbeat(node_id, request);
+
   // Update node realtime resources.
   UpdateNodeRealtimeResources(node_id, *heartbeat_data);
 
@@ -335,6 +337,77 @@ void GcsNodeManager::HandleGetAllAvailableResources(
   GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
 }
 
+void GcsNodeManager::HandleGetAllHeartbeat(const rpc::GetAllHeartbeatRequest &request,
+                                           rpc::GetAllHeartbeatReply *reply,
+                                           rpc::SendReplyCallback send_reply_callback) {
+  if (!node_heartbeats_.empty()) {
+    auto batch = std::make_shared<rpc::HeartbeatBatchTableData>();
+    absl::flat_hash_map<ResourceSet, rpc::ResourceDemand> aggregate_load;
+    for (auto &heartbeat : node_heartbeats_) {
+      // Aggregate the load reported by each raylet.
+      auto load = heartbeat.second.resource_load_by_shape();
+      for (const auto &demand : load.resource_demands()) {
+        auto scheduling_key = ResourceSet(MapFromProtobuf(demand.shape()));
+        auto &aggregate_demand = aggregate_load[scheduling_key];
+        aggregate_demand.set_num_ready_requests_queued(
+            aggregate_demand.num_ready_requests_queued() +
+            demand.num_ready_requests_queued());
+        aggregate_demand.set_num_infeasible_requests_queued(
+            aggregate_demand.num_infeasible_requests_queued() +
+            demand.num_infeasible_requests_queued());
+        if (RayConfig::instance().report_worker_backlog()) {
+          aggregate_demand.set_backlog_size(aggregate_demand.backlog_size() +
+                                            demand.backlog_size());
+        }
+      }
+      heartbeat.second.clear_resource_load_by_shape();
+
+      batch->add_batch()->Swap(&heartbeat.second);
+    }
+
+    for (auto &demand : aggregate_load) {
+      auto demand_proto = batch->mutable_resource_load_by_shape()->add_resource_demands();
+      demand_proto->Swap(&demand.second);
+      for (const auto &resource_pair : demand.first.GetResourceMap()) {
+        (*demand_proto->mutable_shape())[resource_pair.first] = resource_pair.second;
+      }
+    }
+
+    // Update placement group load to heartbeat batch.
+    // This is updated only one per second.
+    if (placement_group_load_.has_value()) {
+      auto placement_group_load = placement_group_load_.value();
+      auto placement_group_load_proto = batch->mutable_placement_group_load();
+      placement_group_load_proto->CopyFrom(*placement_group_load.get());
+    }
+    reply->mutable_heartbeat_data()->CopyFrom(*batch);
+  }
+
+  GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
+}
+
+void GcsNodeManager::UpdateNodeHeartbeat(const NodeID node_id,
+                                         const rpc::ReportHeartbeatRequest &request) {
+  auto iter = node_heartbeats_.find(node_id);
+  if (!RayConfig::instance().light_heartbeat_enabled() ||
+      iter == node_heartbeats_.end()) {
+    auto heartbeat_data = std::make_shared<rpc::HeartbeatTableData>();
+    heartbeat_data->CopyFrom(request.heartbeat());
+    node_heartbeats_[node_id] = *heartbeat_data;
+  } else {
+    if (request.heartbeat().resources_total_size() > 0) {
+      (*iter->second.mutable_resources_total()) = request.heartbeat().resources_total();
+    }
+    if (request.heartbeat().resources_available_changed()) {
+      (*iter->second.mutable_resources_available()) =
+          request.heartbeat().resources_available();
+    }
+    if (request.heartbeat().resource_load_changed()) {
+      (*iter->second.mutable_resource_load()) = request.heartbeat().resource_load();
+    }
+  }
+}
+
 absl::optional<std::shared_ptr<rpc::GcsNodeInfo>> GcsNodeManager::GetNode(
     const ray::NodeID &node_id) const {
   auto iter = alive_nodes_.find(node_id);
@@ -482,44 +555,9 @@ void GcsNodeManager::SendBatchedHeartbeat() {
     auto batch = std::make_shared<rpc::HeartbeatBatchTableData>();
     std::unordered_map<ResourceSet, rpc::ResourceDemand> aggregate_load;
     for (auto &heartbeat : heartbeat_buffer_) {
-      // Aggregate the load reported by each raylet.
-      auto load = heartbeat.second.resource_load_by_shape();
-      for (const auto &demand : load.resource_demands()) {
-        auto scheduling_key = ResourceSet(MapFromProtobuf(demand.shape()));
-        auto &aggregate_demand = aggregate_load[scheduling_key];
-        aggregate_demand.set_num_ready_requests_queued(
-            aggregate_demand.num_ready_requests_queued() +
-            demand.num_ready_requests_queued());
-        aggregate_demand.set_num_infeasible_requests_queued(
-            aggregate_demand.num_infeasible_requests_queued() +
-            demand.num_infeasible_requests_queued());
-        if (RayConfig::instance().report_worker_backlog()) {
-          aggregate_demand.set_backlog_size(aggregate_demand.backlog_size() +
-                                            demand.backlog_size());
-        }
-      }
-      heartbeat.second.clear_resource_load_by_shape();
-
       batch->add_batch()->Swap(&heartbeat.second);
     }
 
-    for (auto &demand : aggregate_load) {
-      auto demand_proto = batch->mutable_resource_load_by_shape()->add_resource_demands();
-      demand_proto->Swap(&demand.second);
-      for (const auto &resource_pair : demand.first.GetResourceMap()) {
-        (*demand_proto->mutable_shape())[resource_pair.first] = resource_pair.second;
-      }
-    }
-
-    // Update placement group load to heartbeat batch.
-    // This is updated only one per second.
-    if (placement_group_load_.has_value()) {
-      auto placement_group_load = placement_group_load_.value();
-      auto placement_group_load_proto = batch->mutable_placement_group_load();
-      placement_group_load_proto->Swap(placement_group_load.get());
-      placement_group_load_.reset();
-    }
-
     RAY_CHECK_OK(gcs_pub_sub_->Publish(HEARTBEAT_BATCH_CHANNEL, "",
                                        batch->SerializeAsString(), nullptr));
     heartbeat_buffer_.clear();

src/ray/gcs/gcs_server/gcs_node_manager.h

@@ -95,6 +95,18 @@ class GcsNodeManager : public rpc::NodeInfoHandler {
       rpc::GetAllAvailableResourcesReply *reply,
       rpc::SendReplyCallback send_reply_callback) override;
 
+  /// Handle get all heartbeat rpc request. Only used when light heartbeat enabled.
+  void HandleGetAllHeartbeat(const rpc::GetAllHeartbeatRequest &request,
+                             rpc::GetAllHeartbeatReply *reply,
+                             rpc::SendReplyCallback send_reply_callback) override;
+
+  /// Update heartbeat of given node.
+  ///
+  /// \param node_id Node id.
+  /// \param request Request containing heartbeat.
+  void UpdateNodeHeartbeat(const NodeID node_id,
+                           const rpc::ReportHeartbeatRequest &request);
+
   /// Add an alive node.
   ///
   /// \param node The info of the node to be added.
@@ -259,6 +271,8 @@ class GcsNodeManager : public rpc::NodeInfoHandler {
   std::list<std::pair<NodeID, int64_t>> sorted_dead_node_list_;
   /// Cluster resources.
   absl::flat_hash_map<NodeID, rpc::ResourceMap> cluster_resources_;
+  /// Newest heartbeat of all nodes.
+  absl::flat_hash_map<NodeID, rpc::HeartbeatTableData> node_heartbeats_;
   /// A buffer containing heartbeats received from node managers in the last tick.
   absl::flat_hash_map<NodeID, rpc::HeartbeatTableData> heartbeat_buffer_;
   /// Listeners which monitors the addition of nodes.

src/ray/gcs/redis_accessor.h

@@ -370,6 +370,11 @@ class RedisNodeInfoAccessor : public NodeInfoAccessor {
 
   void AsyncReReportHeartbeat() override;
 
+  Status AsyncGetAllHeartbeat(
+      const ItemCallback<rpc::HeartbeatBatchTableData> &callback) override {
+    return Status::NotImplemented("AsyncGetAllHeartbeat not implemented");
+  }
+
   Status AsyncSubscribeBatchHeartbeat(
       const ItemCallback<HeartbeatBatchTableData> &subscribe,
       const StatusCallback &done) override;

src/ray/protobuf/gcs_service.proto

@@ -198,6 +198,14 @@ message ReportHeartbeatReply {
   GcsStatus status = 1;
 }
 
+message GetAllHeartbeatRequest {
+}
+
+message GetAllHeartbeatReply {
+  GcsStatus status = 1;
+  HeartbeatBatchTableData heartbeat_data = 2;
+}
+
 message GetResourcesRequest {
   bytes node_id = 1;
 }
@@ -260,6 +268,8 @@ service NodeInfoGcsService {
   // Report heartbeat of a node to GCS Service.
   rpc ReportHeartbeat(ReportHeartbeatRequest) returns (ReportHeartbeatReply);
   // Get node's resources from GCS Service.
+  // Get newest heartbeat of all nodes from GCS Service.
+  rpc GetAllHeartbeat(GetAllHeartbeatRequest) returns (GetAllHeartbeatReply);
   rpc GetResources(GetResourcesRequest) returns (GetResourcesReply);
   // Update resources of a node in GCS Service.
   rpc UpdateResources(UpdateResourcesRequest) returns (UpdateResourcesReply);

src/ray/rpc/gcs_server/gcs_rpc_client.h

@@ -171,6 +171,11 @@ class GcsRpcClient {
   VOID_GCS_RPC_CLIENT_METHOD(NodeInfoGcsService, ReportHeartbeat,
                              node_info_grpc_client_, )
 
+  /// Get newest heartbeat of all nodes from GCS Service. Only used when light heartbeat
+  /// enabled.
+  VOID_GCS_RPC_CLIENT_METHOD(NodeInfoGcsService, GetAllHeartbeat,
+                             node_info_grpc_client_, )
+
   /// Get node's resources from GCS Service.
   VOID_GCS_RPC_CLIENT_METHOD(NodeInfoGcsService, GetResources, node_info_grpc_client_, )
 

src/ray/rpc/gcs_server/gcs_rpc_server.h

@@ -196,6 +196,10 @@ class NodeInfoGcsServiceHandler {
                                      ReportHeartbeatReply *reply,
                                      SendReplyCallback send_reply_callback) = 0;
 
+  virtual void HandleGetAllHeartbeat(const GetAllHeartbeatRequest &request,
+                                     GetAllHeartbeatReply *reply,
+                                     SendReplyCallback send_reply_callback) = 0;
+
   virtual void HandleGetResources(const GetResourcesRequest &request,
                                   GetResourcesReply *reply,
                                   SendReplyCallback send_reply_callback) = 0;
@@ -242,6 +246,7 @@ class NodeInfoGrpcService : public GrpcService {
     NODE_INFO_SERVICE_RPC_HANDLER(UnregisterNode);
     NODE_INFO_SERVICE_RPC_HANDLER(GetAllNodeInfo);
     NODE_INFO_SERVICE_RPC_HANDLER(ReportHeartbeat);
+    NODE_INFO_SERVICE_RPC_HANDLER(GetAllHeartbeat);
     NODE_INFO_SERVICE_RPC_HANDLER(GetResources);
     NODE_INFO_SERVICE_RPC_HANDLER(UpdateResources);
     NODE_INFO_SERVICE_RPC_HANDLER(DeleteResources);