[Placement Group]Trigger placement group scheduling when a new node is added (#9905)

src/ray/common/ray_config_def.h

@@ -285,6 +285,8 @@ RAY_CONFIG(int64_t, gcs_redis_heartbeat_interval_milliseconds, 100)
 RAY_CONFIG(uint32_t, gcs_lease_worker_retry_interval_ms, 200)
 /// Duration to wait between retries for creating actor in gcs server.
 RAY_CONFIG(uint32_t, gcs_create_actor_retry_interval_ms, 200)
+/// Duration to wait between retries for creating placement group in gcs server.
+RAY_CONFIG(uint32_t, gcs_create_placement_group_retry_interval_ms, 200)
 
 /// Maximum number of times to retry putting an object when the plasma store is full.
 /// Can be set to -1 to enable unlimited retries.

src/ray/gcs/gcs_server/gcs_placement_group_manager.cc

@@ -16,6 +16,7 @@
 
 #include "ray/common/ray_config.h"
 #include "ray/gcs/pb_util.h"
+#include "ray/util/asio_util.h"
 #include "src/ray/protobuf/gcs.pb.h"
 
 namespace ray {
@@ -61,9 +62,9 @@ GcsPlacementGroupManager::GcsPlacementGroupManager(
     boost::asio::io_context &io_context,
     std::shared_ptr<GcsPlacementGroupSchedulerInterface> scheduler,
     std::shared_ptr<gcs::GcsTableStorage> gcs_table_storage)
-    : gcs_placement_group_scheduler_(std::move(scheduler)),
-      gcs_table_storage_(std::move(gcs_table_storage)),
-      reschedule_timer_(io_context) {}
+    : io_context_(io_context),
+      gcs_placement_group_scheduler_(std::move(scheduler)),
+      gcs_table_storage_(std::move(gcs_table_storage)) {}
 
 void GcsPlacementGroupManager::RegisterPlacementGroup(
     const ray::rpc::CreatePlacementGroupRequest &request, EmptyCallback callback) {
@@ -105,11 +106,11 @@ void GcsPlacementGroupManager::OnPlacementGroupCreationFailed(
   // registered.
   pending_placement_groups_.emplace_back(std::move(placement_group));
   is_creating_ = false;
-  ScheduleTick();
+  RetryCreatingPlacementGroup();
 }
 
 void GcsPlacementGroupManager::OnPlacementGroupCreationSuccess(
-    std::shared_ptr<GcsPlacementGroup> placement_group) {
+    const std::shared_ptr<GcsPlacementGroup> &placement_group) {
   RAY_LOG(INFO) << "Successfully created placement group " << placement_group->GetName();
   placement_group->UpdateState(rpc::PlacementGroupTableData::ALIVE);
   auto placement_group_id = placement_group->GetPlacementGroupID();
@@ -173,15 +174,9 @@ void GcsPlacementGroupManager::HandleCreatePlacementGroup(
       }));
 }
 
-void GcsPlacementGroupManager::ScheduleTick() {
-  reschedule_timer_.expires_from_now(boost::posix_time::milliseconds(500));
-  reschedule_timer_.async_wait([this](const boost::system::error_code &error) {
-    if (error == boost::asio::error::operation_aborted) {
-      return;
-    } else {
-      SchedulePendingPlacementGroups();
-    }
-  });
+void GcsPlacementGroupManager::RetryCreatingPlacementGroup() {
+  execute_after(io_context_, [this] { SchedulePendingPlacementGroups(); },
+                RayConfig::instance().gcs_create_placement_group_retry_interval_ms());
 }
 
 }  // namespace gcs

src/ray/gcs/gcs_server/gcs_placement_group_manager.h

@@ -139,15 +139,20 @@ class GcsPlacementGroupManager : public rpc::PlacementGroupInfoHandler {
   ///
   /// \param placement_group The placement_group that has been created.
   void OnPlacementGroupCreationSuccess(
-      std::shared_ptr<GcsPlacementGroup> placement_group);
+      const std::shared_ptr<GcsPlacementGroup> &placement_group);
 
  private:
-  /// Schedule another tick after a short time.
-  void ScheduleTick();
+  /// Try to create placement group after a short time.
+  void RetryCreatingPlacementGroup();
+
+  /// The io loop that is used to delay execution of tasks (e.g.,
+  /// execute_after).
+  boost::asio::io_context &io_context_;
 
   /// Callback of placement_group registration requests that are not yet flushed.
   absl::flat_hash_map<PlacementGroupID, EmptyCallback>
       placement_group_to_register_callback_;
+
   /// All registered placement_groups (pending placement_groups are also included).
   absl::flat_hash_map<PlacementGroupID, std::shared_ptr<GcsPlacementGroup>>
       registered_placement_groups_;
@@ -161,9 +166,6 @@ class GcsPlacementGroupManager : public rpc::PlacementGroupInfoHandler {
   std::shared_ptr<gcs::GcsTableStorage> gcs_table_storage_;
   /// If a placement group is creating
   bool is_creating_ = false;
-
-  /// A timer that ticks every schedule failure milliseconds.
-  boost::asio::deadline_timer reschedule_timer_;
 };
 
 }  // namespace gcs

src/ray/gcs/gcs_server/gcs_server.cc

@@ -197,7 +197,8 @@ void GcsServer::InitGcsActorManager() {
   gcs_node_manager_->AddNodeAddedListener(
       [this](const std::shared_ptr<rpc::GcsNodeInfo> &) {
         // Because a new node has been added, we need to try to schedule the pending
-        // actors.
+        // placement groups and the pending actors.
+        gcs_placement_group_manager_->SchedulePendingPlacementGroups();
         gcs_actor_manager_->SchedulePendingActors();
       });
 

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

@@ -59,11 +59,13 @@ class GcsPlacementGroupManagerTest : public ::testing::Test {
     thread_io_service_->join();
   }
 
+  std::shared_ptr<MockPlacementGroupScheduler> mock_placement_group_scheduler_;
+  std::unique_ptr<gcs::GcsPlacementGroupManager> gcs_placement_group_manager_;
+
+ private:
   std::unique_ptr<std::thread> thread_io_service_;
   boost::asio::io_service io_service_;
   std::shared_ptr<gcs::GcsTableStorage> gcs_table_storage_;
-  std::shared_ptr<MockPlacementGroupScheduler> mock_placement_group_scheduler_;
-  std::unique_ptr<gcs::GcsPlacementGroupManager> gcs_placement_group_manager_;
 };
 
 TEST_F(GcsPlacementGroupManagerTest, TestBasic) {
@@ -128,6 +130,25 @@ TEST_F(GcsPlacementGroupManagerTest, TestGetPlacementGroupIDByName) {
           create_placement_group_request.placement_group_spec().placement_group_id()));
 }
 
+TEST_F(GcsPlacementGroupManagerTest, TestRescheduleWhenNodeAdd) {
+  auto create_placement_group_request = Mocker::GenCreatePlacementGroupRequest();
+  std::atomic<int> finished_placement_group_count(0);
+  gcs_placement_group_manager_->RegisterPlacementGroup(
+      create_placement_group_request,
+      [&finished_placement_group_count]() { ++finished_placement_group_count; });
+  ASSERT_EQ(finished_placement_group_count, 0);
+  ASSERT_EQ(mock_placement_group_scheduler_->placement_groups.size(), 1);
+  auto placement_group = mock_placement_group_scheduler_->placement_groups.back();
+  mock_placement_group_scheduler_->placement_groups.pop_back();
+
+  // If the creation of placement group fails, it will be rescheduled after a short time.
+  gcs_placement_group_manager_->OnPlacementGroupCreationFailed(placement_group);
+  auto condition = [this]() {
+    return (int)mock_placement_group_scheduler_->placement_groups.size() == 1;
+  };
+  EXPECT_TRUE(WaitForCondition(condition, 10 * 1000));
+}
+
 }  // namespace ray
 
 int main(int argc, char **argv) {