diff --git a/src/ray/gcs/accessor.h b/src/ray/gcs/accessor.h index 11335f4bd..7c730035d 100644 --- a/src/ray/gcs/accessor.h +++ b/src/ray/gcs/accessor.h @@ -309,6 +309,12 @@ class TaskInfoAccessor { const std::shared_ptr &data_ptr, const StatusCallback &callback) = 0; + /// Reestablish subscription. + /// This should be called when GCS server restarts from a failure. + /// + /// \return Status + virtual Status AsyncReSubscribe() = 0; + protected: TaskInfoAccessor() = default; }; diff --git a/src/ray/gcs/gcs_client/service_based_accessor.cc b/src/ray/gcs/gcs_client/service_based_accessor.cc index 9853d0464..f00b10a9e 100644 --- a/src/ray/gcs/gcs_client/service_based_accessor.cc +++ b/src/ray/gcs/gcs_client/service_based_accessor.cc @@ -830,38 +830,44 @@ Status ServiceBasedTaskInfoAccessor::AsyncSubscribe( const StatusCallback &done) { RAY_LOG(DEBUG) << "Subscribing task, task id = " << task_id; RAY_CHECK(subscribe != nullptr) << "Failed to subscribe task, task id = " << task_id; - auto on_subscribe = [task_id, subscribe](const std::string &id, - const std::string &data) { - TaskTableData task_data; - task_data.ParseFromString(data); - subscribe(task_id, task_data); + + auto subscribe_operation = [this, task_id, subscribe](const StatusCallback &done) { + auto on_subscribe = [task_id, subscribe](const std::string &id, + const std::string &data) { + TaskTableData task_data; + task_data.ParseFromString(data); + subscribe(task_id, task_data); + }; + auto on_done = [this, task_id, subscribe, done](const Status &status) { + if (status.ok()) { + auto callback = [task_id, subscribe, done]( + const Status &status, + const boost::optional &result) { + if (result) { + subscribe(task_id, *result); + } + if (done) { + done(status); + } + }; + RAY_CHECK_OK(AsyncGet(task_id, callback)); + } else if (done) { + done(status); + } + }; + auto status = client_impl_->GetGcsPubSub().Subscribe(TASK_CHANNEL, task_id.Hex(), + on_subscribe, on_done); + RAY_LOG(DEBUG) << "Finished subscribing task, task id = " << task_id; + return status; }; - auto on_done = [this, task_id, subscribe, done](const Status &status) { - if (status.ok()) { - auto callback = [task_id, subscribe, done]( - const Status &status, - const boost::optional &result) { - if (result) { - subscribe(task_id, *result); - } - if (done) { - done(status); - } - }; - RAY_CHECK_OK(AsyncGet(task_id, callback)); - } else if (done) { - done(status); - } - }; - auto status = client_impl_->GetGcsPubSub().Subscribe(TASK_CHANNEL, task_id.Hex(), - on_subscribe, on_done); - RAY_LOG(DEBUG) << "Finished subscribing task, task id = " << task_id; - return status; + subscribe_task_operations_[task_id] = subscribe_operation; + return subscribe_operation(done); } Status ServiceBasedTaskInfoAccessor::AsyncUnsubscribe(const TaskID &task_id) { RAY_LOG(DEBUG) << "Unsubscribing task, task id = " << task_id; auto status = client_impl_->GetGcsPubSub().Unsubscribe(TASK_CHANNEL, task_id.Hex()); + subscribe_task_operations_.erase(task_id); RAY_LOG(DEBUG) << "Finished unsubscribing task, task id = " << task_id; return status; } @@ -912,37 +918,43 @@ Status ServiceBasedTaskInfoAccessor::AsyncSubscribeTaskLease( RAY_LOG(DEBUG) << "Subscribing task lease, task id = " << task_id; RAY_CHECK(subscribe != nullptr) << "Failed to subscribe task lease, task id = " << task_id; - 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); + + auto subscribe_operation = [this, task_id, subscribe](const StatusCallback &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); + }; + auto on_done = [this, task_id, subscribe, done](const Status &status) { + if (status.ok()) { + auto callback = [task_id, subscribe, done]( + const Status &status, + const boost::optional &result) { + subscribe(task_id, result); + if (done) { + done(status); + } + }; + RAY_CHECK_OK(AsyncGetTaskLease(task_id, callback)); + } else if (done) { + done(status); + } + }; + auto status = client_impl_->GetGcsPubSub().Subscribe( + TASK_LEASE_CHANNEL, task_id.Hex(), on_subscribe, on_done); + RAY_LOG(DEBUG) << "Finished subscribing task lease, task id = " << task_id; + return status; }; - auto on_done = [this, task_id, subscribe, done](const Status &status) { - if (status.ok()) { - auto callback = [task_id, subscribe, done]( - const Status &status, - const boost::optional &result) { - subscribe(task_id, result); - if (done) { - done(status); - } - }; - RAY_CHECK_OK(AsyncGetTaskLease(task_id, callback)); - } else if (done) { - done(status); - } - }; - auto status = client_impl_->GetGcsPubSub().Subscribe(TASK_LEASE_CHANNEL, task_id.Hex(), - on_subscribe, on_done); - RAY_LOG(DEBUG) << "Finished subscribing task lease, task id = " << task_id; - return status; + subscribe_task_lease_operations_[task_id] = subscribe_operation; + return subscribe_operation(done); } Status ServiceBasedTaskInfoAccessor::AsyncUnsubscribeTaskLease(const TaskID &task_id) { RAY_LOG(DEBUG) << "Unsubscribing task lease, task id = " << task_id; auto status = client_impl_->GetGcsPubSub().Unsubscribe(TASK_LEASE_CHANNEL, task_id.Hex()); + subscribe_task_lease_operations_.erase(task_id); RAY_LOG(DEBUG) << "Finished unsubscribing task lease, task id = " << task_id; return status; } @@ -969,6 +981,17 @@ Status ServiceBasedTaskInfoAccessor::AttemptTaskReconstruction( return Status::OK(); } +Status ServiceBasedTaskInfoAccessor::AsyncReSubscribe() { + RAY_LOG(INFO) << "Reestablishing subscription for task info."; + for (auto &item : subscribe_task_operations_) { + RAY_CHECK_OK(item.second(nullptr)); + } + for (auto &item : subscribe_task_lease_operations_) { + RAY_CHECK_OK(item.second(nullptr)); + } + return Status::OK(); +} + ServiceBasedObjectInfoAccessor::ServiceBasedObjectInfoAccessor( ServiceBasedGcsClient *client_impl) : client_impl_(client_impl) {} diff --git a/src/ray/gcs/gcs_client/service_based_accessor.h b/src/ray/gcs/gcs_client/service_based_accessor.h index 4948ae882..0115b17c5 100644 --- a/src/ray/gcs/gcs_client/service_based_accessor.h +++ b/src/ray/gcs/gcs_client/service_based_accessor.h @@ -257,7 +257,14 @@ class ServiceBasedTaskInfoAccessor : public TaskInfoAccessor { const std::shared_ptr &data_ptr, const StatusCallback &callback) override; + Status AsyncReSubscribe() override; + private: + /// Save the subscribe operation in this function, so we can call it again when GCS + /// restarts from a failure. + std::unordered_map subscribe_task_operations_; + std::unordered_map subscribe_task_lease_operations_; + ServiceBasedGcsClient *client_impl_; }; diff --git a/src/ray/gcs/gcs_client/service_based_gcs_client.cc b/src/ray/gcs/gcs_client/service_based_gcs_client.cc index b5d4c82ca..c39be0380 100644 --- a/src/ray/gcs/gcs_client/service_based_gcs_client.cc +++ b/src/ray/gcs/gcs_client/service_based_gcs_client.cc @@ -51,6 +51,7 @@ Status ServiceBasedGcsClient::Connect(boost::asio::io_service &io_service) { RAY_CHECK_OK(job_accessor_->AsyncReSubscribe()); RAY_CHECK_OK(actor_accessor_->AsyncReSubscribe()); RAY_CHECK_OK(node_accessor_->AsyncReSubscribe()); + RAY_CHECK_OK(task_accessor_->AsyncReSubscribe()); }; // Connect to gcs service. diff --git a/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc b/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc index 5782f3a9f..7ceb90eeb 100644 --- a/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc +++ b/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc @@ -917,6 +917,45 @@ TEST_F(ServiceBasedGcsClientTest, TestNodeTableReSubscribe) { WaitPendingDone(batch_heartbeat_count, 1); } +TEST_F(ServiceBasedGcsClientTest, TestTaskTableReSubscribe) { + JobID job_id = JobID::FromInt(6); + TaskID task_id = TaskID::ForDriverTask(job_id); + auto task_table_data = Mocker::GenTaskTableData(job_id.Binary(), task_id.Binary()); + + // Subscribe to the event that the given task is added in GCS. + std::atomic task_count(0); + auto task_subscribe = [&task_count](const TaskID &task_id, + const gcs::TaskTableData &data) { ++task_count; }; + ASSERT_TRUE(SubscribeTask(task_id, task_subscribe)); + + // Subscribe to the event that the given task lease is added in GCS. + std::atomic task_lease_count(0); + auto task_lease_subscribe = [&task_lease_count]( + const TaskID &task_id, + const boost::optional &data) { + if (data) { + ++task_lease_count; + } + }; + ASSERT_TRUE(SubscribeTaskLease(task_id, task_lease_subscribe)); + + ASSERT_TRUE(AddTask(task_table_data)); + ClientID node_id = ClientID::FromRandom(); + auto task_lease = Mocker::GenTaskLeaseData(task_id.Binary(), node_id.Binary()); + ASSERT_TRUE(AddTaskLease(task_lease)); + WaitPendingDone(task_count, 1); + WaitPendingDone(task_lease_count, 1); + UnsubscribeTask(task_id); + + RestartGcsServer(); + + node_id = ClientID::FromRandom(); + task_lease = Mocker::GenTaskLeaseData(task_id.Binary(), node_id.Binary()); + ASSERT_TRUE(AddTaskLease(task_lease)); + WaitPendingDone(task_lease_count, 3); + WaitPendingDone(task_count, 1); +} + TEST_F(ServiceBasedGcsClientTest, TestGcsRedisFailureDetector) { // Stop redis. TearDownTestCase(); diff --git a/src/ray/gcs/redis_accessor.cc b/src/ray/gcs/redis_accessor.cc index 239616c6f..179b7d468 100644 --- a/src/ray/gcs/redis_accessor.cc +++ b/src/ray/gcs/redis_accessor.cc @@ -443,7 +443,7 @@ Status RedisTaskInfoAccessor::AsyncGetTaskLease( auto on_failure = [callback](RedisGcsClient *client, const TaskID &task_id) { boost::optional result; - callback(Status::Invalid("Task not exist."), result); + callback(Status::Invalid("Task lease not exist."), result); }; TaskLeaseTable &task_lease_table = client_impl_->task_lease_table(); diff --git a/src/ray/gcs/redis_accessor.h b/src/ray/gcs/redis_accessor.h index cf0ab23d5..6a1bbc4be 100644 --- a/src/ray/gcs/redis_accessor.h +++ b/src/ray/gcs/redis_accessor.h @@ -243,6 +243,10 @@ class RedisTaskInfoAccessor : public TaskInfoAccessor { const std::shared_ptr &data_ptr, const StatusCallback &callback) override; + Status AsyncReSubscribe() override { + return Status::NotImplemented("AsyncReSubscribe not implemented"); + } + private: RedisGcsClient *client_impl_{nullptr}; // Use a random ClientID for task subscription. Because: diff --git a/src/ray/gcs/test/gcs_test_util.h b/src/ray/gcs/test/gcs_test_util.h index c6fd7371d..e7e0f9eeb 100644 --- a/src/ray/gcs/test/gcs_test_util.h +++ b/src/ray/gcs/test/gcs_test_util.h @@ -110,6 +110,7 @@ struct Mocker { auto task_lease_data = std::make_shared(); task_lease_data->set_task_id(task_id); task_lease_data->set_node_manager_id(node_id); + task_lease_data->set_timeout(9999); return task_lease_data; }