Revert the gRPC-based resource broadcast due to check failures during cluster autoscaling (#16910)

Author: ericl

src/ray/common/ray_config_def.h

@@ -252,6 +252,9 @@ RAY_CONFIG(uint64_t, gcs_max_concurrent_resource_pulls, 100)
 // Feature flag to turn on resource report polling. Polling and raylet pushing are
 // mutually exlusive.
 RAY_CONFIG(bool, pull_based_resource_reporting, true)
+// Feature flag to use grpc instead of redis for resource broadcast.
+// TODO(ekl) broken as of https://github.com/ray-project/ray/issues/16858
+RAY_CONFIG(bool, grpc_based_resource_broadcast, false)
 // Feature flag to enable grpc based pubsub in GCS.
 RAY_CONFIG(bool, gcs_grpc_based_pubsub, false)
 
@@ -428,4 +431,4 @@ RAY_CONFIG(bool, gcs_task_scheduling_enabled,
 RAY_CONFIG(uint32_t, max_error_msg_size_bytes, 512 * 1024)
 
 /// If enabled, raylet will report resources only when resources are changed.
-RAY_CONFIG(bool, enable_light_weight_resource_report, true)
+RAY_CONFIG(bool, enable_light_weight_resource_report, true)

src/ray/gcs/accessor.h

@@ -87,6 +87,17 @@ class ActorInfoAccessor {
   virtual Status AsyncCreateActor(const TaskSpecification &task_spec,
                                   const StatusCallback &callback) = 0;
 
+  /// Subscribe to any register or update operations of actors.
+  ///
+  /// \param subscribe Callback that will be called each time when an actor is registered
+  /// or updated.
+  /// \param done Callback that will be called when subscription is complete and we
+  /// are ready to receive notification.
+  /// \return Status
+  virtual Status AsyncSubscribeAll(
+      const SubscribeCallback<ActorID, rpc::ActorTableData> &subscribe,
+      const StatusCallback &done) = 0;
+
   /// Subscribe to any update operations of an actor.
   ///
   /// \param actor_id The ID of actor to be subscribed to.
@@ -182,6 +193,96 @@ class JobInfoAccessor {
   JobInfoAccessor() = default;
 };
 
+/// \class TaskInfoAccessor
+/// `TaskInfoAccessor` is a sub-interface of `GcsClient`.
+/// This class includes all the methods that are related to accessing
+/// task information in the GCS.
+class TaskInfoAccessor {
+ public:
+  virtual ~TaskInfoAccessor() {}
+
+  /// Add a task to GCS asynchronously.
+  ///
+  /// \param data_ptr The task that will be added to GCS.
+  /// \param callback Callback that will be called after task has been added
+  /// to GCS.
+  /// \return Status
+  virtual Status AsyncAdd(const std::shared_ptr<rpc::TaskTableData> &data_ptr,
+                          const StatusCallback &callback) = 0;
+
+  /// Get task information from GCS asynchronously.
+  ///
+  /// \param task_id The ID of the task to look up in GCS.
+  /// \param callback Callback that is called after lookup finished.
+  /// \return Status
+  virtual Status AsyncGet(const TaskID &task_id,
+                          const OptionalItemCallback<rpc::TaskTableData> &callback) = 0;
+
+  /// Add a task lease to GCS asynchronously.
+  ///
+  /// \param data_ptr The task lease that will be added to GCS.
+  /// \param callback Callback that will be called after task lease has been added
+  /// to GCS.
+  /// \return Status
+  virtual Status AsyncAddTaskLease(const std::shared_ptr<rpc::TaskLeaseData> &data_ptr,
+                                   const StatusCallback &callback) = 0;
+
+  /// Get task lease information from GCS asynchronously.
+  ///
+  /// \param task_id The ID of the task to look up in GCS.
+  /// \param callback Callback that is called after lookup finished.
+  /// \return Status
+  virtual Status AsyncGetTaskLease(
+      const TaskID &task_id,
+      const OptionalItemCallback<rpc::TaskLeaseData> &callback) = 0;
+
+  /// Subscribe asynchronously to the event that the given task lease is added in GCS.
+  ///
+  /// \param task_id The ID of the task to be subscribed to.
+  /// \param subscribe Callback that will be called each time when the task lease is
+  /// updated or the task lease is empty currently.
+  /// \param done Callback that will be called when subscription is complete.
+  /// \return Status
+  virtual Status AsyncSubscribeTaskLease(
+      const TaskID &task_id,
+      const SubscribeCallback<TaskID, boost::optional<rpc::TaskLeaseData>> &subscribe,
+      const StatusCallback &done) = 0;
+
+  /// Cancel subscription to a task lease asynchronously.
+  ///
+  /// \param task_id The ID of the task to be unsubscribed to.
+  /// \return Status
+  virtual Status AsyncUnsubscribeTaskLease(const TaskID &task_id) = 0;
+
+  /// Attempt task reconstruction to GCS asynchronously.
+  ///
+  /// \param data_ptr The task reconstruction that will be added to GCS.
+  /// \param callback Callback that will be called after task reconstruction
+  /// has been added to GCS.
+  /// \return Status
+  virtual Status AttemptTaskReconstruction(
+      const std::shared_ptr<rpc::TaskReconstructionData> &data_ptr,
+      const StatusCallback &callback) = 0;
+
+  /// Reestablish subscription.
+  /// This should be called when GCS server restarts from a failure.
+  /// PubSub server restart will cause GCS server restart. In this case, we need to
+  /// resubscribe from PubSub server, otherwise we only need to fetch data from GCS
+  /// server.
+  ///
+  /// \param is_pubsub_server_restarted Whether pubsub server is restarted.
+  virtual void AsyncResubscribe(bool is_pubsub_server_restarted) = 0;
+
+  /// Check if the specified task lease is unsubscribed.
+  ///
+  /// \param task_id The ID of the task.
+  /// \return Whether the specified task lease is unsubscribed.
+  virtual bool IsTaskLeaseUnsubscribed(const TaskID &task_id) = 0;
+
+ protected:
+  TaskInfoAccessor() = default;
+};
+
 /// `ObjectInfoAccessor` is a sub-interface of `GcsClient`.
 /// This class includes all the methods that are related to accessing
 /// object information in the GCS.

src/ray/gcs/gcs_client.h

@@ -123,6 +123,13 @@ class GcsClient : public std::enable_shared_from_this<GcsClient> {
     return *node_resource_accessor_;
   }
 
+  /// Get the sub-interface for accessing task information in GCS.
+  /// This function is thread safe.
+  TaskInfoAccessor &Tasks() {
+    RAY_CHECK(task_accessor_ != nullptr);
+    return *task_accessor_;
+  }
+
   /// Get the sub-interface for accessing error information in GCS.
   /// This function is thread safe.
   ErrorInfoAccessor &Errors() {
@@ -171,6 +178,7 @@ class GcsClient : public std::enable_shared_from_this<GcsClient> {
   std::unique_ptr<ObjectInfoAccessor> object_accessor_;
   std::unique_ptr<NodeInfoAccessor> node_accessor_;
   std::unique_ptr<NodeResourceInfoAccessor> node_resource_accessor_;
+  std::unique_ptr<TaskInfoAccessor> task_accessor_;
   std::unique_ptr<ErrorInfoAccessor> error_accessor_;
   std::unique_ptr<StatsInfoAccessor> stats_accessor_;
   std::unique_ptr<WorkerInfoAccessor> worker_accessor_;

src/ray/gcs/gcs_client/service_based_accessor.cc

@@ -252,6 +252,37 @@ Status ServiceBasedActorInfoAccessor::AsyncCreateActor(
   return Status::OK();
 }
 
+Status ServiceBasedActorInfoAccessor::AsyncSubscribeAll(
+    const SubscribeCallback<ActorID, rpc::ActorTableData> &subscribe,
+    const StatusCallback &done) {
+  RAY_CHECK(subscribe != nullptr);
+  fetch_all_data_operation_ = [this, subscribe](const StatusCallback &done) {
+    auto callback = [subscribe, done](
+                        const Status &status,
+                        const std::vector<rpc::ActorTableData> &actor_info_list) {
+      for (auto &actor_info : actor_info_list) {
+        subscribe(ActorID::FromBinary(actor_info.actor_id()), actor_info);
+      }
+      if (done) {
+        done(status);
+      }
+    };
+    RAY_CHECK_OK(AsyncGetAll(callback));
+  };
+
+  subscribe_all_operation_ = [this, subscribe](const StatusCallback &done) {
+    auto on_subscribe = [subscribe](const std::string &id, const std::string &data) {
+      ActorTableData actor_data;
+      actor_data.ParseFromString(data);
+      subscribe(ActorID::FromBinary(actor_data.actor_id()), actor_data);
+    };
+    return client_impl_->GetGcsPubSub().SubscribeAll(ACTOR_CHANNEL, on_subscribe, done);
+  };
+
+  return subscribe_all_operation_(
+      [this, done](const Status &status) { fetch_all_data_operation_(done); });
+}
+
 Status ServiceBasedActorInfoAccessor::AsyncSubscribe(
     const ActorID &actor_id,
     const SubscribeCallback<ActorID, rpc::ActorTableData> &subscribe,
@@ -835,6 +866,189 @@ Status ServiceBasedNodeResourceInfoAccessor::AsyncGetAllResourceUsage(
   return Status::OK();
 }
 
+ServiceBasedTaskInfoAccessor::ServiceBasedTaskInfoAccessor(
+    ServiceBasedGcsClient *client_impl)
+    : client_impl_(client_impl) {}
+
+Status ServiceBasedTaskInfoAccessor::AsyncAdd(
+    const std::shared_ptr<rpc::TaskTableData> &data_ptr, const StatusCallback &callback) {
+  TaskID task_id = TaskID::FromBinary(data_ptr->task().task_spec().task_id());
+  JobID job_id = JobID::FromBinary(data_ptr->task().task_spec().job_id());
+  RAY_LOG(DEBUG) << "Adding task, task id = " << task_id << ", job id = " << job_id;
+  rpc::AddTaskRequest request;
+  request.mutable_task_data()->CopyFrom(*data_ptr);
+  client_impl_->GetGcsRpcClient().AddTask(
+      request,
+      [task_id, job_id, callback](const Status &status, const rpc::AddTaskReply &reply) {
+        if (callback) {
+          callback(status);
+        }
+        RAY_LOG(DEBUG) << "Finished adding task, status = " << status
+                       << ", task id = " << task_id << ", job id = " << job_id;
+      });
+  return Status::OK();
+}
+
+Status ServiceBasedTaskInfoAccessor::AsyncGet(
+    const TaskID &task_id, const OptionalItemCallback<rpc::TaskTableData> &callback) {
+  RAY_LOG(DEBUG) << "Getting task, task id = " << task_id
+                 << ", job id = " << task_id.JobId();
+  rpc::GetTaskRequest request;
+  request.set_task_id(task_id.Binary());
+  client_impl_->GetGcsRpcClient().GetTask(
+      request, [task_id, callback](const Status &status, const rpc::GetTaskReply &reply) {
+        if (reply.has_task_data()) {
+          callback(status, reply.task_data());
+        } else {
+          callback(status, boost::none);
+        }
+        RAY_LOG(DEBUG) << "Finished getting task, status = " << status
+                       << ", task id = " << task_id << ", job id = " << task_id.JobId();
+      });
+  return Status::OK();
+}
+
+Status ServiceBasedTaskInfoAccessor::AsyncAddTaskLease(
+    const std::shared_ptr<rpc::TaskLeaseData> &data_ptr, const StatusCallback &callback) {
+  TaskID task_id = TaskID::FromBinary(data_ptr->task_id());
+  NodeID node_id = NodeID::FromBinary(data_ptr->node_manager_id());
+  RAY_LOG(DEBUG) << "Adding task lease, task id = " << task_id
+                 << ", node id = " << node_id << ", job id = " << task_id.JobId();
+  rpc::AddTaskLeaseRequest request;
+  request.mutable_task_lease_data()->CopyFrom(*data_ptr);
+  client_impl_->GetGcsRpcClient().AddTaskLease(
+      request, [task_id, node_id, callback](const Status &status,
+                                            const rpc::AddTaskLeaseReply &reply) {
+        if (callback) {
+          callback(status);
+        }
+        RAY_LOG(DEBUG) << "Finished adding task lease, status = " << status
+                       << ", task id = " << task_id << ", node id = " << node_id
+                       << ", job id = " << task_id.JobId();
+      });
+  return Status::OK();
+}
+
+Status ServiceBasedTaskInfoAccessor::AsyncGetTaskLease(
+    const TaskID &task_id, const OptionalItemCallback<rpc::TaskLeaseData> &callback) {
+  RAY_LOG(DEBUG) << "Getting task lease, task id = " << task_id
+                 << ", job id = " << task_id.JobId();
+  rpc::GetTaskLeaseRequest request;
+  request.set_task_id(task_id.Binary());
+  client_impl_->GetGcsRpcClient().GetTaskLease(
+      request,
+      [task_id, callback](const Status &status, const rpc::GetTaskLeaseReply &reply) {
+        if (reply.has_task_lease_data()) {
+          callback(status, reply.task_lease_data());
+        } else {
+          callback(status, boost::none);
+        }
+        RAY_LOG(DEBUG) << "Finished getting task lease, status = " << status
+                       << ", task id = " << task_id << ", job id = " << task_id.JobId();
+      });
+  return Status::OK();
+}
+
+Status ServiceBasedTaskInfoAccessor::AsyncSubscribeTaskLease(
+    const TaskID &task_id,
+    const SubscribeCallback<TaskID, boost::optional<rpc::TaskLeaseData>> &subscribe,
+    const StatusCallback &done) {
+  RAY_CHECK(subscribe != nullptr)
+      << "Failed to subscribe task lease, task id = " << task_id
+      << ", job id = " << task_id.JobId();
+
+  auto fetch_data_operation = [this, task_id,
+                               subscribe](const StatusCallback &fetch_done) {
+    auto callback = [task_id, subscribe, fetch_done](
+                        const Status &status,
+                        const boost::optional<rpc::TaskLeaseData> &result) {
+      subscribe(task_id, result);
+      if (fetch_done) {
+        fetch_done(status);
+      }
+    };
+    RAY_CHECK_OK(AsyncGetTaskLease(task_id, callback));
+  };
+
+  auto subscribe_operation = [this, task_id,
+                              subscribe](const StatusCallback &subscribe_done) {
+    auto on_subscribe = [task_id, subscribe](const std::string &id,
+                                             const std::string &data) {
+      TaskLeaseData task_lease_data;
+      task_lease_data.ParseFromString(data);
+      subscribe(task_id, task_lease_data);
+    };
+    return client_impl_->GetGcsPubSub().Subscribe(TASK_LEASE_CHANNEL, task_id.Hex(),
+                                                  on_subscribe, subscribe_done);
+  };
+
+  subscribe_task_lease_operations_[task_id] = subscribe_operation;
+  fetch_task_lease_data_operations_[task_id] = fetch_data_operation;
+  return subscribe_operation(
+      [fetch_data_operation, done](const Status &status) { fetch_data_operation(done); });
+}
+
+Status ServiceBasedTaskInfoAccessor::AsyncUnsubscribeTaskLease(const TaskID &task_id) {
+  RAY_LOG(DEBUG) << "Unsubscribing task lease, task id = " << task_id
+                 << ", job id = " << task_id.JobId();
+  auto status =
+      client_impl_->GetGcsPubSub().Unsubscribe(TASK_LEASE_CHANNEL, task_id.Hex());
+  subscribe_task_lease_operations_.erase(task_id);
+  fetch_task_lease_data_operations_.erase(task_id);
+  RAY_LOG(DEBUG) << "Finished unsubscribing task lease, task id = " << task_id
+                 << ", job id = " << task_id.JobId();
+  return status;
+}
+
+Status ServiceBasedTaskInfoAccessor::AttemptTaskReconstruction(
+    const std::shared_ptr<rpc::TaskReconstructionData> &data_ptr,
+    const StatusCallback &callback) {
+  auto num_reconstructions = data_ptr->num_reconstructions();
+  NodeID node_id = NodeID::FromBinary(data_ptr->node_manager_id());
+  TaskID task_id = TaskID::FromBinary(data_ptr->task_id());
+  RAY_LOG(DEBUG) << "Reconstructing task, reconstructions num = " << num_reconstructions
+                 << ", node id = " << node_id << ", task id = " << task_id
+                 << ", job id = " << task_id.JobId();
+  rpc::AttemptTaskReconstructionRequest request;
+  request.mutable_task_reconstruction()->CopyFrom(*data_ptr);
+  client_impl_->GetGcsRpcClient().AttemptTaskReconstruction(
+      request,
+      [num_reconstructions, node_id, task_id, callback](
+          const Status &status, const rpc::AttemptTaskReconstructionReply &reply) {
+        if (callback) {
+          callback(status);
+        }
+        RAY_LOG(DEBUG) << "Finished reconstructing task, status = " << status
+                       << ", reconstructions num = " << num_reconstructions
+                       << ", node id = " << node_id << ", task id = " << task_id
+                       << ", job id = " << task_id.JobId();
+      });
+  return Status::OK();
+}
+
+void ServiceBasedTaskInfoAccessor::AsyncResubscribe(bool is_pubsub_server_restarted) {
+  RAY_LOG(DEBUG) << "Reestablishing subscription for task info.";
+  // If only the GCS sever has restarted, we only need to fetch data from the GCS server.
+  // If the pub-sub server has also restarted, we need to resubscribe to the pub-sub
+  // server first, then fetch data from the GCS server.
+  if (is_pubsub_server_restarted) {
+    for (auto &item : subscribe_task_lease_operations_) {
+      auto &task_id = item.first;
+      RAY_CHECK_OK(item.second([this, task_id](const Status &status) {
+        fetch_task_lease_data_operations_[task_id](nullptr);
+      }));
+    }
+  } else {
+    for (auto &item : fetch_task_lease_data_operations_) {
+      item.second(nullptr);
+    }
+  }
+}
+
+bool ServiceBasedTaskInfoAccessor::IsTaskLeaseUnsubscribed(const TaskID &task_id) {
+  return client_impl_->GetGcsPubSub().IsUnsubscribed(TASK_LEASE_CHANNEL, task_id.Hex());
+}
+
 ServiceBasedObjectInfoAccessor::ServiceBasedObjectInfoAccessor(
     ServiceBasedGcsClient *client_impl)
     : client_impl_(client_impl) {}