[GCS]profile info getting implementation based gcs service (#8536)

src/ray/gcs/accessor.h

@@ -570,6 +570,13 @@ class StatsInfoAccessor {
       const std::shared_ptr<rpc::ProfileTableData> &data_ptr,
       const StatusCallback &callback) = 0;
 
+  /// Get all profile info from GCS asynchronously.
+  ///
+  /// \param callback Callback that will be called after lookup finished.
+  /// \return Status
+  virtual Status AsyncGetAll(
+      const MultiItemCallback<rpc::ProfileTableData> &callback) = 0;
+
  protected:
   StatsInfoAccessor() = default;
 };

src/ray/gcs/gcs_client/global_state_accessor.cc

@@ -80,5 +80,22 @@ std::vector<std::string> GlobalStateAccessor::GetAllJobInfo() {
   return job_table_data;
 }
 
+std::vector<std::string> GlobalStateAccessor::GetAllProfileInfo() {
+  std::vector<std::string> profile_table_data;
+  std::promise<bool> promise;
+  auto on_done = [&profile_table_data, &promise](
+                     const Status &status,
+                     const std::vector<rpc::ProfileTableData> &result) {
+    RAY_CHECK_OK(status);
+    for (auto &data : result) {
+      profile_table_data.push_back(data.SerializeAsString());
+    }
+    promise.set_value(true);
+  };
+  RAY_CHECK_OK(gcs_client_->Stats().AsyncGetAll(on_done));
+  promise.get_future().get();
+  return profile_table_data;
+}
+
 }  // namespace gcs
 }  // namespace ray

src/ray/gcs/gcs_client/global_state_accessor.h

@@ -51,6 +51,13 @@ class GlobalStateAccessor {
   /// protobuf function.
   std::vector<std::string> GetAllJobInfo();
 
+  /// Get information of all profiles from GCS Service.
+  ///
+  /// \return All profile info. To support multi-language, we serialized each
+  /// ProfileTableData and returned the serialized string. Where used, it needs to be
+  /// deserialized with protobuf function.
+  std::vector<std::string> GetAllProfileInfo();
+
  private:
   /// Whether this client is connected to gcs server.
   bool is_connected_{false};

src/ray/gcs/gcs_client/service_based_accessor.cc

@@ -1068,6 +1068,21 @@ Status ServiceBasedStatsInfoAccessor::AsyncAddProfileData(
   return Status::OK();
 }
 
+Status ServiceBasedStatsInfoAccessor::AsyncGetAll(
+    const MultiItemCallback<rpc::ProfileTableData> &callback) {
+  RAY_LOG(DEBUG) << "Getting all profile info.";
+  RAY_CHECK(callback);
+  rpc::GetAllProfileInfoRequest request;
+  client_impl_->GetGcsRpcClient().GetAllProfileInfo(
+      request,
+      [callback](const Status &status, const rpc::GetAllProfileInfoReply &reply) {
+        auto result = VectorFromProtobuf(reply.profile_info_list());
+        callback(status, result);
+        RAY_LOG(DEBUG) << "Finished getting all job info.";
+      });
+  return Status::OK();
+}
+
 ServiceBasedErrorInfoAccessor::ServiceBasedErrorInfoAccessor(
     ServiceBasedGcsClient *client_impl)
     : client_impl_(client_impl) {}

src/ray/gcs/gcs_client/service_based_accessor.h

@@ -281,6 +281,8 @@ class ServiceBasedStatsInfoAccessor : public StatsInfoAccessor {
   Status AsyncAddProfileData(const std::shared_ptr<rpc::ProfileTableData> &data_ptr,
                              const StatusCallback &callback) override;
 
+  Status AsyncGetAll(const MultiItemCallback<rpc::ProfileTableData> &callback) override;
+
  private:
   ServiceBasedGcsClient *client_impl_;
 };

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

@@ -105,6 +105,21 @@ TEST_F(GlobalStateAccessorTest, TestJobTable) {
   ASSERT_EQ(global_state_->GetAllJobInfo().size(), job_count);
 }
 
+TEST_F(GlobalStateAccessorTest, TestProfileTable) {
+  int profile_count = 100;
+  ASSERT_EQ(global_state_->GetAllProfileInfo().size(), 0);
+  for (int index = 0; index < profile_count; ++index) {
+    auto client_id = ClientID::FromRandom();
+    auto profile_table_data = Mocker::GenProfileTableData(client_id);
+    std::promise<bool> promise;
+    RAY_CHECK_OK(gcs_client_->Stats().AsyncAddProfileData(
+        profile_table_data,
+        [&promise](Status status) { promise.set_value(status.ok()); }));
+    WaitReady(promise.get_future(), timeout_ms_);
+  }
+  ASSERT_EQ(global_state_->GetAllProfileInfo().size(), profile_count);
+}
+
 }  // namespace ray
 
 int main(int argc, char **argv) {

src/ray/gcs/gcs_server/gcs_server.cc

@@ -229,7 +229,7 @@ std::unique_ptr<rpc::TaskInfoHandler> GcsServer::InitTaskInfoHandler() {
 
 std::unique_ptr<rpc::StatsHandler> GcsServer::InitStatsHandler() {
   return std::unique_ptr<rpc::DefaultStatsHandler>(
-      new rpc::DefaultStatsHandler(*redis_gcs_client_));
+      new rpc::DefaultStatsHandler(gcs_table_storage_));
 }
 
 std::unique_ptr<rpc::ErrorInfoHandler> GcsServer::InitErrorInfoHandler() {

src/ray/gcs/gcs_server/stats_handler_impl.cc

@@ -34,7 +34,8 @@ void DefaultStatsHandler::HandleAddProfileData(const AddProfileDataRequest &requ
     GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
   };
 
-  Status status = gcs_client_.Stats().AsyncAddProfileData(profile_table_data, on_done);
+  Status status = gcs_table_storage_->ProfileTable().Put(UniqueID::FromRandom(),
+                                                         *profile_table_data, on_done);
   if (!status.ok()) {
     on_done(status);
   }
@@ -42,5 +43,24 @@ void DefaultStatsHandler::HandleAddProfileData(const AddProfileDataRequest &requ
                  << request.profile_data().component_type() << ", node id = " << node_id;
 }
 
+void DefaultStatsHandler::HandleGetAllProfileInfo(
+    const rpc::GetAllProfileInfoRequest &request, rpc::GetAllProfileInfoReply *reply,
+    rpc::SendReplyCallback send_reply_callback) {
+  RAY_LOG(DEBUG) << "Getting all profile info.";
+  auto on_done = [reply, send_reply_callback](
+                     const std::unordered_map<UniqueID, ProfileTableData> &result) {
+    for (auto &data : result) {
+      reply->add_profile_info_list()->CopyFrom(data.second);
+    }
+    RAY_LOG(DEBUG) << "Finished getting all profile info.";
+    GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
+  };
+
+  Status status = gcs_table_storage_->ProfileTable().GetAll(on_done);
+  if (!status.ok()) {
+    on_done(std::unordered_map<UniqueID, ProfileTableData>());
+  }
+}
+
 }  // namespace rpc
 }  // namespace ray

src/ray/gcs/gcs_server/stats_handler_impl.h

@@ -15,6 +15,7 @@
 #ifndef RAY_GCS_STATS_HANDLER_IMPL_H
 #define RAY_GCS_STATS_HANDLER_IMPL_H
 
+#include "gcs_table_storage.h"
 #include "ray/gcs/redis_gcs_client.h"
 #include "ray/rpc/gcs_server/gcs_rpc_server.h"
 
@@ -24,15 +25,19 @@ namespace rpc {
 /// This implementation class of `StatsHandler`.
 class DefaultStatsHandler : public rpc::StatsHandler {
  public:
-  explicit DefaultStatsHandler(gcs::RedisGcsClient &gcs_client)
-      : gcs_client_(gcs_client) {}
+  explicit DefaultStatsHandler(std::shared_ptr<gcs::GcsTableStorage> gcs_table_storage)
+      : gcs_table_storage_(std::move(gcs_table_storage)) {}
 
   void HandleAddProfileData(const AddProfileDataRequest &request,
                             AddProfileDataReply *reply,
                             SendReplyCallback send_reply_callback) override;
 
+  void HandleGetAllProfileInfo(const rpc::GetAllProfileInfoRequest &request,
+                               rpc::GetAllProfileInfoReply *reply,
+                               rpc::SendReplyCallback send_reply_callback) override;
+
  private:
-  gcs::RedisGcsClient &gcs_client_;
+  std::shared_ptr<gcs::GcsTableStorage> gcs_table_storage_;
 };
 
 }  // namespace rpc

src/ray/gcs/redis_accessor.h

@@ -404,6 +404,10 @@ class RedisStatsInfoAccessor : public StatsInfoAccessor {
   Status AsyncAddProfileData(const std::shared_ptr<ProfileTableData> &data_ptr,
                              const StatusCallback &callback) override;
 
+  Status AsyncGetAll(const MultiItemCallback<rpc::ProfileTableData> &callback) override {
+    return Status::NotImplemented("AsyncGetAll not implemented");
+  }
+
  private:
   RedisGcsClient *client_impl_{nullptr};
 };

src/ray/protobuf/gcs_service.proto

@@ -368,10 +368,20 @@ message AddProfileDataReply {
   GcsStatus status = 1;
 }
 
+message GetAllProfileInfoRequest {
+}
+
+message GetAllProfileInfoReply {
+  GcsStatus status = 1;
+  repeated ProfileTableData profile_info_list = 2;
+}
+
 // Service for stats access.
 service StatsGcsService {
   // Add profile data to GCS Service.
   rpc AddProfileData(AddProfileDataRequest) returns (AddProfileDataReply);
+  // Get information of all profiles from GCS Service.
+  rpc GetAllProfileInfo(GetAllProfileInfoRequest) returns (GetAllProfileInfoReply);
 }
 
 message ReportJobErrorRequest {

src/ray/rpc/gcs_server/gcs_rpc_client.h

@@ -191,6 +191,9 @@ class GcsRpcClient {
   /// Add profile data to GCS Service.
   VOID_GCS_RPC_CLIENT_METHOD(StatsGcsService, AddProfileData, stats_grpc_client_, )
 
+  /// Get information of all profiles from GCS Service.
+  VOID_GCS_RPC_CLIENT_METHOD(StatsGcsService, GetAllProfileInfo, stats_grpc_client_, )
+
   /// Report a job error to GCS Service.
   VOID_GCS_RPC_CLIENT_METHOD(ErrorInfoGcsService, ReportJobError,
                              error_info_grpc_client_, )

src/ray/rpc/gcs_server/gcs_rpc_server.h

@@ -345,6 +345,10 @@ class StatsGcsServiceHandler {
   virtual void HandleAddProfileData(const AddProfileDataRequest &request,
                                     AddProfileDataReply *reply,
                                     SendReplyCallback send_reply_callback) = 0;
+
+  virtual void HandleGetAllProfileInfo(const GetAllProfileInfoRequest &request,
+                                       GetAllProfileInfoReply *reply,
+                                       SendReplyCallback send_reply_callback) = 0;
 };
 
 /// The `GrpcService` for `StatsGcsService`.
@@ -364,6 +368,7 @@ class StatsGrpcService : public GrpcService {
       const std::unique_ptr<grpc::ServerCompletionQueue> &cq,
       std::vector<std::unique_ptr<ServerCallFactory>> *server_call_factories) override {
     STATS_SERVICE_RPC_HANDLER(AddProfileData);
+    STATS_SERVICE_RPC_HANDLER(GetAllProfileInfo);
   }
 
  private: