Publish unexpected worker failures (#6734)

* Publish unexpected worker failures

* comment

* Move failure handler to NodeManager, update comments

src/ray/gcs/redis_gcs_client.cc

@@ -140,6 +140,7 @@ Status RedisGcsClient::Connect(boost::asio::io_service &io_service) {
   actor_checkpoint_table_.reset(new ActorCheckpointTable(shard_contexts_, this));
   actor_checkpoint_id_table_.reset(new ActorCheckpointIdTable(shard_contexts_, this));
   resource_table_.reset(new DynamicResourceTable({primary_context_}, this));
+  worker_failure_table_.reset(new WorkerFailureTable(shard_contexts_, this));
 
   actor_accessor_.reset(new RedisActorInfoAccessor(this));
   job_accessor_.reset(new RedisJobInfoAccessor(this));
@@ -198,6 +199,10 @@ raylet::TaskTable &RedisGcsClient::raylet_task_table() { return *raylet_task_tab
 
 ActorTable &RedisGcsClient::actor_table() { return *actor_table_; }
 
+WorkerFailureTable &RedisGcsClient::worker_failure_table() {
+  return *worker_failure_table_;
+}
+
 TaskReconstructionLog &RedisGcsClient::task_reconstruction_log() {
   return *task_reconstruction_log_;
 }

src/ray/gcs/redis_gcs_client.h

@@ -66,6 +66,7 @@ class RAY_EXPORT RedisGcsClient : public GcsClient {
 
   // TODO: Some API for getting the error on the driver
   ErrorTable &error_table();
+  WorkerFailureTable &worker_failure_table();
 
   // We also need something to export generic code to run on workers from the
   // driver (to set the PYTHONPATH)
@@ -125,6 +126,7 @@ class RAY_EXPORT RedisGcsClient : public GcsClient {
   std::unique_ptr<ActorCheckpointTable> actor_checkpoint_table_;
   std::unique_ptr<ActorCheckpointIdTable> actor_checkpoint_id_table_;
   std::unique_ptr<DynamicResourceTable> resource_table_;
+  std::unique_ptr<WorkerFailureTable> worker_failure_table_;
   // The following contexts write to the data shard
   std::vector<std::shared_ptr<RedisContext>> shard_contexts_;
   std::vector<std::unique_ptr<RedisAsioClient>> shard_asio_async_clients_;

src/ray/gcs/tables.cc

@@ -790,6 +790,7 @@ template class Log<JobID, JobTableData>;
 template class Log<UniqueID, ProfileTableData>;
 template class Table<ActorCheckpointID, ActorCheckpointData>;
 template class Table<ActorID, ActorCheckpointIdData>;
+template class Table<WorkerID, WorkerFailureData>;
 
 template class Log<ClientID, ResourceTableData>;
 template class Hash<ClientID, ResourceTableData>;

src/ray/gcs/tables.h

@@ -40,6 +40,7 @@ using rpc::TablePubsub;
 using rpc::TaskLeaseData;
 using rpc::TaskReconstructionData;
 using rpc::TaskTableData;
+using rpc::WorkerFailureData;
 
 class RedisContext;
 
@@ -716,6 +717,17 @@ class ActorTable : public Log<ActorID, ActorTableData> {
   }
 };
 
+class WorkerFailureTable : public Table<WorkerID, WorkerFailureData> {
+ public:
+  WorkerFailureTable(const std::vector<std::shared_ptr<RedisContext>> &contexts,
+                     RedisGcsClient *client)
+      : Table(contexts, client) {
+    pubsub_channel_ = TablePubsub::WORKER_FAILURE_PUBSUB;
+    prefix_ = TablePrefix::WORKER_FAILURE;
+  }
+  virtual ~WorkerFailureTable() {}
+};
+
 class TaskReconstructionLog : public Log<TaskID, TaskReconstructionData> {
  public:
   TaskReconstructionLog(const std::vector<std::shared_ptr<RedisContext>> &contexts,

src/ray/protobuf/gcs.proto

@@ -27,7 +27,8 @@ enum TablePrefix {
   ACTOR_CHECKPOINT_ID = 16;
   NODE_RESOURCE = 17;
   DIRECT_ACTOR = 18;
-  TABLE_PREFIX_MAX = 19;
+  WORKER_FAILURE = 19;
+  TABLE_PREFIX_MAX = 20;
 }
 
 // The channel that Add operations to the Table should be published on, if any.
@@ -46,7 +47,8 @@ enum TablePubsub {
   JOB_PUBSUB = 11;
   NODE_RESOURCE_PUBSUB = 12;
   DIRECT_ACTOR_PUBSUB = 13;
-  TABLE_PUBSUB_MAX = 14;
+  WORKER_FAILURE_PUBSUB = 14;
+  TABLE_PUBSUB_MAX = 15;
 }
 
 enum GcsChangeMode {
@@ -272,6 +274,13 @@ message ActorCheckpointIdData {
   repeated uint64 timestamps = 3;
 }
 
+message WorkerFailureData {
+  // Address of the worker that failed.
+  Address worker_address = 1;
+  // The UNIX timestamp at which the worker failed.
+  int64 timestamp = 3;
+}
+
 // This enum type is used as object's metadata to indicate the object's creating
 // task has failed because of a certain error.
 // TODO(hchen): We may want to make these errors more specific. E.g., we may want

src/ray/raylet/node_manager.cc

@@ -191,6 +191,18 @@ ray::Status NodeManager::RegisterGcs() {
   RAY_RETURN_NOT_OK(gcs_client_->Nodes().AsyncSubscribeBatchHeartbeat(
       heartbeat_batch_added, /*done*/ nullptr));
 
+  // Subscribe to all unexpected failure notifications from the local and
+  // remote raylets. Note that this does not include workers that failed due to
+  // node failure. These workers can be identified by comparing the raylet_id
+  // in their rpc::Address to the ID of a failed raylet.
+  const auto &failure_handler = [this](gcs::RedisGcsClient *client, const WorkerID &id,
+                                       const gcs::WorkerFailureData &worker_failure) {
+    HandleUnexpectedWorkerFailure(id, worker_failure);
+  };
+  RAY_CHECK_OK(gcs_client_->worker_failure_table().Subscribe(
+      JobID::Nil(), ClientID::Nil(), failure_handler,
+      /*done_callback=*/nullptr));
+
   // Subscribe to job updates.
   const auto job_subscribe_handler = [this](const JobID &job_id,
                                             const JobTableData &job_data) {
@@ -210,6 +222,12 @@ ray::Status NodeManager::RegisterGcs() {
   return ray::Status::OK();
 }
 
+void NodeManager::HandleUnexpectedWorkerFailure(
+    const WorkerID &worker_id, const gcs::WorkerFailureData &worker_failed_data) {
+  RAY_LOG(DEBUG) << "Worker " << worker_id << " failed";
+  // TODO: Clean up after the failure: If the failed worker is our owner, then exit.
+}
+
 void NodeManager::KillWorker(std::shared_ptr<Worker> worker) {
   // If we're just cleaning up a single worker, allow it some time to clean
   // up its state before force killing. The client socket will be closed
@@ -1111,6 +1129,16 @@ void NodeManager::ProcessDisconnectClientMessage(
     }
     // Erase any lease metadata.
     leased_workers_.erase(worker->WorkerId());
+
+    // Publish the worker failure.
+    auto data = std::make_shared<rpc::WorkerFailureData>();
+    data->mutable_worker_address()->set_ip_address(initial_config_.node_manager_address);
+    data->mutable_worker_address()->set_port(worker->Port());
+    data->mutable_worker_address()->set_worker_id(worker->WorkerId().Binary());
+    data->mutable_worker_address()->set_raylet_id(self_node_id_.Binary());
+    data->set_timestamp(std::time(nullptr));
+    RAY_CHECK_OK(gcs_client_->worker_failure_table().Add(JobID::Nil(), worker->WorkerId(),
+                                                         data, nullptr));
   }
 
   if (is_worker) {

src/ray/raylet/node_manager.h

@@ -119,6 +119,13 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
  private:
   /// Methods for handling clients.
 
+  /// Handle an unexpected failure notification from GCS pubsub.
+  ///
+  /// \param worker_id The ID of the failed worker.
+  /// \param worker_data Data associated with the worker failure.
+  void HandleUnexpectedWorkerFailure(const WorkerID &worker_id,
+                                     const gcs::WorkerFailureData &worker_failed_data);
+
   /// Handler for the addition of a new node.
   ///
   /// \param data Data associated with the new node.