Repository: antgroup/ant-ray
Pull request #599: [VirtualCluster] Fix remove nodes from virtual cluster: return to primary cluster
Author: login: Chong-Li
Created: 2025-04-28T15:19:02Z | Merged: 2025-04-30T03:49:21Z

<!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?

Currently, when removing nodes from a virtual cluster, their corresponding replica instances are not correctly returned to the primary cluster. This PR fixes this bug.

<!-- Please give a short summary of the change and the problem this solves. -->

## Related issue number

<!-- For example: "Closes #1234" -->
#557 

## Checks

- [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR.
- [x] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for https://docs.ray.io/en/master/.
    - [ ] I've added any new APIs to the API Reference. For example, if I added a
           method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [x] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(



Diff:
diff --git a/python/ray/dashboard/modules/virtual_cluster/tests/test_virtual_cluster.py b/python/ray/dashboard/modules/virtual_cluster/tests/test_virtual_cluster.py
index 7a1192f28d77..95ffd24f7715 100644
--- a/python/ray/dashboard/modules/virtual_cluster/tests/test_virtual_cluster.py
+++ b/python/ray/dashboard/modules/virtual_cluster/tests/test_virtual_cluster.py
@@ -347,6 +347,7 @@ def test_remove_nodes_from_virtual_cluster(
         revision=0,
     )
     assert result["result"] is True
+    revision = result["data"]["revision"]
 
     resp = requests.get(webui_url + "/virtual_clusters")
     resp.raise_for_status()
@@ -401,8 +402,8 @@ def test_remove_nodes_from_virtual_cluster(
     actor = SmallActor.options(num_cpus=1, resources={"8c16g": 1}).remote()
     ray.get(actor.pid.remote(), timeout=10)
 
-    # This time, we try to remove the first `4c8g` node (which now
-    # belongs to the primary cluster) and the `8c16g` node (which is not idle).
+    # This time, we try to remove the first `4c8g` node (which has already been removed)
+    # and the `8c16g` node (which is not idle).
     nodes_to_remove = [template_to_nodes["4c8g"][0], template_to_nodes["8c16g"][0]]
     result = remove_nodes_from_virtual_cluster(
         webui_url=webui_url,
@@ -411,11 +412,22 @@ def test_remove_nodes_from_virtual_cluster(
     )
     assert result["result"] is False
     node_set_to_remove = set(nodes_to_remove)
-    # The two nodes we were trying to remove should fail.
+    # The two nodes we were trying to remove should both fail.
     assert len(result["data"]["nodesWithFailure"]) == 2
     for node_id in result["data"]["nodesWithFailure"]:
         assert node_id in node_set_to_remove
 
+    # If the first `4c8g` node now belongs to the primary cluster, then adding it back
+    # to the virtual cluster (update with two `4c8g` replicas) should succeed.
+    result = create_or_update_virtual_cluster(
+        webui_url=webui_url,
+        virtual_cluster_id="virtual_cluster_1",
+        divisible=divisible,
+        replica_sets={"4c8g": 2, "8c16g": 1},
+        revision=revision,
+    )
+    assert result["result"] is True
+
 
 @pytest.mark.parametrize(
     "ray_start_cluster_head",
diff --git a/src/ray/gcs/gcs_server/gcs_virtual_cluster.cc b/src/ray/gcs/gcs_server/gcs_virtual_cluster.cc
index 5afd4a661505..d7871b1684d2 100644
--- a/src/ray/gcs/gcs_server/gcs_virtual_cluster.cc
+++ b/src/ray/gcs/gcs_server/gcs_virtual_cluster.cc
@@ -153,17 +153,19 @@ void VirtualCluster::RemoveNodeInstances(ReplicaInstances replica_instances) {
 }
 
 Status VirtualCluster::RemoveNodeInstances(
-    const std::vector<std::string> &nodes_to_remove) {
+    const std::vector<std::string> &nodes_to_remove,
+    ReplicaInstances *removed_replica_instances) {
   absl::flat_hash_set<std::string> node_set_to_remove(nodes_to_remove.begin(),
                                                       nodes_to_remove.end());
-  ReplicaInstances replica_instances_to_remove;
   for (auto &[template_id, job_node_instances] : visible_node_instances_) {
     for (auto &[job_id, node_instances] : job_node_instances) {
       for (auto &[node_instance_id, node_instance] : node_instances) {
         auto removing_node_iter = node_set_to_remove.find(node_instance_id);
         if (removing_node_iter != node_set_to_remove.end() &&
-            IsNodeInstanceIdle(node_instance_id)) {
-          replica_instances_to_remove[template_id][job_id][node_instance_id] =
+            (!cluster_resource_manager_.HasNode(
+                 scheduling::NodeID(NodeID::FromHex(node_instance_id).Binary())) ||
+             IsNodeInstanceIdle(node_instance_id))) {
+          (*removed_replica_instances)[template_id][job_id][node_instance_id] =
               node_instance;
           node_set_to_remove.erase(removing_node_iter);
         }
@@ -179,7 +181,7 @@ Status VirtualCluster::RemoveNodeInstances(
         std::any(nodes_with_failure));
   }
 
-  RemoveNodeInstances(replica_instances_to_remove);
+  RemoveNodeInstances(*removed_replica_instances);
   return Status::OK();
 }
 
@@ -960,11 +962,19 @@ Status PrimaryCluster::RemoveNodesFromVirtualCluster(
     const rpc::RemoveNodesFromVirtualClusterRequest &request,
     RemoveNodesFromVirtualClusterCallback callback) {
   auto logical_cluster = GetLogicalCluster(request.virtual_cluster_id());
-  auto status = logical_cluster->RemoveNodeInstances(std::vector<std::string>(
-      request.nodes_to_remove().begin(), request.nodes_to_remove().end()));
+  ReplicaInstances removed_replica_instances;
+  auto status = logical_cluster->RemoveNodeInstances(
+      std::vector<std::string>(request.nodes_to_remove().begin(),
+                               request.nodes_to_remove().end()),
+      &removed_replica_instances);
   if (!status.ok()) {
     return status;
   }
+
+  // Add the instances that just removed from the virtual cluster to the primary cluster.
+  UpdateNodeInstances(/*replica_instances_to_add=*/std::move(removed_replica_instances),
+                      /*replica_instances_to_remove=*/ReplicaInstances());
+
   return async_data_flusher_(logical_cluster->ToProto(), std::move(callback));
 }
 
diff --git a/src/ray/gcs/gcs_server/gcs_virtual_cluster.h b/src/ray/gcs/gcs_server/gcs_virtual_cluster.h
index 59f2b3ad68a4..b4c12bfe1550 100644
--- a/src/ray/gcs/gcs_server/gcs_virtual_cluster.h
+++ b/src/ray/gcs/gcs_server/gcs_virtual_cluster.h
@@ -258,7 +258,14 @@ class VirtualCluster {
   std::shared_ptr<NodeInstance> ReplenishUndividedNodeInstance(
       std::shared_ptr<NodeInstance> node_instance_to_replenish);
 
-  Status RemoveNodeInstances(const std::vector<std::string> &nodes_to_remove);
+  /// Remove specified nodes from a virtual cluster.
+  ///
+  /// \param[in] nodes_to_remove The nodes to be removed.
+  /// \param[out] removed_replica_instances The replica instances that were
+  /// actually removed based on the `nodes_to_remove`.
+  /// \return The status of the removal.
+  Status RemoveNodeInstances(const std::vector<std::string> &nodes_to_remove,
+                             ReplicaInstances *removed_replica_instances);
 
  protected:
   /// Insert the node instances to the cluster.
diff --git a/src/ray/gcs/gcs_server/test/gcs_virtual_cluster_manager_test.cc b/src/ray/gcs/gcs_server/test/gcs_virtual_cluster_manager_test.cc
index 3528799b323f..172ccf519873 100644
--- a/src/ray/gcs/gcs_server/test/gcs_virtual_cluster_manager_test.cc
+++ b/src/ray/gcs/gcs_server/test/gcs_virtual_cluster_manager_test.cc
@@ -30,9 +30,9 @@ class GcsVirtualClusterManagerTest : public ::testing::Test {
   GcsVirtualClusterManagerTest() : cluster_resource_manager_(io_service_) {
     gcs_publisher_ = std::make_unique<gcs::GcsPublisher>(
         std::make_unique<ray::pubsub::MockPublisher>());
-    gcs_table_storage_ = std::make_unique<gcs::InMemoryGcsTableStorage>(io_service_);
+    gcs_table_storage_ = std::make_unique<gcs::InMemoryGcsTableStorage>();
     gcs_virtual_cluster_manager_ = std::make_unique<gcs::GcsVirtualClusterManager>(
-        *gcs_table_storage_, *gcs_publisher_, cluster_resource_manager_);
+        io_service_, *gcs_table_storage_, *gcs_publisher_, cluster_resource_manager_);
   }
 
   instrumented_io_context io_service_;
@@ -813,6 +813,148 @@ TEST_F(PrimaryClusterTest, GetVirtualClusters) {
   }
 }
 
+TEST_F(PrimaryClusterTest, RemoveNodesFromVirtualCluster) {
+  size_t node_count = 200;
+  size_t template_count = 10;
+  auto primary_cluster = InitPrimaryCluster(node_count, template_count);
+
+  std::string template_id_0 = "0";
+  std::string template_id_1 = "1";
+  size_t node_count_per_template = node_count / template_count;
+
+  // Create virtual_cluster_id_0 and check that the status is successful.
+  std::string virtual_cluster_id_0 = "virtual_cluster_id_0";
+  ASSERT_TRUE(CreateVirtualCluster(primary_cluster,
+                                   virtual_cluster_id_0,
+                                   {{template_id_0, 5}, {template_id_1, 10}})
+                  .ok());
+
+  {
+    // Check the primary cluster visible node instances.
+    const auto &visiable_node_instances = primary_cluster->GetVisibleNodeInstances();
+    // Check that template_id_0 remains template_count - 5 nodes, template_id_1 has
+    // template_count - 10 nodes.
+    EXPECT_EQ(visiable_node_instances.size(), template_count);
+    EXPECT_EQ(visiable_node_instances.at(template_id_0).size(), 1);
+    EXPECT_EQ(visiable_node_instances.at(template_id_0).at(kUndividedClusterId).size(),
+              node_count_per_template - 5);
+
+    EXPECT_EQ(visiable_node_instances.at(template_id_1).size(), 1);
+    EXPECT_EQ(visiable_node_instances.at(template_id_1).at(kUndividedClusterId).size(),
+              node_count_per_template - 10);
+
+    // Check that the revision unchanged.
+    EXPECT_NE(primary_cluster->GetRevision(), 0);
+  }
+
+  // Expecting to remove the first five `template_id_1` nodes from the virtual cluster.
+  size_t removing_node_count = 5;
+  absl::flat_hash_set<std::string> nodes_to_remove;
+  absl::flat_hash_set<std::string> remaining_nodes;
+  {
+    // Check the logical cluster virtual_cluster_id_0 visible node instances.
+    auto logical_cluster = primary_cluster->GetLogicalCluster(virtual_cluster_id_0);
+    ASSERT_NE(logical_cluster, nullptr);
+    const auto &visiable_node_instances = logical_cluster->GetVisibleNodeInstances();
+    // Check that template_id_0 has 5 nodes, template_id_1 has 10 nodes.
+    EXPECT_EQ(visiable_node_instances.size(), 2);
+    EXPECT_EQ(visiable_node_instances.at(template_id_0).size(), 1);
+    EXPECT_EQ(visiable_node_instances.at(template_id_0).at(kUndividedClusterId).size(),
+              5);
+
+    EXPECT_EQ(visiable_node_instances.at(template_id_1).size(), 1);
+    EXPECT_EQ(visiable_node_instances.at(template_id_1).at(kUndividedClusterId).size(),
+              10);
+    for (const auto &[node_id, _] :
+         visiable_node_instances.at(template_id_1).at(kUndividedClusterId)) {
+      if (removing_node_count > 0) {
+        nodes_to_remove.emplace(node_id);
+        removing_node_count--;
+      } else {
+        remaining_nodes.emplace(node_id);
+      }
+    }
+
+    // Check that the revision changed.
+    EXPECT_NE(logical_cluster->GetRevision(), 0);
+  }
+
+  {
+    rpc::RemoveNodesFromVirtualClusterRequest request;
+    request.set_virtual_cluster_id(virtual_cluster_id_0);
+    request.mutable_nodes_to_remove()->Assign(nodes_to_remove.begin(),
+                                              nodes_to_remove.end());
+    auto status = primary_cluster->RemoveNodesFromVirtualCluster(
+        request,
+        [this](const Status &status,
+               std::shared_ptr<rpc::VirtualClusterTableData> data,
+               const ReplicaSets *replica_sets_to_recommend) {
+          ASSERT_TRUE(status.ok());
+        });
+    ASSERT_TRUE(status.ok());
+  }
+
+  {
+    // Check the primary cluster visible node instances.
+    const auto &visiable_node_instances = primary_cluster->GetVisibleNodeInstances();
+    // Check that template_id_0 remains template_count - 5 nodes, template_id_1 has
+    // template_count - 5 nodes.
+    EXPECT_EQ(visiable_node_instances.size(), template_count);
+    EXPECT_EQ(visiable_node_instances.at(template_id_0).size(), 1);
+    EXPECT_EQ(visiable_node_instances.at(template_id_0).at(kUndividedClusterId).size(),
+              node_count_per_template - 5);
+
+    EXPECT_EQ(visiable_node_instances.at(template_id_1).size(), 1);
+    EXPECT_EQ(visiable_node_instances.at(template_id_1).at(kUndividedClusterId).size(),
+              node_count_per_template - 5);
+  }
+
+  {
+    // Check the logical cluster virtual_cluster_id_0 visible node instances.
+    auto logical_cluster = primary_cluster->GetLogicalCluster(virtual_cluster_id_0);
+    ASSERT_NE(logical_cluster, nullptr);
+    const auto &visiable_node_instances = logical_cluster->GetVisibleNodeInstances();
+    // Check that template_id_0 has 5 nodes, template_id_1 has 5 nodes.
+    EXPECT_EQ(visiable_node_instances.size(), 2);
+    EXPECT_EQ(visiable_node_instances.at(template_id_0).size(), 1);
+    EXPECT_EQ(visiable_node_instances.at(template_id_0).at(kUndividedClusterId).size(),
+              5);
+
+    EXPECT_EQ(visiable_node_instances.at(template_id_1).size(), 1);
+    EXPECT_EQ(visiable_node_instances.at(template_id_1).at(kUndividedClusterId).size(),
+              5);
+
+    // The nodes left in the virtual cluster should be exactly the same with
+    // `remaining_nodes`. Otherwise, wrong nodes were removed from the virtual cluster.
+    for (const auto &[node_id, _] :
+         visiable_node_instances.at(template_id_1).at(kUndividedClusterId)) {
+      ASSERT_TRUE(remaining_nodes.contains(node_id));
+    }
+  }
+
+  {
+    // Remove the `nodes_to_remove` again, we shall see failure.
+    rpc::RemoveNodesFromVirtualClusterRequest request;
+    request.set_virtual_cluster_id(virtual_cluster_id_0);
+    request.mutable_nodes_to_remove()->Assign(nodes_to_remove.begin(),
+                                              nodes_to_remove.end());
+    // Remove nodes that do not exist in the virtual cluster.
+    auto status = primary_cluster->RemoveNodesFromVirtualCluster(
+        request,
+        [this, &nodes_to_remove](const Status &status,
+                                 std::shared_ptr<rpc::VirtualClusterTableData> data,
+                                 const ReplicaSets *replica_sets_to_recommend) {});
+    ASSERT_TRUE(status.IsUnsafeToRemove());
+    // The failure message should tell which nodes are not eligible for removal.
+    auto nodes_with_failure = std::any_cast<std::vector<std::string>>(status.data());
+    ASSERT_EQ(nodes_with_failure.size(), 5);
+    for (const auto &node_id : nodes_with_failure) {
+      ASSERT_EQ(nodes_to_remove.erase(node_id), 1);
+    }
+    ASSERT_TRUE(nodes_to_remove.empty());
+  }
+}
+
 // ┌───────────────────────────────────────────────────┐
 // │ ┌───────────────────┐ ┌─────────┐                 │
 // │ │ virtual_cluster_1 │ │         │    Exclusive    │
@@ -899,8 +1041,7 @@ TEST_F(FailoverTest, FailoverNormal) {
   ASSERT_EQ(virtual_clusters_data_.size(), 4);
 
   // Mock a gcs_init_data.
-  instrumented_io_context io_service;
-  gcs::InMemoryGcsTableStorage gcs_table_storage(io_service);
+  gcs::InMemoryGcsTableStorage gcs_table_storage;
   MockGcsInitData gcs_init_data(gcs_table_storage);
   gcs_init_data.SetNodes(nodes_);
   gcs_init_data.SetVirtualClusters(virtual_clusters_data_);
@@ -961,8 +1102,7 @@ TEST_F(FailoverTest, FailoverWithDeadNodes) {
 
   {
     // Mock a gcs_init_data.
-    instrumented_io_context io_service;
-    gcs::InMemoryGcsTableStorage gcs_table_storage(io_service);
+    gcs::InMemoryGcsTableStorage gcs_table_storage;
     MockGcsInitData gcs_init_data(gcs_table_storage);
     gcs_init_data.SetNodes(nodes_);
     gcs_init_data.SetVirtualClusters(virtual_clusters_data_);
@@ -1008,8 +1148,7 @@ TEST_F(FailoverTest, FailoverWithDeadNodes) {
     ASSERT_TRUE(async_data_flusher_(virtual_cluster_1->ToProto(), nullptr).ok());
 
     // Mock a gcs_init_data.
-    instrumented_io_context io_service;
-    gcs::InMemoryGcsTableStorage gcs_table_storage(io_service);
+    gcs::InMemoryGcsTableStorage gcs_table_storage;
     MockGcsInitData gcs_init_data(gcs_table_storage);
     gcs_init_data.SetNodes(nodes_);
     gcs_init_data.SetVirtualClusters(virtual_clusters_data_);
@@ -1081,8 +1220,7 @@ TEST_F(FailoverTest, OnlyFlushJobClusters) {
     // async_data_flusher_(virtual_cluster_1->ToProto(), nullptr);
 
     // Mock a gcs_init_data.
-    instrumented_io_context io_service;
-    gcs::InMemoryGcsTableStorage gcs_table_storage(io_service);
+    gcs::InMemoryGcsTableStorage gcs_table_storage;
     MockGcsInitData gcs_init_data(gcs_table_storage);
     gcs_init_data.SetNodes(nodes_);
     gcs_init_data.SetVirtualClusters(virtual_clusters_data_);


Reviews:
user: login: xsuler
body: LGTM
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-29T08:33:12Z
user: login: wumuzi520
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-29T09:49:47Z
user: login: Chong-Li
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-29T11:58:31Z
user: login: wumuzi520
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T00:58:37Z
user: login: wumuzi520
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-30T01:03:32Z

Review comments:
diff_hunk: @@ -258,7 +258,14 @@ class VirtualCluster {
   std::shared_ptr<NodeInstance> ReplenishUndividedNodeInstance(
       std::shared_ptr<NodeInstance> node_instance_to_replenish);
 
-  Status RemoveNodeInstances(const std::vector<std::string> &nodes_to_remove);
+  /// Remove specified nodes from a virtual cluster.
+  ///
+  /// \param nodes_to_remove The nodes to be removed.
+  /// \param removed_replica_instances The replica instances that were
+  /// actually removed based on the `nodes_to_remove`.
+  /// \return The status of the removal.
+  Status RemoveNodeInstances(const std::vector<std::string> &nodes_to_remove,
path: src/ray/gcs/gcs_server/gcs_virtual_cluster.h
user: login: wumuzi520
body: It's better to redefine the method to
```
ReplicaInstances RemoveNodeInstances(const std::vector<std::string> &nodes_to_remove);
```
created_at: 2025-04-29T09:49:47Z
author_association: COLLABORATOR
diff_hunk: @@ -258,7 +258,14 @@ class VirtualCluster {
   std::shared_ptr<NodeInstance> ReplenishUndividedNodeInstance(
       std::shared_ptr<NodeInstance> node_instance_to_replenish);
 
-  Status RemoveNodeInstances(const std::vector<std::string> &nodes_to_remove);
+  /// Remove specified nodes from a virtual cluster.
+  ///
+  /// \param nodes_to_remove The nodes to be removed.
+  /// \param removed_replica_instances The replica instances that were
+  /// actually removed based on the `nodes_to_remove`.
+  /// \return The status of the removal.
+  Status RemoveNodeInstances(const std::vector<std::string> &nodes_to_remove,
path: src/ray/gcs/gcs_server/gcs_virtual_cluster.h
user: login: Chong-Li
body: The tricky part is: the function have to return both nodes_with_failure (if exists) and the replica instances. That is why it needs an output param. How about adding a `\param[out]` label to the comment?
created_at: 2025-04-29T11:58:26Z
author_association: COLLABORATOR
diff_hunk: @@ -258,7 +258,14 @@ class VirtualCluster {
   std::shared_ptr<NodeInstance> ReplenishUndividedNodeInstance(
       std::shared_ptr<NodeInstance> node_instance_to_replenish);
 
-  Status RemoveNodeInstances(const std::vector<std::string> &nodes_to_remove);
+  /// Remove specified nodes from a virtual cluster.
+  ///
+  /// \param nodes_to_remove The nodes to be removed.
+  /// \param removed_replica_instances The replica instances that were
+  /// actually removed based on the `nodes_to_remove`.
+  /// \return The status of the removal.
+  Status RemoveNodeInstances(const std::vector<std::string> &nodes_to_remove,
path: src/ray/gcs/gcs_server/gcs_virtual_cluster.h
user: login: wumuzi520
body: Okey, BTW I found that there are some compile error in this CI: https://github.com/antgroup/ant-ray/actions/runs/14730932680/job/41360948939?pr=599  I'm not sure is it related to this PR.
created_at: 2025-04-30T00:58:36Z
author_association: COLLABORATOR

PR comments:
user: login: wumuzi520
created_at: 2025-04-30T01:54:34Z
body: BTW, please add an extra remote nodes from virtual cluster API to the document: https://antgroup.github.io/ant-ray/virtual-cluster/management.html
author_association: COLLABORATOR
user: login: Chong-Li
created_at: 2025-04-30T09:21:53Z
body: > BTW, please add an extra remote nodes from virtual cluster API to the document: https://antgroup.github.io/ant-ray/virtual-cluster/management.html

I missed this comment when merging the PR. A new one (https://github.com/antgroup/ant-ray/pull/600) is created for the doc.
author_association: COLLABORATOR