[Placement Group] Support named placement group (#13755)

src/ray/gcs/accessor.h

@@ -727,7 +727,7 @@ class PlacementGroupInfoAccessor {
   virtual Status AsyncCreatePlacementGroup(
       const PlacementGroupSpecification &placement_group_spec) = 0;
 
-  /// Get a placement group data from GCS asynchronously.
+  /// Get a placement group data from GCS asynchronously by id.
   ///
   /// \param placement_group_id The id of a placement group to obtain from GCS.
   /// \return Status.
@@ -735,6 +735,14 @@ class PlacementGroupInfoAccessor {
       const PlacementGroupID &placement_group_id,
       const OptionalItemCallback<rpc::PlacementGroupTableData> &callback) = 0;
 
+  /// Get a placement group data from GCS asynchronously by name.
+  ///
+  /// \param placement_group_name The name of a placement group to obtain from GCS.
+  /// \return Status.
+  virtual Status AsyncGetByName(
+      const std::string &placement_group_name,
+      const OptionalItemCallback<rpc::PlacementGroupTableData> &callback) = 0;
+
   /// Get all placement group info from GCS asynchronously.
   ///
   /// \param callback Callback that will be called after lookup finished.

src/ray/gcs/gcs_client/global_state_accessor.cc

@@ -259,5 +259,17 @@ std::unique_ptr<std::string> GlobalStateAccessor::GetPlacementGroupInfo(
   return placement_group_table_data;
 }
 
+std::unique_ptr<std::string> GlobalStateAccessor::GetPlacementGroupByName(
+    const std::string &placement_group_name) {
+  std::unique_ptr<std::string> placement_group_table_data;
+  std::promise<bool> promise;
+  RAY_CHECK_OK(gcs_client_->PlacementGroups().AsyncGetByName(
+      placement_group_name,
+      TransformForOptionalItemCallback<rpc::PlacementGroupTableData>(
+          placement_group_table_data, promise)));
+  promise.get_future().get();
+  return placement_group_table_data;
+}
+
 }  // namespace gcs
 }  // namespace ray

src/ray/gcs/gcs_client/global_state_accessor.h

@@ -151,15 +151,24 @@ class GlobalStateAccessor {
   /// deserialized with protobuf function.
   std::vector<std::string> GetAllPlacementGroupInfo();
 
-  /// Get information of a placement group from GCS Service.
+  /// Get information of a placement group from GCS Service by ID.
   ///
-  /// \param placement_group The ID of placement group to look up in the GCS Service.
+  /// \param placement_group_id The ID of placement group to look up in the GCS Service.
   /// \return Placement group info. To support multi-language, we serialize each
   /// PlacementGroupTableData and return the serialized string. Where used, it needs to be
   /// deserialized with protobuf function.
   std::unique_ptr<std::string> GetPlacementGroupInfo(
       const PlacementGroupID &placement_group_id);
 
+  /// Get information of a placement group from GCS Service by name.
+  ///
+  /// \param placement_group_name The name of placement group to look up in the GCS
+  /// Service. \return Placement group info. To support multi-language, we serialize each
+  /// PlacementGroupTableData and return the serialized string. Where used, it needs to be
+  /// deserialized with protobuf function.
+  std::unique_ptr<std::string> GetPlacementGroupByName(
+      const std::string &placement_group_name);
+
  private:
   /// MultiItem transformation helper in template style.
   ///

src/ray/gcs/gcs_client/service_based_accessor.cc

@@ -1466,6 +1466,26 @@ Status ServiceBasedPlacementGroupInfoAccessor::AsyncGet(
   return Status::OK();
 }
 
+Status ServiceBasedPlacementGroupInfoAccessor::AsyncGetByName(
+    const std::string &name,
+    const OptionalItemCallback<rpc::PlacementGroupTableData> &callback) {
+  RAY_LOG(DEBUG) << "Getting named placement group info, name = " << name;
+  rpc::GetNamedPlacementGroupRequest request;
+  request.set_name(name);
+  client_impl_->GetGcsRpcClient().GetNamedPlacementGroup(
+      request, [name, callback](const Status &status,
+                                const rpc::GetNamedPlacementGroupReply &reply) {
+        if (reply.has_placement_group_table_data()) {
+          callback(status, reply.placement_group_table_data());
+        } else {
+          callback(status, boost::none);
+        }
+        RAY_LOG(DEBUG) << "Finished getting named placement group info, status = "
+                       << status << ", name = " << name;
+      });
+  return Status::OK();
+}
+
 Status ServiceBasedPlacementGroupInfoAccessor::AsyncGetAll(
     const MultiItemCallback<rpc::PlacementGroupTableData> &callback) {
   RAY_LOG(DEBUG) << "Getting all placement group info.";

src/ray/gcs/gcs_client/service_based_accessor.h

@@ -453,6 +453,10 @@ class ServiceBasedPlacementGroupInfoAccessor : public PlacementGroupInfoAccessor
       const PlacementGroupID &placement_group_id,
       const OptionalItemCallback<rpc::PlacementGroupTableData> &callback) override;
 
+  Status AsyncGetByName(
+      const std::string &name,
+      const OptionalItemCallback<rpc::PlacementGroupTableData> &callback) override;
+
   Status AsyncGetAll(
       const MultiItemCallback<rpc::PlacementGroupTableData> &callback) override;
 

src/ray/gcs/gcs_server/gcs_placement_group_manager.cc

@@ -65,7 +65,8 @@ rpc::PlacementStrategy GcsPlacementGroup::GetStrategy() const {
   return placement_group_table_data_.strategy();
 }
 
-const rpc::PlacementGroupTableData &GcsPlacementGroup::GetPlacementGroupTableData() {
+const rpc::PlacementGroupTableData &GcsPlacementGroup::GetPlacementGroupTableData()
+    const {
   return placement_group_table_data_;
 }
 
@@ -147,6 +148,21 @@ void GcsPlacementGroupManager::RegisterPlacementGroup(
     }
     return;
   }
+  if (!placement_group->GetName().empty()) {
+    auto it = named_placement_groups_.find(placement_group->GetName());
+    if (it == named_placement_groups_.end()) {
+      named_placement_groups_.emplace(placement_group->GetName(),
+                                      placement_group->GetPlacementGroupID());
+    } else {
+      std::stringstream stream;
+      stream << "Failed to create placement group '"
+             << placement_group->GetPlacementGroupID() << "' because name '"
+             << placement_group->GetName() << "' already exists.";
+      RAY_LOG(WARNING) << stream.str();
+      callback(Status::Invalid(stream.str()));
+      return;
+    }
+  }
 
   // Mark the callback as pending and invoke it after the placement_group has been
   // successfully created.
@@ -178,11 +194,9 @@ void GcsPlacementGroupManager::RegisterPlacementGroup(
 PlacementGroupID GcsPlacementGroupManager::GetPlacementGroupIDByName(
     const std::string &name) {
   PlacementGroupID placement_group_id = PlacementGroupID::Nil();
-  for (const auto &iter : registered_placement_groups_) {
-    if (iter.second->GetName() == name) {
-      placement_group_id = iter.first;
-      break;
-    }
+  auto it = named_placement_groups_.find(name);
+  if (it != named_placement_groups_.end()) {
+    placement_group_id = it->second;
   }
   return placement_group_id;
 }
@@ -315,10 +329,19 @@ void GcsPlacementGroupManager::RemovePlacementGroup(
     on_placement_group_removed(Status::OK());
     return;
   }
-  auto placement_group = placement_group_it->second;
+  auto placement_group = std::move(placement_group_it->second);
   registered_placement_groups_.erase(placement_group_it);
   placement_group_to_create_callbacks_.erase(placement_group_id);
 
+  // Remove placement group from `named_placement_groups_` if its name is not empty.
+  if (!placement_group->GetName().empty()) {
+    auto it = named_placement_groups_.find(placement_group->GetName());
+    if (it != named_placement_groups_.end() &&
+        it->second == placement_group->GetPlacementGroupID()) {
+      named_placement_groups_.erase(it);
+    }
+  }
+
   // Destroy all bundles.
   gcs_placement_group_scheduler_->DestroyPlacementGroupBundleResourcesIfExists(
       placement_group_id);
@@ -385,6 +408,30 @@ void GcsPlacementGroupManager::HandleGetPlacementGroup(
   ++counts_[CountType::GET_PLACEMENT_GROUP_REQUEST];
 }
 
+void GcsPlacementGroupManager::HandleGetNamedPlacementGroup(
+    const rpc::GetNamedPlacementGroupRequest &request,
+    rpc::GetNamedPlacementGroupReply *reply, rpc::SendReplyCallback send_reply_callback) {
+  const std::string &name = request.name();
+  RAY_LOG(DEBUG) << "Getting named placement group info, name = " << name;
+
+  // Try to look up the placement Group ID for the named placement group.
+  auto placement_group_id = GetPlacementGroupIDByName(name);
+
+  if (placement_group_id.IsNil()) {
+    // The placement group was not found.
+    RAY_LOG(DEBUG) << "Placement Group with name '" << name << "' was not found";
+  } else {
+    const auto &iter = registered_placement_groups_.find(placement_group_id);
+    RAY_CHECK(iter != registered_placement_groups_.end());
+    reply->mutable_placement_group_table_data()->CopyFrom(
+        iter->second->GetPlacementGroupTableData());
+    RAY_LOG(DEBUG) << "Finished get named placement group info, placement group id = "
+                   << placement_group_id;
+  }
+  GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
+  ++counts_[CountType::GET_NAMED_PLACEMENT_GROUP_REQUEST];
+}
+
 void GcsPlacementGroupManager::HandleGetAllPlacementGroup(
     const rpc::GetAllPlacementGroupRequest &request,
     rpc::GetAllPlacementGroupReply *reply, rpc::SendReplyCallback send_reply_callback) {
@@ -550,6 +597,10 @@ void GcsPlacementGroupManager::Initialize(const GcsInitData &gcs_init_data) {
     auto placement_group = std::make_shared<GcsPlacementGroup>(item.second);
     if (item.second.state() != rpc::PlacementGroupTableData::REMOVED) {
       registered_placement_groups_.emplace(item.first, placement_group);
+      if (!placement_group->GetName().empty()) {
+        named_placement_groups_.emplace(placement_group->GetName(),
+                                        placement_group->GetPlacementGroupID());
+      }
 
       if (item.second.state() == rpc::PlacementGroupTableData::PENDING ||
           item.second.state() == rpc::PlacementGroupTableData::RESCHEDULING) {
@@ -587,6 +638,7 @@ std::string GcsPlacementGroupManager::DebugString() const {
          << ", WaitPlacementGroupUntilReady request count: "
          << counts_[CountType::WAIT_PLACEMENT_GROUP_UNTIL_READY_REQUEST]
          << ", Registered placement groups count: " << registered_placement_groups_.size()
+         << ", Named placement group count: " << named_placement_groups_.size()
          << ", Pending placement groups count: " << pending_placement_groups_.size()
          << "}";
   return stream.str();

src/ray/gcs/gcs_server/gcs_placement_group_manager.h

@@ -65,7 +65,7 @@ class GcsPlacementGroup {
   }
 
   /// Get the immutable PlacementGroupTableData of this placement group.
-  const rpc::PlacementGroupTableData &GetPlacementGroupTableData();
+  const rpc::PlacementGroupTableData &GetPlacementGroupTableData() const;
 
   /// Get the mutable bundle of this placement group.
   rpc::Bundle *GetMutableBundle(int bundle_index);
@@ -155,10 +155,13 @@ class GcsPlacementGroupManager : public rpc::PlacementGroupInfoHandler {
                                rpc::GetPlacementGroupReply *reply,
                                rpc::SendReplyCallback send_reply_callback) override;
 
+  void HandleGetNamedPlacementGroup(const rpc::GetNamedPlacementGroupRequest &request,
+                                    rpc::GetNamedPlacementGroupReply *reply,
+                                    rpc::SendReplyCallback send_reply_callback) override;
+
   void HandleGetAllPlacementGroup(const rpc::GetAllPlacementGroupRequest &request,
                                   rpc::GetAllPlacementGroupReply *reply,
                                   rpc::SendReplyCallback send_reply_callback) override;
-
   void HandleWaitPlacementGroupUntilReady(
       const rpc::WaitPlacementGroupUntilReadyRequest &request,
       rpc::WaitPlacementGroupUntilReadyReply *reply,
@@ -315,6 +318,9 @@ class GcsPlacementGroupManager : public rpc::PlacementGroupInfoHandler {
   /// Reference of GcsResourceManager.
   GcsResourceManager &gcs_resource_manager_;
 
+  /// Maps placement group names to their placement group ID for lookups by name.
+  absl::flat_hash_map<std::string, PlacementGroupID> named_placement_groups_;
+
   // Debug info.
   enum CountType {
     CREATE_PLACEMENT_GROUP_REQUEST = 0,
@@ -322,7 +328,8 @@ class GcsPlacementGroupManager : public rpc::PlacementGroupInfoHandler {
     GET_PLACEMENT_GROUP_REQUEST = 2,
     GET_ALL_PLACEMENT_GROUP_REQUEST = 3,
     WAIT_PLACEMENT_GROUP_UNTIL_READY_REQUEST = 4,
-    CountType_MAX = 5,
+    GET_NAMED_PLACEMENT_GROUP_REQUEST = 5,
+    CountType_MAX = 6,
   };
   uint64_t counts_[CountType::CountType_MAX] = {0};
 };

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

@@ -174,6 +174,31 @@ TEST_F(GcsPlacementGroupManagerTest, TestGetPlacementGroupIDByName) {
       PlacementGroupID::FromBinary(request.placement_group_spec().placement_group_id()));
 }
 
+TEST_F(GcsPlacementGroupManagerTest, TestRemoveNamedPlacementGroup) {
+  auto request = Mocker::GenCreatePlacementGroupRequest("test_name");
+  std::atomic<int> finished_placement_group_count(0);
+  gcs_placement_group_manager_->RegisterPlacementGroup(
+      std::make_shared<gcs::GcsPlacementGroup>(request),
+      [&finished_placement_group_count](const Status &status) {
+        ++finished_placement_group_count;
+      });
+
+  ASSERT_EQ(finished_placement_group_count, 0);
+  WaitForExpectedPgCount(1);
+  auto placement_group = mock_placement_group_scheduler_->placement_groups_.back();
+  mock_placement_group_scheduler_->placement_groups_.pop_back();
+
+  gcs_placement_group_manager_->OnPlacementGroupCreationSuccess(placement_group);
+  WaitForExpectedCount(finished_placement_group_count, 1);
+  ASSERT_EQ(placement_group->GetState(), rpc::PlacementGroupTableData::CREATED);
+  // Remove the named placement group.
+  gcs_placement_group_manager_->RemovePlacementGroup(
+      placement_group->GetPlacementGroupID(),
+      [](const Status &status) { ASSERT_TRUE(status.ok()); });
+  ASSERT_EQ(gcs_placement_group_manager_->GetPlacementGroupIDByName("test_name"),
+            PlacementGroupID::Nil());
+}
+
 TEST_F(GcsPlacementGroupManagerTest, TestRescheduleWhenNodeAdd) {
   auto request = Mocker::GenCreatePlacementGroupRequest();
   std::atomic<int> finished_placement_group_count(0);

src/ray/protobuf/gcs_service.proto

@@ -504,6 +504,17 @@ message WaitPlacementGroupUntilReadyReply {
   GcsStatus status = 1;
 }
 
+message GetNamedPlacementGroupRequest {
+  // Name of the placement group.
+  string name = 1;
+}
+
+message GetNamedPlacementGroupReply {
+  GcsStatus status = 1;
+  // Data of placement group.
+  PlacementGroupTableData placement_group_table_data = 2;
+}
+
 // Service for placement group info access.
 service PlacementGroupInfoGcsService {
   // Create placement group via gcs service.
@@ -514,6 +525,9 @@ service PlacementGroupInfoGcsService {
       returns (RemovePlacementGroupReply);
   // Get placement group information via gcs service.
   rpc GetPlacementGroup(GetPlacementGroupRequest) returns (GetPlacementGroupReply);
+  // Get named placement group information via gcs service.
+  rpc GetNamedPlacementGroup(GetNamedPlacementGroupRequest)
+      returns (GetNamedPlacementGroupReply);
   // Get information of all placement group from GCS Service.
   rpc GetAllPlacementGroup(GetAllPlacementGroupRequest)
       returns (GetAllPlacementGroupReply);

src/ray/rpc/gcs_server/gcs_rpc_client.h

@@ -254,6 +254,10 @@ class GcsRpcClient {
   VOID_GCS_RPC_CLIENT_METHOD(PlacementGroupInfoGcsService, GetPlacementGroup,
                              placement_group_info_grpc_client_, )
 
+  /// Get placement group data from GCS Service by name.
+  VOID_GCS_RPC_CLIENT_METHOD(PlacementGroupInfoGcsService, GetNamedPlacementGroup,
+                             placement_group_info_grpc_client_, )
+
   /// Get information of all placement group from GCS Service.
   VOID_GCS_RPC_CLIENT_METHOD(PlacementGroupInfoGcsService, GetAllPlacementGroup,
                              placement_group_info_grpc_client_, )

src/ray/rpc/gcs_server/gcs_rpc_server.h

@@ -522,6 +522,10 @@ class PlacementGroupInfoGcsServiceHandler {
       const WaitPlacementGroupUntilReadyRequest &request,
       WaitPlacementGroupUntilReadyReply *reply,
       SendReplyCallback send_reply_callback) = 0;
+
+  virtual void HandleGetNamedPlacementGroup(const GetNamedPlacementGroupRequest &request,
+                                            GetNamedPlacementGroupReply *reply,
+                                            SendReplyCallback send_reply_callback) = 0;
 };
 
 /// The `GrpcService` for `PlacementGroupInfoGcsService`.
@@ -543,6 +547,7 @@ class PlacementGroupInfoGrpcService : public GrpcService {
     PLACEMENT_GROUP_INFO_SERVICE_RPC_HANDLER(CreatePlacementGroup);
     PLACEMENT_GROUP_INFO_SERVICE_RPC_HANDLER(RemovePlacementGroup);
     PLACEMENT_GROUP_INFO_SERVICE_RPC_HANDLER(GetPlacementGroup);
+    PLACEMENT_GROUP_INFO_SERVICE_RPC_HANDLER(GetNamedPlacementGroup);
     PLACEMENT_GROUP_INFO_SERVICE_RPC_HANDLER(GetAllPlacementGroup);
     PLACEMENT_GROUP_INFO_SERVICE_RPC_HANDLER(WaitPlacementGroupUntilReady);
   }