[Placement Group]Add get all placement group api (#11460)

* add get all interface for placement group

* add get all interface for placement group

* make it work

* fix lint

* fix lint

* fix comment

* add cpp test

* fix python lint

src/ray/gcs/accessor.h

@@ -769,6 +769,13 @@ class PlacementGroupInfoAccessor {
       const PlacementGroupID &placement_group_id,
       const OptionalItemCallback<rpc::PlacementGroupTableData> &callback) = 0;
 
+  /// Get all placement group info from GCS asynchronously.
+  ///
+  /// \param callback Callback that will be called after lookup finished.
+  /// \return Status
+  virtual Status AsyncGetAll(
+      const MultiItemCallback<rpc::PlacementGroupTableData> &callback) = 0;
+
   /// Remove a placement group to GCS synchronously.
   ///
   /// \param placement_group_id The id for the placement group to remove.

src/ray/gcs/gcs_client/global_state_accessor.cc

@@ -237,6 +237,16 @@ bool GlobalStateAccessor::AddWorkerInfo(const std::string &serialized_string) {
   return true;
 }
 
+std::vector<std::string> GlobalStateAccessor::GetAllPlacementGroupInfo() {
+  std::vector<std::string> placement_group_table_data;
+  std::promise<bool> promise;
+  RAY_CHECK_OK(gcs_client_->PlacementGroups().AsyncGetAll(
+      TransformForMultiItemCallback<rpc::PlacementGroupTableData>(
+          placement_group_table_data, promise)));
+  promise.get_future().get();
+  return placement_group_table_data;
+}
+
 std::unique_ptr<std::string> GlobalStateAccessor::GetPlacementGroupInfo(
     const PlacementGroupID &placement_group_id) {
   std::unique_ptr<std::string> placement_group_table_data;

src/ray/gcs/gcs_client/global_state_accessor.h

@@ -144,6 +144,13 @@ class GlobalStateAccessor {
   /// \return Is operation success.
   bool AddWorkerInfo(const std::string &serialized_string);
 
+  /// Get information of all placement group from GCS Service.
+  ///
+  /// \return All 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::vector<std::string> GetAllPlacementGroupInfo();
+
   /// Get information of a placement group from GCS Service.
   ///
   /// \param placement_group The ID of placement group to look up in the GCS Service.

src/ray/gcs/gcs_client/service_based_accessor.cc

@@ -1545,5 +1545,19 @@ Status ServiceBasedPlacementGroupInfoAccessor::AsyncGet(
   return Status::OK();
 }
 
+Status ServiceBasedPlacementGroupInfoAccessor::AsyncGetAll(
+    const MultiItemCallback<rpc::PlacementGroupTableData> &callback) {
+  RAY_LOG(DEBUG) << "Getting all placement group info.";
+  rpc::GetAllPlacementGroupRequest request;
+  client_impl_->GetGcsRpcClient().GetAllPlacementGroup(
+      request,
+      [callback](const Status &status, const rpc::GetAllPlacementGroupReply &reply) {
+        callback(status, VectorFromProtobuf(reply.placement_group_table_data()));
+        RAY_LOG(DEBUG) << "Finished getting all placement group info, status = "
+                       << status;
+      });
+  return Status::OK();
+}
+
 }  // namespace gcs
 }  // namespace ray

src/ray/gcs/gcs_client/service_based_accessor.h

@@ -460,6 +460,9 @@ class ServiceBasedPlacementGroupInfoAccessor : public PlacementGroupInfoAccessor
       const PlacementGroupID &placement_group_id,
       const OptionalItemCallback<rpc::PlacementGroupTableData> &callback) override;
 
+  Status AsyncGetAll(
+      const MultiItemCallback<rpc::PlacementGroupTableData> &callback) override;
+
  private:
   ServiceBasedGcsClient *client_impl_;
 };

src/ray/gcs/gcs_client/test/global_state_accessor_test.cc

@@ -250,6 +250,9 @@ TEST_F(GlobalStateAccessorTest, TestWorkerTable) {
 }
 
 // TODO(sang): Add tests after adding asyncAdd
+TEST_F(GlobalStateAccessorTest, TestPlacementGroupTable) {
+  ASSERT_EQ(global_state_->GetAllPlacementGroupInfo().size(), 0);
+}
 
 }  // namespace ray
 

src/ray/gcs/gcs_server/gcs_placement_group_manager.cc

@@ -328,6 +328,25 @@ void GcsPlacementGroupManager::HandleGetPlacementGroup(
   }
 }
 
+void GcsPlacementGroupManager::HandleGetAllPlacementGroup(
+    const rpc::GetAllPlacementGroupRequest &request,
+    rpc::GetAllPlacementGroupReply *reply, rpc::SendReplyCallback send_reply_callback) {
+  RAY_LOG(DEBUG) << "Getting all placement group info.";
+  auto on_done =
+      [reply, send_reply_callback](
+          const std::unordered_map<PlacementGroupID, PlacementGroupTableData> &result) {
+        for (auto &data : result) {
+          reply->add_placement_group_table_data()->CopyFrom(data.second);
+        }
+        RAY_LOG(DEBUG) << "Finished getting all placement group info.";
+        GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
+      };
+  Status status = gcs_table_storage_->PlacementGroupTable().GetAll(on_done);
+  if (!status.ok()) {
+    on_done(std::unordered_map<PlacementGroupID, PlacementGroupTableData>());
+  }
+}
+
 void GcsPlacementGroupManager::RetryCreatingPlacementGroup() {
   execute_after(io_context_, [this] { SchedulePendingPlacementGroups(); },
                 RayConfig::instance().gcs_create_placement_group_retry_interval_ms());

src/ray/gcs/gcs_server/gcs_placement_group_manager.h

@@ -129,6 +129,10 @@ class GcsPlacementGroupManager : public rpc::PlacementGroupInfoHandler {
                                rpc::GetPlacementGroupReply *reply,
                                rpc::SendReplyCallback send_reply_callback) override;
 
+  void HandleGetAllPlacementGroup(const rpc::GetAllPlacementGroupRequest &request,
+                                  rpc::GetAllPlacementGroupReply *reply,
+                                  rpc::SendReplyCallback send_reply_callback) override;
+
   /// Register placement_group asynchronously.
   ///
   /// \param placement_group The placement group to be created.

src/ray/gcs/redis_accessor.cc

@@ -781,6 +781,11 @@ Status RedisPlacementGroupInfoAccessor::AsyncGet(
   return Status::Invalid("Not implemented");
 }
 
+Status RedisPlacementGroupInfoAccessor::AsyncGetAll(
+    const MultiItemCallback<rpc::PlacementGroupTableData> &callback) {
+  return Status::Invalid("Not implemented");
+}
+
 }  // namespace gcs
 
 }  // namespace ray

src/ray/gcs/redis_accessor.h

@@ -490,6 +490,9 @@ class RedisPlacementGroupInfoAccessor : public PlacementGroupInfoAccessor {
   Status AsyncGet(
       const PlacementGroupID &placement_group_id,
       const OptionalItemCallback<rpc::PlacementGroupTableData> &callback) override;
+
+  Status AsyncGetAll(
+      const MultiItemCallback<rpc::PlacementGroupTableData> &callback) override;
 };
 
 }  // namespace gcs

src/ray/protobuf/gcs_service.proto

@@ -164,20 +164,6 @@ service ActorInfoGcsService {
       returns (GetActorCheckpointIDReply);
 }
 
-// Service for placement group info access.
-service PlacementGroupInfoGcsService {
-  // Create placement group via gcs service.
-  rpc CreatePlacementGroup(CreatePlacementGroupRequest)
-      returns (CreatePlacementGroupReply);
-
-  // Remove placement group via gcs service.
-  rpc RemovePlacementGroup(RemovePlacementGroupRequest)
-      returns (RemovePlacementGroupReply);
-
-  // Get placement group information via gcs service.
-  rpc GetPlacementGroup(GetPlacementGroupRequest) returns (GetPlacementGroupReply);
-}
-
 message RegisterNodeRequest {
   // Info of node.
   GcsNodeInfo node_info = 1;
@@ -543,3 +529,30 @@ enum GcsServiceFailureType {
   RPC_DISCONNECT = 0;
   GCS_SERVER_RESTART = 1;
 }
+
+message GetAllPlacementGroupRequest {
+}
+
+message GetAllPlacementGroupReply {
+  GcsStatus status = 1;
+  // Data of placement group
+  repeated PlacementGroupTableData placement_group_table_data = 2;
+}
+
+// Service for placement group info access.
+service PlacementGroupInfoGcsService {
+  // Create placement group via gcs service.
+  rpc CreatePlacementGroup(CreatePlacementGroupRequest)
+      returns (CreatePlacementGroupReply);
+
+  // Remove placement group via gcs service.
+  rpc RemovePlacementGroup(RemovePlacementGroupRequest)
+      returns (RemovePlacementGroupReply);
+
+  // Get placement group information via gcs service.
+  rpc GetPlacementGroup(GetPlacementGroupRequest) returns (GetPlacementGroupReply);
+
+  // Get information of all placement group from GCS Service.
+  rpc GetAllPlacementGroup(GetAllPlacementGroupRequest)
+      returns (GetAllPlacementGroupReply);
+}

src/ray/rpc/gcs_server/gcs_rpc_client.h

@@ -262,6 +262,10 @@ class GcsRpcClient {
   VOID_GCS_RPC_CLIENT_METHOD(PlacementGroupInfoGcsService, GetPlacementGroup,
                              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_, )
+
  private:
   std::function<void(GcsServiceFailureType)> gcs_service_failure_detected_;
 

src/ray/rpc/gcs_server/gcs_rpc_server.h

@@ -470,6 +470,10 @@ class PlacementGroupInfoGcsServiceHandler {
   virtual void HandleGetPlacementGroup(const GetPlacementGroupRequest &request,
                                        GetPlacementGroupReply *reply,
                                        SendReplyCallback send_reply_callback) = 0;
+
+  virtual void HandleGetAllPlacementGroup(const GetAllPlacementGroupRequest &request,
+                                          GetAllPlacementGroupReply *reply,
+                                          SendReplyCallback send_reply_callback) = 0;
 };
 
 /// The `GrpcService` for `PlacementGroupInfoGcsService`.
@@ -491,6 +495,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(GetAllPlacementGroup);
   }
 
  private: