Add error info handler to gcs service (#6754)

* add error info accessor

* rebase master

* add function comments

* capture type instead of request

src/ray/gcs/gcs_server/error_info_handler_impl.cc

@@ -0,0 +1,31 @@
+#include "error_info_handler_impl.h"
+
+namespace ray {
+namespace rpc {
+
+void DefaultErrorInfoHandler::HandleReportJobError(
+    const ReportJobErrorRequest &request, ReportJobErrorReply *reply,
+    SendReplyCallback send_reply_callback) {
+  JobID job_id = JobID::FromBinary(request.error_data().job_id());
+  std::string type = request.error_data().type();
+  RAY_LOG(DEBUG) << "Reporting job error, job id = " << job_id << ", type = " << type;
+  auto error_table_data = std::make_shared<ErrorTableData>();
+  error_table_data->CopyFrom(request.error_data());
+  auto on_done = [job_id, type, send_reply_callback](Status status) {
+    if (!status.ok()) {
+      RAY_LOG(ERROR) << "Failed to report job error, job id = " << job_id
+                     << ", type = " << type;
+    }
+    send_reply_callback(status, nullptr, nullptr);
+  };
+
+  Status status = gcs_client_.Errors().AsyncReportJobError(error_table_data, on_done);
+  if (!status.ok()) {
+    on_done(status);
+  }
+  RAY_LOG(DEBUG) << "Finished reporting job error, job id = " << job_id
+                 << ", type = " << type;
+}
+
+}  // namespace rpc
+}  // namespace ray

src/ray/gcs/gcs_server/error_info_handler_impl.h

@@ -0,0 +1,27 @@
+#ifndef RAY_GCS_ERROR_INFO_HANDLER_IMPL_H
+#define RAY_GCS_ERROR_INFO_HANDLER_IMPL_H
+
+#include "ray/gcs/redis_gcs_client.h"
+#include "ray/rpc/gcs_server/gcs_rpc_server.h"
+
+namespace ray {
+namespace rpc {
+
+/// This implementation class of `ErrorInfoHandler`.
+class DefaultErrorInfoHandler : public rpc::ErrorInfoHandler {
+ public:
+  explicit DefaultErrorInfoHandler(gcs::RedisGcsClient &gcs_client)
+      : gcs_client_(gcs_client) {}
+
+  void HandleReportJobError(const ReportJobErrorRequest &request,
+                            ReportJobErrorReply *reply,
+                            SendReplyCallback send_reply_callback) override;
+
+ private:
+  gcs::RedisGcsClient &gcs_client_;
+};
+
+}  // namespace rpc
+}  // namespace ray
+
+#endif  // RAY_GCS_ERROR_INFO_HANDLER_IMPL_H

src/ray/gcs/gcs_server/gcs_server.cc

@@ -1,5 +1,6 @@
 #include "gcs_server.h"
 #include "actor_info_handler_impl.h"
+#include "error_info_handler_impl.h"
 #include "job_info_handler_impl.h"
 #include "node_info_handler_impl.h"
 #include "object_info_handler_impl.h"
@@ -49,6 +50,11 @@ void GcsServer::Start() {
   stats_service_.reset(new rpc::StatsGrpcService(main_service_, *stats_handler_));
   rpc_server_.RegisterService(*stats_service_);
 
+  error_info_handler_ = InitErrorInfoHandler();
+  error_info_service_.reset(
+      new rpc::ErrorInfoGrpcService(main_service_, *error_info_handler_));
+  rpc_server_.RegisterService(*error_info_service_);
+
   // Run rpc server.
   rpc_server_.Run();
 
@@ -105,5 +111,10 @@ std::unique_ptr<rpc::StatsHandler> GcsServer::InitStatsHandler() {
       new rpc::DefaultStatsHandler(*redis_gcs_client_));
 }
 
+std::unique_ptr<rpc::ErrorInfoHandler> GcsServer::InitErrorInfoHandler() {
+  return std::unique_ptr<rpc::DefaultErrorInfoHandler>(
+      new rpc::DefaultErrorInfoHandler(*redis_gcs_client_));
+}
+
 }  // namespace gcs
 }  // namespace ray

src/ray/gcs/gcs_server/gcs_server.h

@@ -62,6 +62,9 @@ class GcsServer {
   /// The stats handler
   virtual std::unique_ptr<rpc::StatsHandler> InitStatsHandler();
 
+  /// The error info handler
+  virtual std::unique_ptr<rpc::ErrorInfoHandler> InitErrorInfoHandler();
+
  private:
   /// Gcs server configuration
   GcsServerConfig config_;
@@ -87,6 +90,9 @@ class GcsServer {
   /// Stats handler and service
   std::unique_ptr<rpc::StatsHandler> stats_handler_;
   std::unique_ptr<rpc::StatsGrpcService> stats_service_;
+  /// Error info handler and service
+  std::unique_ptr<rpc::ErrorInfoHandler> error_info_handler_;
+  std::unique_ptr<rpc::ErrorInfoGrpcService> error_info_service_;
   /// Backend client
   std::shared_ptr<RedisGcsClient> redis_gcs_client_;
 };

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

@@ -355,6 +355,16 @@ class GcsServerTest : public RedisServiceManagerForTest {
     return WaitReady(promise.get_future(), timeout_ms_);
   }
 
+  bool ReportJobError(const rpc::ReportJobErrorRequest &request) {
+    std::promise<bool> promise;
+    client_->ReportJobError(
+        request, [&promise](const Status &status, const rpc::ReportJobErrorReply &reply) {
+          RAY_CHECK_OK(status);
+          promise.set_value(true);
+        });
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
   bool WaitReady(const std::future<bool> &future, uint64_t timeout_ms) {
     auto status = future.wait_for(std::chrono::milliseconds(timeout_ms));
     return status == std::future_status::ready;
@@ -613,6 +623,16 @@ TEST_F(GcsServerTest, TestStats) {
   ASSERT_TRUE(AddProfileData(add_profile_data_request));
 }
 
+TEST_F(GcsServerTest, TestErrorInfo) {
+  // Report error
+  rpc::ReportJobErrorRequest report_error_request;
+  rpc::ErrorTableData error_table_data;
+  JobID job_id = JobID::FromInt(1);
+  error_table_data.set_job_id(job_id.Binary());
+  report_error_request.mutable_error_data()->CopyFrom(error_table_data);
+  ASSERT_TRUE(ReportJobError(report_error_request));
+}
+
 }  // namespace ray
 
 int main(int argc, char **argv) {

src/ray/protobuf/gcs_service.proto

@@ -282,3 +282,16 @@ service StatsGcsService {
   // Add profile data to GCS Service.
   rpc AddProfileData(AddProfileDataRequest) returns (AddProfileDataReply);
 }
+
+message ReportJobErrorRequest {
+  ErrorTableData error_data = 1;
+}
+
+message ReportJobErrorReply {
+}
+
+// Service for error info access.
+service ErrorInfoGcsService {
+  // Report a job error to GCS Service.
+  rpc ReportJobError(ReportJobErrorRequest) returns (ReportJobErrorReply);
+}

src/ray/rpc/gcs_server/gcs_rpc_client.h

@@ -28,6 +28,8 @@ class GcsRpcClient {
         new GrpcClient<TaskInfoGcsService>(address, port, client_call_manager));
     stats_grpc_client_ = std::unique_ptr<GrpcClient<StatsGcsService>>(
         new GrpcClient<StatsGcsService>(address, port, client_call_manager));
+    error_info_grpc_client_ = std::unique_ptr<GrpcClient<ErrorInfoGcsService>>(
+        new GrpcClient<ErrorInfoGcsService>(address, port, client_call_manager));
   };
 
   /// Add job info to gcs server.
@@ -114,6 +116,9 @@ class GcsRpcClient {
   /// Add profile data to GCS Service.
   VOID_RPC_CLIENT_METHOD(StatsGcsService, AddProfileData, stats_grpc_client_, )
 
+  /// Report a job error to GCS Service.
+  VOID_RPC_CLIENT_METHOD(ErrorInfoGcsService, ReportJobError, error_info_grpc_client_, )
+
  private:
   /// The gRPC-generated stub.
   std::unique_ptr<GrpcClient<JobInfoGcsService>> job_info_grpc_client_;
@@ -122,6 +127,7 @@ class GcsRpcClient {
   std::unique_ptr<GrpcClient<ObjectInfoGcsService>> object_info_grpc_client_;
   std::unique_ptr<GrpcClient<TaskInfoGcsService>> task_info_grpc_client_;
   std::unique_ptr<GrpcClient<StatsGcsService>> stats_grpc_client_;
+  std::unique_ptr<GrpcClient<ErrorInfoGcsService>> error_info_grpc_client_;
 };
 
 }  // namespace rpc

src/ray/rpc/gcs_server/gcs_rpc_server.h

@@ -27,6 +27,9 @@ namespace rpc {
 #define STATS_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \
   RPC_SERVICE_HANDLER(StatsGcsService, HANDLER, CONCURRENCY)
 
+#define ERROR_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \
+  RPC_SERVICE_HANDLER(ErrorInfoGcsService, HANDLER, CONCURRENCY)
+
 class JobInfoGcsServiceHandler {
  public:
   virtual ~JobInfoGcsServiceHandler() = default;
@@ -335,12 +338,49 @@ class StatsGrpcService : public GrpcService {
   StatsGcsServiceHandler &service_handler_;
 };
 
+class ErrorInfoGcsServiceHandler {
+ public:
+  virtual ~ErrorInfoGcsServiceHandler() = default;
+
+  virtual void HandleReportJobError(const ReportJobErrorRequest &request,
+                                    ReportJobErrorReply *reply,
+                                    SendReplyCallback send_reply_callback) = 0;
+};
+
+/// The `GrpcService` for `ErrorInfoGcsService`.
+class ErrorInfoGrpcService : public GrpcService {
+ public:
+  /// Constructor.
+  ///
+  /// \param[in] handler The service handler that actually handle the requests.
+  explicit ErrorInfoGrpcService(boost::asio::io_service &io_service,
+                                ErrorInfoGcsServiceHandler &handler)
+      : GrpcService(io_service), service_handler_(handler){};
+
+ protected:
+  grpc::Service &GetGrpcService() override { return service_; }
+
+  void InitServerCallFactories(
+      const std::unique_ptr<grpc::ServerCompletionQueue> &cq,
+      std::vector<std::pair<std::unique_ptr<ServerCallFactory>, int>>
+          *server_call_factories_and_concurrencies) override {
+    ERROR_INFO_SERVICE_RPC_HANDLER(ReportJobError, 1);
+  }
+
+ private:
+  /// The grpc async service object.
+  ErrorInfoGcsService::AsyncService service_;
+  /// The service handler that actually handle the requests.
+  ErrorInfoGcsServiceHandler &service_handler_;
+};
+
 using JobInfoHandler = JobInfoGcsServiceHandler;
 using ActorInfoHandler = ActorInfoGcsServiceHandler;
 using NodeInfoHandler = NodeInfoGcsServiceHandler;
 using ObjectInfoHandler = ObjectInfoGcsServiceHandler;
 using TaskInfoHandler = TaskInfoGcsServiceHandler;
 using StatsHandler = StatsGcsServiceHandler;
+using ErrorInfoHandler = ErrorInfoGcsServiceHandler;
 
 }  // namespace rpc
 }  // namespace ray