[GCS]Add node realtime resource view (#10043)

src/ray/gcs/gcs_server/gcs_node_manager.cc

@@ -21,6 +21,18 @@
 namespace ray {
 namespace gcs {
 
+bool GcsNodeResource::IsSubset(
+    const std::unordered_map<std::string, double> &request_resources) const {
+  for (const auto &request_resource : request_resources) {
+    auto iter = resources_available_.find(request_resource.first);
+    if (iter == resources_available_.end() || iter->second < request_resource.second) {
+      return false;
+    }
+  }
+  return true;
+}
+
+//////////////////////////////////////////////////////////////////////////////////////////
 GcsNodeManager::NodeFailureDetector::NodeFailureDetector(
     boost::asio::io_service &io_service,
     std::shared_ptr<gcs::GcsTableStorage> gcs_table_storage,
@@ -231,6 +243,10 @@ void GcsNodeManager::HandleReportHeartbeat(const rpc::ReportHeartbeatRequest &re
   ClientID node_id = ClientID::FromBinary(request.heartbeat().client_id());
   auto heartbeat_data = std::make_shared<rpc::HeartbeatTableData>();
   heartbeat_data->CopyFrom(request.heartbeat());
+
+  // Update node realtime resources.
+  UpdateNodeRealtimeResources(node_id, *heartbeat_data);
+
   // Note: To avoid heartbeats being delayed by main thread, make sure heartbeat is always
   // handled by its own IO service.
   node_failure_detector_service_.post([this, node_id, heartbeat_data] {
@@ -451,5 +467,22 @@ void GcsNodeManager::StartNodeFailureDetector() {
   node_failure_detector_service_.post([this] { node_failure_detector_->Start(); });
 }
 
+void GcsNodeManager::UpdateNodeRealtimeResources(
+    const ClientID &node_id, const rpc::HeartbeatTableData &heartbeat) {
+  auto resources_available = MapFromProtobuf(heartbeat.resources_available());
+  auto iter = cluster_realtime_resources_.find(node_id);
+  if (iter != cluster_realtime_resources_.end()) {
+    iter->second->resources_available_ = resources_available;
+  } else {
+    cluster_realtime_resources_[node_id] =
+        std::make_shared<GcsNodeResource>(resources_available);
+  }
+}
+
+const absl::flat_hash_map<ClientID, std::shared_ptr<GcsNodeResource>>
+    &GcsNodeManager::GetClusterRealtimeResources() const {
+  return cluster_realtime_resources_;
+}
+
 }  // namespace gcs
 }  // namespace ray

src/ray/gcs/gcs_server/gcs_node_manager.h

@@ -27,6 +27,18 @@
 namespace ray {
 namespace gcs {
 
+/// Node resource information.
+class GcsNodeResource {
+ public:
+  explicit GcsNodeResource(std::unordered_map<std::string, double> resources_available)
+      : resources_available_(std::move(resources_available)) {}
+
+  bool IsSubset(const std::unordered_map<std::string, double> &request_resources) const;
+
+  /// Dynamic resource capacity.
+  std::unordered_map<std::string, double> resources_available_;
+};
+
 /// GcsNodeManager is responsible for managing and monitoring nodes as well as handing
 /// node and resource related rpc requests.
 /// This class is not thread-safe.
@@ -143,6 +155,14 @@ class GcsNodeManager : public rpc::NodeInfoHandler {
   /// Start node failure detector.
   void StartNodeFailureDetector();
 
+  // Update node realtime resources.
+  void UpdateNodeRealtimeResources(const ClientID &node_id,
+                                   const rpc::HeartbeatTableData &heartbeat);
+
+  /// Get cluster realtime resources.
+  const absl::flat_hash_map<ClientID, std::shared_ptr<GcsNodeResource>>
+      &GetClusterRealtimeResources() const;
+
  protected:
   class NodeFailureDetector {
    public:
@@ -239,6 +259,9 @@ class GcsNodeManager : public rpc::NodeInfoHandler {
   std::shared_ptr<gcs::GcsPubSub> gcs_pub_sub_;
   /// Storage for GCS tables.
   std::shared_ptr<gcs::GcsTableStorage> gcs_table_storage_;
+  /// Cluster realtime resources.
+  absl::flat_hash_map<ClientID, std::shared_ptr<GcsNodeResource>>
+      cluster_realtime_resources_;
 };
 
 }  // namespace gcs

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

@@ -86,6 +86,28 @@ TEST_F(GcsNodeManagerTest, TestListener) {
   }
 }
 
+TEST_F(GcsNodeManagerTest, TestGetClusterRealtimeResources) {
+  boost::asio::io_service io_service;
+  gcs::GcsNodeManager node_manager(io_service, io_service, gcs_pub_sub_,
+                                   gcs_table_storage_);
+
+  auto node_id = ClientID::FromRandom();
+  rpc::HeartbeatTableData heartbeat;
+  (*heartbeat.mutable_resources_available())["CPU"] = 10;
+  node_manager.UpdateNodeRealtimeResources(node_id, heartbeat);
+  auto node_resources = node_manager.GetClusterRealtimeResources();
+  std::unordered_map<std::string, double> request_resources;
+  request_resources["CPU"] = 9;
+  ASSERT_TRUE(node_resources[node_id]->IsSubset(request_resources));
+  request_resources["CPU"] = 10;
+  ASSERT_TRUE(node_resources[node_id]->IsSubset(request_resources));
+  request_resources["CPU"] = 10.1;
+  ASSERT_FALSE(node_resources[node_id]->IsSubset(request_resources));
+  request_resources.clear();
+  request_resources["GPU"] = 1;
+  ASSERT_FALSE(node_resources[node_id]->IsSubset(request_resources));
+}
+
 }  // namespace ray
 
 int main(int argc, char **argv) {