[14/N][VirtualCluster] Raylet only subscribe mixed cluster data (#431)

Signed-off-by: 黑驰 <[email protected]>

src/mock/ray/gcs/gcs_client/accessor.h

@@ -365,7 +365,9 @@ class MockVirtualClusterInfoAccessor : public VirtualClusterInfoAccessor {
               (override));
   MOCK_METHOD(Status,
               AsyncGetAll,
-              (const MultiItemCallback<rpc::VirtualClusterTableData> &callback),
+              (bool include_job_clusters,
+               bool only_include_mixed_clusters,
+               (const MultiItemCallback<rpc::VirtualClusterTableData> &callback)),
               (override));
   MOCK_METHOD(Status,
               AsyncSubscribeAll,

src/ray/gcs/gcs_client/accessor.ant.cc

@@ -47,10 +47,13 @@ Status VirtualClusterInfoAccessor::AsyncGet(
 }
 
 Status VirtualClusterInfoAccessor::AsyncGetAll(
+    bool include_job_clusters,
+    bool only_include_mixed_clusters,
     const MultiItemCallback<rpc::VirtualClusterTableData> &callback) {
   RAY_LOG(DEBUG) << "Getting all virtual cluster info.";
   rpc::GetVirtualClustersRequest request;
   request.set_include_job_clusters(true);
+  request.set_only_include_mixed_clusters(true);
   client_impl_->GetGcsRpcClient().GetVirtualClusters(
       request, [callback](const Status &status, rpc::GetVirtualClustersReply &&reply) {
         callback(
@@ -79,7 +82,10 @@ Status VirtualClusterInfoAccessor::AsyncSubscribeAll(
             done(status);
           }
         };
-    RAY_CHECK_OK(AsyncGetAll(callback));
+    RAY_CHECK_OK(AsyncGetAll(
+        /*include_job_clusters=*/true,
+        /*only_include_mixed_clusters=*/true,
+        callback));
   };
   subscribe_operation_ = [this, subscribe](const StatusCallback &done) {
     return client_impl_->GetGcsSubscriber().SubscribeAllVirtualClusters(subscribe, done);

src/ray/gcs/gcs_client/accessor.h

@@ -1016,6 +1016,8 @@ class VirtualClusterInfoAccessor {
   /// \param callback Callback that will be called after lookup finished.
   /// \return Status
   virtual Status AsyncGetAll(
+      bool include_job_clusters,
+      bool only_include_mixed_clusters,
       const MultiItemCallback<rpc::VirtualClusterTableData> &callback);
 
   /// Subscribe to virtual cluster updates.

src/ray/gcs/gcs_server/gcs_virtual_cluster.cc

@@ -519,6 +519,7 @@ void PrimaryCluster::GetVirtualClustersData(rpc::GetVirtualClustersRequest reque
   std::vector<std::shared_ptr<rpc::VirtualClusterTableData>> virtual_cluster_data_list;
   auto virtual_cluster_id = request.virtual_cluster_id();
   bool include_job_clusters = request.include_job_clusters();
+  bool only_include_mixed_cluster = request.only_include_mixed_clusters();
 
   auto visit_proto_data = [&](const VirtualCluster *cluster) {
     if (include_job_clusters && cluster->GetMode() == rpc::AllocationMode::EXCLUSIVE) {
@@ -528,6 +529,10 @@ void PrimaryCluster::GetVirtualClustersData(rpc::GetVirtualClustersRequest reque
             callback(job_cluster->ToProto());
           });
     }
+    if (only_include_mixed_cluster &&
+        cluster->GetMode() == rpc::AllocationMode::EXCLUSIVE) {
+      return;
+    }
     if (cluster->GetID() != kPrimaryClusterID) {
       // Skip the primary cluster's proto data.
       callback(cluster->ToProto());

src/ray/gcs/gcs_server/gcs_virtual_cluster_manager.cc

@@ -216,6 +216,12 @@ Status GcsVirtualClusterManager::FlushAndPublish(
   auto on_done = [this, data, callback = std::move(callback)](const Status &status) {
     // The backend storage is supposed to be reliable, so the status must be ok.
     RAY_CHECK_OK(status);
+    if (data->mode() != rpc::AllocationMode::MIXED) {
+      // Tasks can only be scheduled on the nodes in the mixed cluster, so we just need to
+      // publish the mixed cluster data.
+      return;
+    }
+
     RAY_CHECK_OK(gcs_publisher_.PublishVirtualCluster(
         VirtualClusterID::FromBinary(data->id()), *data, nullptr));
     if (callback) {

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

@@ -58,7 +58,7 @@ class VirtualClusterTest : public ::testing::Test {
     for (size_t i = 0; i < node_count; ++i) {
       auto node = Mocker::GenNodeInfo();
       auto template_id = std::to_string(i % template_count);
-      node->set_template_id(template_id);
+      node->set_node_type_name(template_id);
       primary_cluster->OnNodeAdd(*node);
       if (template_id_to_nodes != nullptr) {
         (*template_id_to_nodes)[template_id].emplace(NodeID::FromBinary(node->node_id()),
@@ -628,7 +628,6 @@ TEST_F(PrimaryClusterTest, GetVirtualClusters) {
         virtual_clusters_data_map;
     primary_cluster->GetVirtualClustersData(
         request, [this, &virtual_clusters_data_map](auto data) {
-          RAY_LOG(INFO) << "xxx: " << data->id();
           virtual_clusters_data_map.emplace(data->id(), data);
         });
     ASSERT_EQ(virtual_clusters_data_map.size(), 1);
@@ -646,6 +645,17 @@ TEST_F(PrimaryClusterTest, GetVirtualClusters) {
     auto job_cluster = virtual_cluster_0->GetJobCluster("job_1");
     ASSERT_TRUE(job_cluster != nullptr);
     ASSERT_TRUE(virtual_clusters_data_map.contains(job_cluster->GetID()));
+
+    virtual_clusters_data_map.clear();
+    request.set_include_job_clusters(true);
+    request.set_only_include_mixed_clusters(true);
+    primary_cluster->GetVirtualClustersData(
+        request, [this, &virtual_clusters_data_map](auto data) {
+          virtual_clusters_data_map.emplace(data->id(), data);
+        });
+    ASSERT_EQ(virtual_clusters_data_map.size(), 1);
+    ASSERT_FALSE(virtual_clusters_data_map.contains(virtual_cluster_id_0));
+    ASSERT_TRUE(virtual_clusters_data_map.contains(job_cluster->GetID()));
   }
 }
 

src/ray/protobuf/gcs_service.proto

@@ -852,6 +852,8 @@ message GetVirtualClustersRequest {
   string virtual_cluster_id = 1;
   // Wether include job clusters.
   bool include_job_clusters = 2;
+  // It will reply mixed clusters if only_include_mixed_clusters is true.
+  bool only_include_mixed_clusters = 3;
 }
 
 message GetVirtualClustersReply {