[PlacementGroup]Fix placement group wait api disorder bug (#12827)

* [PlacementGroup]Fix placment group wait api disorder bug

* fix review comment

* fix review comment

* fix review comment

* fix review comments

* increase num_heartbeats_timeout

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

src/ray/core_worker/core_worker.cc

@@ -1460,14 +1460,14 @@ Status CoreWorker::RemovePlacementGroup(const PlacementGroupID &placement_group_
 }
 
 Status CoreWorker::WaitPlacementGroupReady(const PlacementGroupID &placement_group_id,
-                                           int timeout_ms) {
+                                           int timeout_seconds) {
   std::shared_ptr<std::promise<Status>> status_promise =
       std::make_shared<std::promise<Status>>();
   RAY_CHECK_OK(gcs_client_->PlacementGroups().AsyncWaitUntilReady(
       placement_group_id,
       [status_promise](const Status &status) { status_promise->set_value(status); }));
   auto status_future = status_promise->get_future();
-  if (status_future.wait_for(std::chrono::milliseconds(timeout_ms)) !=
+  if (status_future.wait_for(std::chrono::seconds(timeout_seconds)) !=
       std::future_status::ready) {
     std::ostringstream stream;
     stream << "There was timeout in waiting for placement group " << placement_group_id

src/ray/core_worker/core_worker.h

@@ -682,11 +682,11 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
   /// Returns once the placement group is created or the timeout expires.
   ///
   /// \param placement_group The id of a placement group to wait for.
-  /// \param timeout_ms Timeout in milliseconds.
+  /// \param timeout_seconds Timeout in seconds.
   /// \return Status OK if the placement group is created. TimedOut if request to GCS
   /// server times out. NotFound if placement group is already removed or doesn't exist.
   Status WaitPlacementGroupReady(const PlacementGroupID &placement_group_id,
-                                 int timeout_ms);
+                                 int timeout_seconds);
 
   /// Submit an actor task.
   ///

src/ray/core_worker/lib/java/io_ray_runtime_task_NativeTaskSubmitter.cc

@@ -297,11 +297,11 @@ Java_io_ray_runtime_task_NativeTaskSubmitter_nativeRemovePlacementGroup(
 
 JNIEXPORT jboolean JNICALL
 Java_io_ray_runtime_task_NativeTaskSubmitter_nativeWaitPlacementGroupReady(
-    JNIEnv *env, jclass p, jbyteArray placement_group_id_bytes, jint timeout_ms) {
+    JNIEnv *env, jclass p, jbyteArray placement_group_id_bytes, jint timeout_seconds) {
   const auto placement_group_id =
       JavaByteArrayToId<ray::PlacementGroupID>(env, placement_group_id_bytes);
   auto status = ray::CoreWorkerProcess::GetCoreWorker().WaitPlacementGroupReady(
-      placement_group_id, timeout_ms);
+      placement_group_id, timeout_seconds);
   if (status.IsNotFound()) {
     env->ThrowNew(java_ray_exception_class, status.message().c_str());
   }

src/ray/gcs/gcs_server/gcs_placement_group_manager.cc

@@ -410,26 +410,48 @@ void GcsPlacementGroupManager::HandleWaitPlacementGroupUntilReady(
   RAY_LOG(DEBUG) << "Waiting for placement group until ready, placement group id = "
                  << placement_group_id;
 
+  auto callback = [placement_group_id, reply, send_reply_callback](const Status &status) {
+    RAY_LOG(DEBUG)
+        << "Finished waiting for placement group until ready, placement group id = "
+        << placement_group_id;
+    GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
+  };
+
   // If the placement group does not exist or it has been successfully created, return
   // directly.
   const auto &iter = registered_placement_groups_.find(placement_group_id);
   if (iter == registered_placement_groups_.end()) {
-    RAY_LOG(DEBUG) << "Placement group is not exist, placement group id = "
-                   << placement_group_id;
-    GCS_RPC_SEND_REPLY(send_reply_callback, reply,
-                       Status::NotFound("Placement group is not exist."));
+    // Check whether the placement group does not exist or is removed.
+    auto on_done = [this, placement_group_id, reply, callback, send_reply_callback](
+                       const Status &status,
+                       const boost::optional<PlacementGroupTableData> &result) {
+      if (result) {
+        RAY_LOG(DEBUG) << "Placement group is removed, placement group id = "
+                       << placement_group_id;
+        GCS_RPC_SEND_REPLY(send_reply_callback, reply,
+                           Status::NotFound("Placement group is removed."));
+      } else {
+        // `wait` is a method of placement group object. Placement group object is
+        // obtained by create placement group api, so it can guarantee the existence of
+        // placement group.
+        // GCS client does not guarantee the order of placement group creation and
+        // wait, so GCS may call wait placement group first and then create placement
+        // group.
+        placement_group_to_create_callbacks_[placement_group_id].emplace_back(
+            std::move(callback));
+      }
+    };
+
+    Status status =
+        gcs_table_storage_->PlacementGroupTable().Get(placement_group_id, on_done);
+    if (!status.ok()) {
+      on_done(status, boost::none);
+    }
   } else if (iter->second->GetState() == rpc::PlacementGroupTableData::CREATED) {
     RAY_LOG(DEBUG) << "Placement group is created, placement group id = "
                    << placement_group_id;
     GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
   } else {
-    auto callback = [placement_group_id, reply,
-                     send_reply_callback](const Status &status) {
-      RAY_LOG(DEBUG)
-          << "Finished waiting for placement group until ready, placement group id = "
-          << placement_group_id;
-      GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
-    };
     placement_group_to_create_callbacks_[placement_group_id].emplace_back(
         std::move(callback));
   }