diff --git a/src/ray/gcs/accessor.h b/src/ray/gcs/accessor.h index 330a761b0..5497384c3 100644 --- a/src/ray/gcs/accessor.h +++ b/src/ray/gcs/accessor.h @@ -232,6 +232,16 @@ class TaskInfoAccessor { virtual Status AsyncUnsubscribeTaskLease(const TaskID &task_id, const StatusCallback &done) = 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 &data_ptr, + const StatusCallback &callback) = 0; + protected: TaskInfoAccessor() = default; }; diff --git a/src/ray/gcs/redis_accessor.cc b/src/ray/gcs/redis_accessor.cc index bffd71392..ebd3d1ebc 100644 --- a/src/ray/gcs/redis_accessor.cc +++ b/src/ray/gcs/redis_accessor.cc @@ -331,6 +331,30 @@ Status RedisTaskInfoAccessor::AsyncUnsubscribeTaskLease(const TaskID &task_id, return task_lease_sub_executor_.AsyncUnsubscribe(subscribe_id_, task_id, done); } +Status RedisTaskInfoAccessor::AttemptTaskReconstruction( + const std::shared_ptr &data_ptr, + const StatusCallback &callback) { + TaskReconstructionLog::WriteCallback on_success = nullptr; + TaskReconstructionLog::WriteCallback on_failure = nullptr; + if (callback != nullptr) { + on_success = [callback](RedisGcsClient *client, const TaskID &id, + const TaskReconstructionData &data) { + callback(Status::OK()); + }; + on_failure = [callback](RedisGcsClient *client, const TaskID &id, + const TaskReconstructionData &data) { + callback(Status::Invalid("Updating task reconstruction failed.")); + }; + } + + TaskID task_id = TaskID::FromBinary(data_ptr->task_id()); + int reconstruction_attempt = data_ptr->num_reconstructions(); + TaskReconstructionLog &task_reconstruction_log = + client_impl_->task_reconstruction_log(); + return task_reconstruction_log.AppendAt(JobID::Nil(), task_id, data_ptr, on_success, + on_failure, reconstruction_attempt); +} + RedisObjectInfoAccessor::RedisObjectInfoAccessor(RedisGcsClient *client_impl) : client_impl_(client_impl), object_sub_executor_(client_impl->object_table()) {} diff --git a/src/ray/gcs/redis_accessor.h b/src/ray/gcs/redis_accessor.h index 4a61b7c2c..f7a6529c2 100644 --- a/src/ray/gcs/redis_accessor.h +++ b/src/ray/gcs/redis_accessor.h @@ -150,6 +150,10 @@ class RedisTaskInfoAccessor : public TaskInfoAccessor { Status AsyncUnsubscribeTaskLease(const TaskID &task_id, const StatusCallback &done) override; + Status AttemptTaskReconstruction( + const std::shared_ptr &data_ptr, + const StatusCallback &callback) override; + private: RedisGcsClient *client_impl_{nullptr}; // Use a random ClientID for task subscription. Because: diff --git a/src/ray/gcs/redis_gcs_client.h b/src/ray/gcs/redis_gcs_client.h index f5d48fda6..573d563a9 100644 --- a/src/ray/gcs/redis_gcs_client.h +++ b/src/ray/gcs/redis_gcs_client.h @@ -27,6 +27,8 @@ class RAY_EXPORT RedisGcsClient : public GcsClient { friend class RedisObjectInfoAccessor; friend class SubscriptionExecutorTest; friend class LogSubscribeTestHelper; + friend class LogLookupTestHelper; + friend class LogDeleteTestHelper; friend class TaskTableTestHelper; friend class ClientTableTestHelper; friend class SetTestHelper; @@ -62,7 +64,6 @@ class RAY_EXPORT RedisGcsClient : public GcsClient { void Disconnect(); // TODO: Some API for getting the error on the driver - TaskReconstructionLog &task_reconstruction_log(); ErrorTable &error_table(); ProfileTable &profile_table(); @@ -100,9 +101,10 @@ class RAY_EXPORT RedisGcsClient : public GcsClient { HeartbeatTable &heartbeat_table(); HeartbeatBatchTable &heartbeat_batch_table(); DynamicResourceTable &resource_table(); - /// The following two methods will be deprecated, use method Tasks() instead. + /// The following three methods will be deprecated, use method Tasks() instead. raylet::TaskTable &raylet_task_table(); TaskLeaseTable &task_lease_table(); + TaskReconstructionLog &task_reconstruction_log(); // GCS command type. If CommandType::kChain, chain-replicated versions of the tables // might be used, if available. diff --git a/src/ray/gcs/test/redis_gcs_client_test.cc b/src/ray/gcs/test/redis_gcs_client_test.cc index 8f45e55c3..7d73d9ff0 100644 --- a/src/ray/gcs/test/redis_gcs_client_test.cc +++ b/src/ray/gcs/test/redis_gcs_client_test.cc @@ -364,109 +364,115 @@ class TaskTableTestHelper { TEST_TASK_TABLE_MACRO(TestGcsWithAsio, TestTableLookup); -void TestLogLookup(const JobID &job_id, std::shared_ptr client) { - // Append some entries to the log at an object ID. - TaskID task_id = RandomTaskId(); - std::vector node_manager_ids = {"abc", "def", "ghi"}; - for (auto &node_manager_id : node_manager_ids) { - auto data = std::make_shared(); - data->set_node_manager_id(node_manager_id); - // Check that we added the correct object entries. - auto add_callback = [task_id, data](gcs::RedisGcsClient *client, const TaskID &id, - const TaskReconstructionData &d) { +class LogLookupTestHelper { + public: + static void TestLogLookup(const JobID &job_id, + std::shared_ptr client) { + // Append some entries to the log at an object ID. + TaskID task_id = RandomTaskId(); + std::vector node_manager_ids = {"abc", "def", "ghi"}; + for (auto &node_manager_id : node_manager_ids) { + auto data = std::make_shared(); + data->set_node_manager_id(node_manager_id); + // Check that we added the correct object entries. + auto add_callback = [task_id, data](gcs::RedisGcsClient *client, const TaskID &id, + const TaskReconstructionData &d) { + ASSERT_EQ(id, task_id); + ASSERT_EQ(data->node_manager_id(), d.node_manager_id()); + }; + RAY_CHECK_OK( + client->task_reconstruction_log().Append(job_id, task_id, data, add_callback)); + } + + // Check that lookup returns the added object entries. + auto lookup_callback = [task_id, node_manager_ids]( + gcs::RedisGcsClient *client, const TaskID &id, + const std::vector &data) { ASSERT_EQ(id, task_id); - ASSERT_EQ(data->node_manager_id(), d.node_manager_id()); + for (const auto &entry : data) { + ASSERT_EQ(entry.node_manager_id(), node_manager_ids[test->NumCallbacks()]); + test->IncrementNumCallbacks(); + } + if (test->NumCallbacks() == node_manager_ids.size()) { + test->Stop(); + } }; + + // Do a lookup at the object ID. RAY_CHECK_OK( - client->task_reconstruction_log().Append(job_id, task_id, data, add_callback)); + client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback)); + // Run the event loop. The loop will only stop if the Lookup callback is + // called (or an assertion failure). + test->Start(); + ASSERT_EQ(test->NumCallbacks(), node_manager_ids.size()); } - // Check that lookup returns the added object entries. - auto lookup_callback = [task_id, node_manager_ids]( - gcs::RedisGcsClient *client, const TaskID &id, - const std::vector &data) { - ASSERT_EQ(id, task_id); - for (const auto &entry : data) { - ASSERT_EQ(entry.node_manager_id(), node_manager_ids[test->NumCallbacks()]); - test->IncrementNumCallbacks(); + static void TestLogAppendAt(const JobID &job_id, + std::shared_ptr client) { + TaskID task_id = RandomTaskId(); + std::vector node_manager_ids = {"A", "B"}; + std::vector> data_log; + for (const auto &node_manager_id : node_manager_ids) { + auto data = std::make_shared(); + data->set_node_manager_id(node_manager_id); + data_log.push_back(data); } - if (test->NumCallbacks() == node_manager_ids.size()) { - test->Stop(); - } - }; - // Do a lookup at the object ID. - RAY_CHECK_OK( - client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback)); - // Run the event loop. The loop will only stop if the Lookup callback is - // called (or an assertion failure). - test->Start(); - ASSERT_EQ(test->NumCallbacks(), node_manager_ids.size()); -} + // Check that we added the correct task. + auto failure_callback = [task_id](gcs::RedisGcsClient *client, const TaskID &id, + const TaskReconstructionData &d) { + ASSERT_EQ(id, task_id); + test->IncrementNumCallbacks(); + }; + + // Will succeed. + RAY_CHECK_OK(client->task_reconstruction_log().Append(job_id, task_id, + data_log.front(), + /*done callback=*/nullptr)); + // Append at index 0 will fail. + RAY_CHECK_OK(client->task_reconstruction_log().AppendAt( + job_id, task_id, data_log[1], + /*done callback=*/nullptr, failure_callback, /*log_length=*/0)); + + // Append at index 2 will fail. + RAY_CHECK_OK(client->task_reconstruction_log().AppendAt( + job_id, task_id, data_log[1], + /*done callback=*/nullptr, failure_callback, /*log_length=*/2)); + + // Append at index 1 will succeed. + RAY_CHECK_OK(client->task_reconstruction_log().AppendAt( + job_id, task_id, data_log[1], + /*done callback=*/nullptr, failure_callback, /*log_length=*/1)); + + auto lookup_callback = [node_manager_ids]( + gcs::RedisGcsClient *client, const TaskID &id, + const std::vector &data) { + std::vector appended_managers; + for (const auto &entry : data) { + appended_managers.push_back(entry.node_manager_id()); + } + ASSERT_EQ(appended_managers, node_manager_ids); + test->Stop(); + }; + RAY_CHECK_OK( + client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback)); + // Run the event loop. The loop will only stop if the Lookup callback is + // called (or an assertion failure). + test->Start(); + ASSERT_EQ(test->NumCallbacks(), 2); + } +}; TEST_F(TestGcsWithAsio, TestLogLookup) { test = this; - TestLogLookup(job_id_, client_); + LogLookupTestHelper::TestLogLookup(job_id_, client_); } TEST_TASK_TABLE_MACRO(TestGcsWithAsio, TestTableLookupFailure); -void TestLogAppendAt(const JobID &job_id, std::shared_ptr client) { - TaskID task_id = RandomTaskId(); - std::vector node_manager_ids = {"A", "B"}; - std::vector> data_log; - for (const auto &node_manager_id : node_manager_ids) { - auto data = std::make_shared(); - data->set_node_manager_id(node_manager_id); - data_log.push_back(data); - } - - // Check that we added the correct task. - auto failure_callback = [task_id](gcs::RedisGcsClient *client, const TaskID &id, - const TaskReconstructionData &d) { - ASSERT_EQ(id, task_id); - test->IncrementNumCallbacks(); - }; - - // Will succeed. - RAY_CHECK_OK(client->task_reconstruction_log().Append(job_id, task_id, data_log.front(), - /*done callback=*/nullptr)); - // Append at index 0 will fail. - RAY_CHECK_OK(client->task_reconstruction_log().AppendAt( - job_id, task_id, data_log[1], - /*done callback=*/nullptr, failure_callback, /*log_length=*/0)); - - // Append at index 2 will fail. - RAY_CHECK_OK(client->task_reconstruction_log().AppendAt( - job_id, task_id, data_log[1], - /*done callback=*/nullptr, failure_callback, /*log_length=*/2)); - - // Append at index 1 will succeed. - RAY_CHECK_OK(client->task_reconstruction_log().AppendAt( - job_id, task_id, data_log[1], - /*done callback=*/nullptr, failure_callback, /*log_length=*/1)); - - auto lookup_callback = [node_manager_ids]( - gcs::RedisGcsClient *client, const TaskID &id, - const std::vector &data) { - std::vector appended_managers; - for (const auto &entry : data) { - appended_managers.push_back(entry.node_manager_id()); - } - ASSERT_EQ(appended_managers, node_manager_ids); - test->Stop(); - }; - RAY_CHECK_OK( - client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback)); - // Run the event loop. The loop will only stop if the Lookup callback is - // called (or an assertion failure). - test->Start(); - ASSERT_EQ(test->NumCallbacks(), 2); -} - TEST_F(TestGcsWithAsio, TestLogAppendAt) { test = this; - TestLogAppendAt(job_id_, client_); + LogLookupTestHelper::TestLogAppendAt(job_id_, client_); } class SetTestHelper { @@ -806,52 +812,55 @@ TEST_F(TestGcsWithAsio, TestSet) { SetTestHelper::TestSet(job_id_, client_); } -void TestDeleteKeysFromLog( - const JobID &job_id, std::shared_ptr client, - std::vector> &data_vector) { - std::vector ids; - TaskID task_id; - for (auto &data : data_vector) { - task_id = RandomTaskId(); - ids.push_back(task_id); - // Check that we added the correct object entries. - auto add_callback = [task_id, data](gcs::RedisGcsClient *client, const TaskID &id, - const TaskReconstructionData &d) { - ASSERT_EQ(id, task_id); - ASSERT_EQ(data->node_manager_id(), d.node_manager_id()); - test->IncrementNumCallbacks(); - }; - RAY_CHECK_OK( - client->task_reconstruction_log().Append(job_id, task_id, data, add_callback)); +class LogDeleteTestHelper { + public: + static void TestDeleteKeysFromLog( + const JobID &job_id, std::shared_ptr client, + std::vector> &data_vector) { + std::vector ids; + TaskID task_id; + for (auto &data : data_vector) { + task_id = RandomTaskId(); + ids.push_back(task_id); + // Check that we added the correct object entries. + auto add_callback = [task_id, data](gcs::RedisGcsClient *client, const TaskID &id, + const TaskReconstructionData &d) { + ASSERT_EQ(id, task_id); + ASSERT_EQ(data->node_manager_id(), d.node_manager_id()); + test->IncrementNumCallbacks(); + }; + RAY_CHECK_OK( + client->task_reconstruction_log().Append(job_id, task_id, data, add_callback)); + } + for (const auto &task_id : ids) { + // Check that lookup returns the added object entries. + auto lookup_callback = [task_id, data_vector]( + gcs::RedisGcsClient *client, const TaskID &id, + const std::vector &data) { + ASSERT_EQ(id, task_id); + ASSERT_EQ(data.size(), 1); + test->IncrementNumCallbacks(); + }; + RAY_CHECK_OK( + client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback)); + } + if (ids.size() == 1) { + client->task_reconstruction_log().Delete(job_id, ids[0]); + } else { + client->task_reconstruction_log().Delete(job_id, ids); + } + for (const auto &task_id : ids) { + auto lookup_callback = [task_id](gcs::RedisGcsClient *client, const TaskID &id, + const std::vector &data) { + ASSERT_EQ(id, task_id); + ASSERT_TRUE(data.size() == 0); + test->IncrementNumCallbacks(); + }; + RAY_CHECK_OK( + client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback)); + } } - for (const auto &task_id : ids) { - // Check that lookup returns the added object entries. - auto lookup_callback = [task_id, data_vector]( - gcs::RedisGcsClient *client, const TaskID &id, - const std::vector &data) { - ASSERT_EQ(id, task_id); - ASSERT_EQ(data.size(), 1); - test->IncrementNumCallbacks(); - }; - RAY_CHECK_OK( - client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback)); - } - if (ids.size() == 1) { - client->task_reconstruction_log().Delete(job_id, ids[0]); - } else { - client->task_reconstruction_log().Delete(job_id, ids); - } - for (const auto &task_id : ids) { - auto lookup_callback = [task_id](gcs::RedisGcsClient *client, const TaskID &id, - const std::vector &data) { - ASSERT_EQ(id, task_id); - ASSERT_TRUE(data.size() == 0); - test->IncrementNumCallbacks(); - }; - RAY_CHECK_OK( - client->task_reconstruction_log().Lookup(job_id, task_id, lookup_callback)); - } -} +}; // Test delete function for keys of Log or Table. void TestDeleteKeys(const JobID &job_id, std::shared_ptr client) { @@ -867,7 +876,7 @@ void TestDeleteKeys(const JobID &job_id, std::shared_ptr cl // Test one element case. AppendTaskReconstructionData(1); ASSERT_EQ(task_reconstruction_vector.size(), 1); - TestDeleteKeysFromLog(job_id, client, task_reconstruction_vector); + LogDeleteTestHelper::TestDeleteKeysFromLog(job_id, client, task_reconstruction_vector); // Test the case for more than one elements and less than // maximum_gcs_deletion_batch_size. AppendTaskReconstructionData(RayConfig::instance().maximum_gcs_deletion_batch_size() / @@ -875,14 +884,14 @@ void TestDeleteKeys(const JobID &job_id, std::shared_ptr cl ASSERT_GT(task_reconstruction_vector.size(), 1); ASSERT_LT(task_reconstruction_vector.size(), RayConfig::instance().maximum_gcs_deletion_batch_size()); - TestDeleteKeysFromLog(job_id, client, task_reconstruction_vector); + LogDeleteTestHelper::TestDeleteKeysFromLog(job_id, client, task_reconstruction_vector); // Test the case for more than maximum_gcs_deletion_batch_size. // The Delete function will split the data into two commands. AppendTaskReconstructionData(RayConfig::instance().maximum_gcs_deletion_batch_size() / 2); ASSERT_GT(task_reconstruction_vector.size(), RayConfig::instance().maximum_gcs_deletion_batch_size()); - TestDeleteKeysFromLog(job_id, client, task_reconstruction_vector); + LogDeleteTestHelper::TestDeleteKeysFromLog(job_id, client, task_reconstruction_vector); // Test delete function for keys of Table. std::vector> task_vector; diff --git a/src/ray/protobuf/gcs.proto b/src/ray/protobuf/gcs.proto index cc436dfd3..b646ce8d4 100644 --- a/src/ray/protobuf/gcs.proto +++ b/src/ray/protobuf/gcs.proto @@ -68,10 +68,12 @@ message ObjectTableData { } message TaskReconstructionData { + // The ID of task. + bytes task_id = 1; // The number of times this task has been reconstructed so far. - uint64 num_reconstructions = 1; + uint64 num_reconstructions = 2; // The node manager that is trying to reconstruct the task. - bytes node_manager_id = 2; + bytes node_manager_id = 3; } message TaskTableData { diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index 25d21825a..c6aab817e 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -98,8 +98,7 @@ NodeManager::NodeManager(boost::asio::io_service &io_service, HandleTaskReconstruction(task_id, required_object_id); }, RayConfig::instance().initial_reconstruction_timeout_milliseconds(), - self_node_id_, gcs_client_, object_directory_, - gcs_client_->task_reconstruction_log()), + self_node_id_, gcs_client_, object_directory_), task_dependency_manager_( object_manager, reconstruction_policy_, io_service, self_node_id_, RayConfig::instance().initial_reconstruction_timeout_milliseconds(), diff --git a/src/ray/raylet/reconstruction_policy.cc b/src/ray/raylet/reconstruction_policy.cc index 6302bf4c9..d53331cfa 100644 --- a/src/ray/raylet/reconstruction_policy.cc +++ b/src/ray/raylet/reconstruction_policy.cc @@ -11,15 +11,13 @@ ReconstructionPolicy::ReconstructionPolicy( std::function reconstruction_handler, int64_t initial_reconstruction_timeout_ms, const ClientID &client_id, std::shared_ptr gcs_client, - std::shared_ptr object_directory, - gcs::LogInterface &task_reconstruction_log) + std::shared_ptr object_directory) : io_service_(io_service), reconstruction_handler_(reconstruction_handler), initial_reconstruction_timeout_ms_(initial_reconstruction_timeout_ms), client_id_(client_id), gcs_client_(gcs_client), - object_directory_(std::move(object_directory)), - task_reconstruction_log_(task_reconstruction_log) {} + object_directory_(std::move(object_directory)) {} void ReconstructionPolicy::SetTaskTimeout( std::unordered_map::iterator task_it, @@ -134,21 +132,19 @@ void ReconstructionPolicy::AttemptReconstruction(const TaskID &task_id, // reconstruction log. This will fail if another node has already inserted // an entry for this reconstruction. auto reconstruction_entry = std::make_shared(); + reconstruction_entry->set_task_id(task_id.Binary()); reconstruction_entry->set_num_reconstructions(reconstruction_attempt); reconstruction_entry->set_node_manager_id(client_id_.Binary()); - RAY_CHECK_OK(task_reconstruction_log_.AppendAt( - JobID::Nil(), task_id, reconstruction_entry, - /*success_callback=*/ - [this, required_object_id](gcs::RedisGcsClient *client, const TaskID &task_id, - const TaskReconstructionData &data) { - HandleReconstructionLogAppend(task_id, required_object_id, /*success=*/true); - }, - /*failure_callback=*/ - [this, required_object_id](gcs::RedisGcsClient *client, const TaskID &task_id, - const TaskReconstructionData &data) { - HandleReconstructionLogAppend(task_id, required_object_id, /*success=*/false); - }, - reconstruction_attempt)); + RAY_CHECK_OK(gcs_client_->Tasks().AttemptTaskReconstruction( + reconstruction_entry, + /*done=*/ + [this, task_id, required_object_id](Status status) { + if (status.ok()) { + HandleReconstructionLogAppend(task_id, required_object_id, /*success=*/true); + } else { + HandleReconstructionLogAppend(task_id, required_object_id, /*success=*/false); + } + })); // Increment the number of times reconstruction has been attempted. This is // used to suppress duplicate reconstructions of the same task. If diff --git a/src/ray/raylet/reconstruction_policy.h b/src/ray/raylet/reconstruction_policy.h index be11ff1d0..5b6438365 100644 --- a/src/ray/raylet/reconstruction_policy.h +++ b/src/ray/raylet/reconstruction_policy.h @@ -44,8 +44,7 @@ class ReconstructionPolicy : public ReconstructionPolicyInterface { std::function reconstruction_handler, int64_t initial_reconstruction_timeout_ms, const ClientID &client_id, std::shared_ptr gcs_client, - std::shared_ptr object_directory, - gcs::LogInterface &task_reconstruction_log); + std::shared_ptr object_directory); /// Listen for task lease notifications about an object that may require /// reconstruction. If no notifications are received within the initial @@ -146,7 +145,6 @@ class ReconstructionPolicy : public ReconstructionPolicyInterface { std::shared_ptr gcs_client_; /// The object directory used to lookup object locations. std::shared_ptr object_directory_; - gcs::LogInterface &task_reconstruction_log_; /// The tasks that we are currently subscribed to in the GCS. std::unordered_map listening_tasks_; }; diff --git a/src/ray/raylet/reconstruction_policy_test.cc b/src/ray/raylet/reconstruction_policy_test.cc index 829ad11a6..0e6a27bac 100644 --- a/src/ray/raylet/reconstruction_policy_test.cc +++ b/src/ray/raylet/reconstruction_policy_test.cc @@ -131,14 +131,33 @@ class MockTaskInfoAccessor : public gcs::RedisTaskInfoAccessor { return Status::OK(); } + Status AttemptTaskReconstruction( + const std::shared_ptr &task_data, + const gcs::StatusCallback &done) override { + int log_index = task_data->num_reconstructions(); + TaskID task_id = TaskID::FromBinary(task_data->task_id()); + if (task_reconstruction_log_[task_id].size() == static_cast(log_index)) { + task_reconstruction_log_[task_id].push_back(*task_data); + if (done != nullptr) { + done(Status::OK()); + } + } else { + if (done != nullptr) { + done(Status::Invalid("Updating task reconstruction failed.")); + } + } + return Status::OK(); + } + private: gcs::SubscribeCallback> subscribe_callback_; std::unordered_map> task_lease_table_; std::unordered_set subscribed_tasks_; + std::unordered_map> + task_reconstruction_log_; }; -class MockGcs : public gcs::RedisGcsClient, - public ray::gcs::LogInterface { +class MockGcs : public gcs::RedisGcsClient { public: MockGcs() : gcs::RedisGcsClient(gcs::GcsClientOptions("", 0, "")){}; @@ -146,37 +165,6 @@ class MockGcs : public gcs::RedisGcsClient, task_accessor_.reset(task_accessor); node_accessor_.reset(node_accessor); } - - Status AppendAt( - const JobID &job_id, const TaskID &task_id, - const std::shared_ptr &task_data, - const ray::gcs::LogInterface::WriteCallback - &success_callback, - const ray::gcs::LogInterface::WriteCallback - &failure_callback, - int log_index) { - if (task_reconstruction_log_[task_id].size() == static_cast(log_index)) { - task_reconstruction_log_[task_id].push_back(*task_data); - if (success_callback != nullptr) { - success_callback(nullptr, task_id, *task_data); - } - } else { - if (failure_callback != nullptr) { - failure_callback(nullptr, task_id, *task_data); - } - } - return Status::OK(); - } - - MOCK_METHOD4( - Append, - ray::Status( - const JobID &, const TaskID &, const std::shared_ptr &, - const ray::gcs::LogInterface::WriteCallback &)); - - private: - std::unordered_map> - task_reconstruction_log_; }; class ReconstructionPolicyTest : public ::testing::Test { @@ -194,7 +182,7 @@ class ReconstructionPolicyTest : public ::testing::Test { TriggerReconstruction(task_id); }, reconstruction_timeout_ms_, ClientID::FromRandom(), mock_gcs_, - mock_object_directory_, *mock_gcs_)), + mock_object_directory_)), timer_canceled_(false) { subscribe_callback_ = [this](const TaskID &task_id, const boost::optional &task_lease) { @@ -449,14 +437,13 @@ TEST_F(ReconstructionPolicyTest, TestSimultaneousReconstructionSuppressed) { // reconstruction first. This should suppress this node's first attempt at // reconstruction. auto task_reconstruction_data = std::make_shared(); + task_reconstruction_data->set_task_id(task_id.Binary()); task_reconstruction_data->set_node_manager_id(ClientID::FromRandom().Binary()); task_reconstruction_data->set_num_reconstructions(0); - RAY_CHECK_OK( - mock_gcs_->AppendAt(JobID::Nil(), task_id, task_reconstruction_data, nullptr, - /*failure_callback=*/ - [](ray::gcs::RedisGcsClient *client, const TaskID &task_id, - const TaskReconstructionData &data) { ASSERT_TRUE(false); }, - /*log_index=*/0)); + RAY_CHECK_OK(mock_gcs_->Tasks().AttemptTaskReconstruction( + task_reconstruction_data, + /*done=*/ + [](Status status) { ASSERT_TRUE(status.ok()); })); // Listen for an object. reconstruction_policy_->ListenAndMaybeReconstruct(object_id);