gcs adapts to worker table pub sub (#8182)

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

src/ray/gcs/gcs_client/service_based_accessor.cc

@@ -858,17 +858,20 @@ Status ServiceBasedErrorInfoAccessor::AsyncReportJobError(
 
 ServiceBasedWorkerInfoAccessor::ServiceBasedWorkerInfoAccessor(
     ServiceBasedGcsClient *client_impl)
-    : client_impl_(client_impl),
-      worker_failure_sub_executor_(
-          client_impl->GetRedisGcsClient().worker_failure_table()) {}
+    : client_impl_(client_impl) {}
 
 Status ServiceBasedWorkerInfoAccessor::AsyncSubscribeToWorkerFailures(
     const SubscribeCallback<WorkerID, rpc::WorkerFailureData> &subscribe,
     const StatusCallback &done) {
   RAY_LOG(DEBUG) << "Subscribing worker failures.";
   RAY_CHECK(subscribe != nullptr);
-  auto status =
-      worker_failure_sub_executor_.AsyncSubscribeAll(ClientID::Nil(), subscribe, done);
+  auto on_subscribe = [subscribe](const std::string &id, const std::string &data) {
+    rpc::WorkerFailureData worker_failure_data;
+    worker_failure_data.ParseFromString(data);
+    subscribe(WorkerID::FromBinary(id), worker_failure_data);
+  };
+  auto status = client_impl_->GetGcsPubSub().SubscribeAll(WORKER_FAILURE_CHANNEL,
+                                                          on_subscribe, done);
   RAY_LOG(DEBUG) << "Finished subscribing worker failures.";
   return status;
 }

src/ray/gcs/gcs_client/service_based_accessor.h

@@ -337,10 +337,6 @@ class ServiceBasedWorkerInfoAccessor : public WorkerInfoAccessor {
 
  private:
   ServiceBasedGcsClient *client_impl_;
-
-  typedef SubscriptionExecutor<WorkerID, WorkerFailureData, WorkerFailureTable>
-      WorkerFailureSubscriptionExecutor;
-  WorkerFailureSubscriptionExecutor worker_failure_sub_executor_;
 };
 
 }  // namespace gcs

src/ray/gcs/gcs_server/gcs_server.cc

@@ -188,8 +188,8 @@ std::unique_ptr<rpc::ErrorInfoHandler> GcsServer::InitErrorInfoHandler() {
 }
 
 std::unique_ptr<rpc::WorkerInfoHandler> GcsServer::InitWorkerInfoHandler() {
-  return std::unique_ptr<rpc::DefaultWorkerInfoHandler>(
-      new rpc::DefaultWorkerInfoHandler(*redis_gcs_client_, *gcs_actor_manager_));
+  return std::unique_ptr<rpc::DefaultWorkerInfoHandler>(new rpc::DefaultWorkerInfoHandler(
+      *redis_gcs_client_, *gcs_actor_manager_, gcs_pub_sub_));
 }
 
 }  // namespace gcs

src/ray/gcs/gcs_server/worker_info_handler_impl.cc

@@ -29,10 +29,17 @@ void DefaultWorkerInfoHandler::HandleReportWorkerFailure(
   auto worker_id = WorkerID::FromBinary(worker_address.worker_id());
   gcs_actor_manager_.ReconstructActorOnWorker(node_id, worker_id, need_reschedule);
 
-  auto on_done = [worker_address, reply, send_reply_callback](Status status) {
+  auto on_done = [this, worker_address, worker_id, worker_failure_data, reply,
+                  send_reply_callback](const Status &status) {
     if (!status.ok()) {
       RAY_LOG(ERROR) << "Failed to report worker failure, "
                      << worker_address.DebugString();
+    } else {
+      RAY_CHECK_OK(gcs_pub_sub_->Publish(WORKER_FAILURE_CHANNEL, worker_id.Binary(),
+                                         worker_failure_data->SerializeAsString(),
+                                         nullptr));
+      RAY_LOG(DEBUG) << "Finished reporting worker failure, "
+                     << worker_address.DebugString();
     }
     GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
   };
@@ -42,7 +49,6 @@ void DefaultWorkerInfoHandler::HandleReportWorkerFailure(
   if (!status.ok()) {
     on_done(status);
   }
-  RAY_LOG(DEBUG) << "Finished reporting worker failure, " << worker_address.DebugString();
 }
 
 void DefaultWorkerInfoHandler::HandleRegisterWorker(
@@ -52,9 +58,11 @@ void DefaultWorkerInfoHandler::HandleRegisterWorker(
   auto worker_id = WorkerID::FromBinary(request.worker_id());
   auto worker_info = MapFromProtobuf(request.worker_info());
 
-  auto on_done = [worker_id, reply, send_reply_callback](Status status) {
+  auto on_done = [worker_id, reply, send_reply_callback](const Status &status) {
     if (!status.ok()) {
       RAY_LOG(ERROR) << "Failed to register worker " << worker_id;
+    } else {
+      RAY_LOG(DEBUG) << "Finished registering worker " << worker_id;
     }
     GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
   };
@@ -64,7 +72,6 @@ void DefaultWorkerInfoHandler::HandleRegisterWorker(
   if (!status.ok()) {
     on_done(status);
   }
-  RAY_LOG(DEBUG) << "Finished registering worker " << worker_id;
 }
 
 }  // namespace rpc

src/ray/gcs/gcs_server/worker_info_handler_impl.h

@@ -16,6 +16,7 @@
 #define RAY_GCS_WORKER_INFO_HANDLER_IMPL_H
 
 #include "gcs_actor_manager.h"
+#include "ray/gcs/pubsub/gcs_pub_sub.h"
 #include "ray/gcs/redis_gcs_client.h"
 #include "ray/rpc/gcs_server/gcs_rpc_server.h"
 
@@ -26,8 +27,11 @@ namespace rpc {
 class DefaultWorkerInfoHandler : public rpc::WorkerInfoHandler {
  public:
   explicit DefaultWorkerInfoHandler(gcs::RedisGcsClient &gcs_client,
-                                    gcs::GcsActorManager &gcs_actor_manager)
-      : gcs_client_(gcs_client), gcs_actor_manager_(gcs_actor_manager) {}
+                                    gcs::GcsActorManager &gcs_actor_manager,
+                                    std::shared_ptr<gcs::GcsPubSub> &gcs_pub_sub)
+      : gcs_client_(gcs_client),
+        gcs_actor_manager_(gcs_actor_manager),
+        gcs_pub_sub_(gcs_pub_sub) {}
 
   void HandleReportWorkerFailure(const ReportWorkerFailureRequest &request,
                                  ReportWorkerFailureReply *reply,
@@ -40,6 +44,7 @@ class DefaultWorkerInfoHandler : public rpc::WorkerInfoHandler {
  private:
   gcs::RedisGcsClient &gcs_client_;
   gcs::GcsActorManager &gcs_actor_manager_;
+  std::shared_ptr<gcs::GcsPubSub> gcs_pub_sub_;
 };
 
 }  // namespace rpc

src/ray/gcs/pubsub/gcs_pub_sub.h

@@ -26,6 +26,7 @@ namespace ray {
 namespace gcs {
 
 #define JOB_CHANNEL "JOB"
+#define WORKER_FAILURE_CHANNEL "WORKER_FAILURE"
 
 /// \class GcsPubSub
 ///