[GCS]GCS resource manager support scheduling resource (#12780)

* add part code

* add part code

* fix review comments

* rebase master

* add part code

* add part code

* fix review comments

* add part code

* fix code style

* fix ut bug

* fix ut bug

* fix review comments

* fix review comment

Co-authored-by: 灵洵 <fengbin.ffb@antgroup.com>

src/ray/gcs/gcs_server/gcs_node_manager.cc

@@ -317,7 +317,7 @@ void GcsNodeManager::UpdateNodeRealtimeResources(
   if (!light_report_resource_usage_enabled_ ||
       gcs_resource_manager_->GetClusterResources().count(node_id) == 0 ||
       resource_data.resources_available_changed()) {
-    gcs_resource_manager_->UpdateResources(
+    gcs_resource_manager_->SetAvailableResources(
         node_id, ResourceSet(MapFromProtobuf(resource_data.resources_available())));
   }
 }

src/ray/gcs/gcs_server/gcs_placement_group_scheduler.cc

@@ -66,7 +66,7 @@ ScheduleMap GcsStrictPackStrategy::Schedule(
   // Filter candidate nodes.
   std::vector<std::pair<int64_t, NodeID>> candidate_nodes;
   for (auto &node : context->cluster_resources_) {
-    if (required_resources.IsSubset(node.second)) {
+    if (required_resources.IsSubset(node.second.GetAvailableResources())) {
       candidate_nodes.emplace_back((*context->node_to_bundles_)[node.first], node.first);
     }
   }
@@ -99,7 +99,8 @@ ScheduleMap GcsPackStrategy::Schedule(
   for (const auto &bundle : bundles) {
     const auto &required_resources = bundle->GetRequiredResources();
     for (const auto &node : context->cluster_resources_) {
-      if (IsAvailableResourceSufficient(node.second, allocated_resources[node.first],
+      if (IsAvailableResourceSufficient(node.second.GetAvailableResources(),
+                                        allocated_resources[node.first],
                                         required_resources)) {
         schedule_map[bundle->BundleId()] = node.first;
         allocated_resources[node.first].AddResources(required_resources);
@@ -135,7 +136,8 @@ ScheduleMap GcsSpreadStrategy::Schedule(
     // that meets the resource requirements. `iter_begin` is the next node of the last
     // selected node.
     for (; iter != candidate_nodes.end(); ++iter) {
-      if (IsAvailableResourceSufficient(iter->second, allocated_resources[iter->first],
+      if (IsAvailableResourceSufficient(iter->second.GetAvailableResources(),
+                                        allocated_resources[iter->first],
                                         required_resources)) {
         schedule_map[bundle->BundleId()] = iter->first;
         allocated_resources[iter->first].AddResources(required_resources);
@@ -153,8 +155,9 @@ ScheduleMap GcsSpreadStrategy::Schedule(
       if (iter_begin != candidate_nodes.begin()) {
         // Traverse all the nodes from `candidate_nodes.begin()` to `iter_begin`.
         for (iter = candidate_nodes.begin(); iter != iter_begin; ++iter) {
-          if (IsAvailableResourceSufficient(
-                  iter->second, allocated_resources[iter->first], required_resources)) {
+          if (IsAvailableResourceSufficient(iter->second.GetAvailableResources(),
+                                            allocated_resources[iter->first],
+                                            required_resources)) {
             schedule_map[bundle->BundleId()] = iter->first;
             allocated_resources[iter->first].AddResources(required_resources);
             break;
@@ -207,7 +210,8 @@ ScheduleMap GcsStrictSpreadStrategy::Schedule(
     auto iter = candidate_nodes.begin();
     for (; iter != candidate_nodes.end(); ++iter) {
       if (!allocated_resources.contains(iter->first) &&
-          IsAvailableResourceSufficient(iter->second, allocated_resources[iter->first],
+          IsAvailableResourceSufficient(iter->second.GetAvailableResources(),
+                                        allocated_resources[iter->first],
                                         required_resources)) {
         schedule_map[bundle->BundleId()] = iter->first;
         allocated_resources[iter->first].AddResources(required_resources);

src/ray/gcs/gcs_server/gcs_placement_group_scheduler.h

@@ -89,9 +89,10 @@ class GcsPlacementGroupSchedulerInterface {
 /// ScheduleContext provides information that are needed for bundle scheduling decision.
 class ScheduleContext {
  public:
-  ScheduleContext(std::shared_ptr<absl::flat_hash_map<NodeID, int64_t>> node_to_bundles,
-                  const absl::optional<std::shared_ptr<BundleLocations>> bundle_locations,
-                  const absl::flat_hash_map<NodeID, ResourceSet> &cluster_resources)
+  ScheduleContext(
+      std::shared_ptr<absl::flat_hash_map<NodeID, int64_t>> node_to_bundles,
+      const absl::optional<std::shared_ptr<BundleLocations>> bundle_locations,
+      const absl::flat_hash_map<NodeID, SchedulingResources> &cluster_resources)
       : node_to_bundles_(std::move(node_to_bundles)),
         bundle_locations_(bundle_locations),
         cluster_resources_(cluster_resources) {}
@@ -101,7 +102,7 @@ class ScheduleContext {
   // The locations of existing bundles for this placement group.
   const absl::optional<std::shared_ptr<BundleLocations>> bundle_locations_;
   // The available resources of all nodes.
-  const absl::flat_hash_map<NodeID, ResourceSet> &cluster_resources_;
+  const absl::flat_hash_map<NodeID, SchedulingResources> &cluster_resources_;
 };
 
 class GcsScheduleStrategy {

src/ray/gcs/gcs_server/gcs_resource_manager.cc

@@ -28,7 +28,7 @@ void GcsResourceManager::HandleGetResources(const rpc::GetResourcesRequest &requ
   NodeID node_id = NodeID::FromBinary(request.node_id());
   auto iter = cluster_resources_.find(node_id);
   if (iter != cluster_resources_.end()) {
-    for (auto &resource : iter->second.items()) {
+    for (const auto &resource : iter->second.items()) {
       (*reply->mutable_resources())[resource.first] = resource.second;
     }
   }
@@ -42,19 +42,27 @@ void GcsResourceManager::HandleUpdateResources(
   NodeID node_id = NodeID::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();
+  std::unordered_map<std::string, double> to_be_updated_resources;
+  for (const auto &entry : request.resources()) {
+    to_be_updated_resources.emplace(entry.first, entry.second.resource_capacity());
+  }
+
   if (iter != cluster_resources_.end()) {
-    for (auto &entry : to_be_updated_resources) {
+    for (const auto &entry : request.resources()) {
       (*iter->second.mutable_items())[entry.first] = entry.second;
     }
+    UpdateResourceCapacity(node_id, to_be_updated_resources);
     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) {
-        (*node_resource_change.mutable_updated_resources())[it.first] =
-            it.second.resource_capacity();
+      for (const auto &it : to_be_updated_resources) {
+        const auto &resource_name = it.first;
+        const auto &resource_capacity = it.second;
+        auto &node_updated_resources =
+            (*node_resource_change.mutable_updated_resources());
+        node_updated_resources[resource_name] = resource_capacity;
       }
       RAY_CHECK_OK(gcs_pub_sub_->Publish(NODE_RESOURCE_CHANNEL, node_id.Hex(),
                                          node_resource_change.SerializeAsString(),
@@ -82,7 +90,9 @@ void GcsResourceManager::HandleDeleteResources(
   auto resource_names = VectorFromProtobuf(request.resource_name_list());
   auto iter = cluster_resources_.find(node_id);
   if (iter != cluster_resources_.end()) {
-    for (auto &resource_name : resource_names) {
+    DeleteResources(node_id, resource_names);
+
+    for (const auto &resource_name : resource_names) {
       RAY_IGNORE_EXPR(iter->second.mutable_items()->erase(resource_name));
     }
     auto on_done = [this, node_id, resource_names, reply,
@@ -115,7 +125,7 @@ void GcsResourceManager::HandleGetAllAvailableResources(
   for (const auto &iter : cluster_scheduling_resources_) {
     rpc::AvailableResources resource;
     resource.set_node_id(iter.first.Binary());
-    for (const auto &res : iter.second.GetResourceAmountMap()) {
+    for (const auto &res : iter.second.GetAvailableResources().GetResourceAmountMap()) {
       (*resource.mutable_resources_available())[res.first] = res.second.ToDouble();
     }
     reply->add_resources_list()->CopyFrom(resource);
@@ -134,14 +144,39 @@ void GcsResourceManager::Initialize(const GcsInitData &gcs_init_data) {
   }
 }
 
-const absl::flat_hash_map<NodeID, ResourceSet> &GcsResourceManager::GetClusterResources()
-    const {
+const absl::flat_hash_map<NodeID, SchedulingResources>
+    &GcsResourceManager::GetClusterResources() const {
   return cluster_scheduling_resources_;
 }
 
-void GcsResourceManager::UpdateResources(const NodeID &node_id,
-                                         const ResourceSet &resources) {
-  cluster_scheduling_resources_[node_id] = resources;
+void GcsResourceManager::SetAvailableResources(const NodeID &node_id,
+                                               const ResourceSet &resources) {
+  cluster_scheduling_resources_[node_id].SetAvailableResources(ResourceSet(resources));
+}
+
+void GcsResourceManager::UpdateResourceCapacity(
+    const NodeID &node_id,
+    const std::unordered_map<std::string, double> &changed_resources) {
+  auto iter = cluster_scheduling_resources_.find(node_id);
+  if (iter != cluster_scheduling_resources_.end()) {
+    SchedulingResources &scheduling_resources = iter->second;
+    for (const auto &entry : changed_resources) {
+      scheduling_resources.UpdateResourceCapacity(entry.first, entry.second);
+    }
+  } else {
+    cluster_scheduling_resources_.emplace(
+        node_id, SchedulingResources(ResourceSet(changed_resources)));
+  }
+}
+
+void GcsResourceManager::DeleteResources(
+    const NodeID &node_id, const std::vector<std::string> &deleted_resources) {
+  auto iter = cluster_scheduling_resources_.find(node_id);
+  if (iter != cluster_scheduling_resources_.end()) {
+    for (auto &resource_name : deleted_resources) {
+      iter->second.DeleteResource(resource_name);
+    }
+  }
 }
 
 void GcsResourceManager::OnNodeAdd(const NodeID &node_id) {
@@ -158,10 +193,10 @@ bool GcsResourceManager::AcquireResources(const NodeID &node_id,
                                           const ResourceSet &required_resources) {
   auto iter = cluster_scheduling_resources_.find(node_id);
   if (iter != cluster_scheduling_resources_.end()) {
-    if (!required_resources.IsSubset(iter->second)) {
+    if (!required_resources.IsSubset(iter->second.GetAvailableResources())) {
       return false;
     }
-    iter->second.SubtractResourcesStrict(required_resources);
+    iter->second.Acquire(required_resources);
   }
   // If node dead, we will not find the node. This is a normal scenario, so it returns
   // true.
@@ -172,7 +207,7 @@ bool GcsResourceManager::ReleaseResources(const NodeID &node_id,
                                           const ResourceSet &acquired_resources) {
   auto iter = cluster_scheduling_resources_.find(node_id);
   if (iter != cluster_scheduling_resources_.end()) {
-    iter->second.AddResources(acquired_resources);
+    iter->second.Release(acquired_resources);
   }
   // If node dead, we will not find the node. This is a normal scenario, so it returns
   // true.

src/ray/gcs/gcs_server/gcs_resource_manager.h

@@ -68,7 +68,7 @@ class GcsResourceManager : public rpc::NodeResourceInfoHandler {
   /// Get the resources of all nodes in the cluster.
   ///
   /// \return The resources of all nodes in the cluster.
-  const absl::flat_hash_map<NodeID, ResourceSet> &GetClusterResources() const;
+  const absl::flat_hash_map<NodeID, SchedulingResources> &GetClusterResources() const;
 
   /// Handle a node registration.
   ///
@@ -80,11 +80,11 @@ class GcsResourceManager : public rpc::NodeResourceInfoHandler {
   /// \param node_id The specified node id.
   void OnNodeDead(const NodeID &node_id);
 
-  /// Update the resources of the specified node.
+  /// Set the available resources of the specified node.
   ///
   /// \param node_id Id of a node.
-  /// \param resources Resources of a node.
-  void UpdateResources(const NodeID &node_id, const ResourceSet &resources);
+  /// \param resources Available resources of a node.
+  void SetAvailableResources(const NodeID &node_id, const ResourceSet &resources);
 
   /// Acquire resources from the specified node. It will deduct directly from the node
   /// resources.
@@ -110,15 +110,30 @@ class GcsResourceManager : public rpc::NodeResourceInfoHandler {
 
   std::string DebugString() const;
 
+  /// Update the total resources and available resources of the specified node.
+  ///
+  /// \param node_id Id of a node.
+  /// \param changed_resources Changed resources of a node.
+  void UpdateResourceCapacity(
+      const NodeID &node_id,
+      const std::unordered_map<std::string, double> &changed_resources);
+
  private:
+  /// Delete the scheduling resources of the specified node.
+  ///
+  /// \param node_id Id of a node.
+  /// \param deleted_resources Deleted resources of a node.
+  void DeleteResources(const NodeID &node_id,
+                       const std::vector<std::string> &deleted_resources);
+
   /// A publisher for publishing gcs messages.
   std::shared_ptr<gcs::GcsPubSub> gcs_pub_sub_;
   /// Storage for GCS tables.
   std::shared_ptr<gcs::GcsTableStorage> gcs_table_storage_;
   /// Cluster resources.
   absl::flat_hash_map<NodeID, rpc::ResourceMap> cluster_resources_;
-  /// Map from node id to the resources of the node.
-  absl::flat_hash_map<NodeID, ResourceSet> cluster_scheduling_resources_;
+  /// Map from node id to the scheduling resources of the node.
+  absl::flat_hash_map<NodeID, SchedulingResources> cluster_scheduling_resources_;
 
   /// Debug info.
   enum CountType {

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

@@ -38,7 +38,7 @@ TEST_F(GcsResourceManagerTest, TestBasic) {
   std::unordered_map<std::string, double> resource_map;
   resource_map[cpu_resource] = 10;
   ResourceSet resource_set(resource_map);
-  gcs_resource_manager_->UpdateResources(node_id, resource_set);
+  gcs_resource_manager_->UpdateResourceCapacity(node_id, resource_map);
 
   // Get and check cluster resources.
   const auto &cluster_resource = gcs_resource_manager_->GetClusterResources();