From cc110ff1e3c17af491474220533868173e4497e7 Mon Sep 17 00:00:00 2001 From: micafan <550435771@qq.com> Date: Sun, 5 Jan 2020 13:54:33 +0800 Subject: [PATCH] [GCS]Add task lease methods to TaskInfoAccessor (#6645) --- src/ray/gcs/accessor.h | 30 ++++- src/ray/gcs/redis_accessor.cc | 29 ++++- src/ray/gcs/redis_accessor.h | 28 +++- src/ray/gcs/redis_gcs_client.h | 4 +- src/ray/gcs/subscription_executor.cc | 2 + src/ray/gcs/subscription_executor.h | 2 +- src/ray/gcs/tables.cc | 19 +++ src/ray/gcs/tables.h | 11 ++ src/ray/protobuf/gcs.proto | 8 +- src/ray/raylet/node_manager.cc | 29 +---- src/ray/raylet/reconstruction_policy.cc | 41 ++++-- src/ray/raylet/reconstruction_policy.h | 12 +- src/ray/raylet/reconstruction_policy_test.cc | 120 +++++++++++------- src/ray/raylet/task_dependency_manager.cc | 8 +- src/ray/raylet/task_dependency_manager.h | 7 +- .../raylet/task_dependency_manager_test.cc | 50 +++++--- 16 files changed, 280 insertions(+), 120 deletions(-) diff --git a/src/ray/gcs/accessor.h b/src/ray/gcs/accessor.h index 2f0b0251a..330a761b0 100644 --- a/src/ray/gcs/accessor.h +++ b/src/ray/gcs/accessor.h @@ -197,13 +197,41 @@ class TaskInfoAccessor { const StatusCallback &done) = 0; /// Cancel subscription to a task asynchronously. - /// This method is for node only (core worker shouldn't use this method). /// /// \param task_id The ID of the task to be unsubscribed to. /// \param done Callback that will be called when unsubscribe is complete. /// \return Status virtual Status AsyncUnsubscribe(const TaskID &task_id, const StatusCallback &done) = 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 &data_ptr, + const StatusCallback &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> &subscribe, + const StatusCallback &done) = 0; + + /// Cancel subscription to a task lease asynchronously. + /// + /// \param task_id The ID of the task to be unsubscribed to. + /// \param done Callback that will be called when unsubscribe is complete. + /// \return Status + virtual Status AsyncUnsubscribeTaskLease(const TaskID &task_id, + const StatusCallback &done) = 0; + protected: TaskInfoAccessor() = default; }; diff --git a/src/ray/gcs/redis_accessor.cc b/src/ray/gcs/redis_accessor.cc index a84da3923..a1fb72aab 100644 --- a/src/ray/gcs/redis_accessor.cc +++ b/src/ray/gcs/redis_accessor.cc @@ -247,7 +247,9 @@ Status RedisJobInfoAccessor::AsyncSubscribeToFinishedJobs( } RedisTaskInfoAccessor::RedisTaskInfoAccessor(RedisGcsClient *client_impl) - : client_impl_(client_impl), task_sub_executor_(client_impl->raylet_task_table()) {} + : client_impl_(client_impl), + task_sub_executor_(client_impl->raylet_task_table()), + task_lease_sub_executor_(client_impl->task_lease_table()) {} Status RedisTaskInfoAccessor::AsyncAdd(const std::shared_ptr &data_ptr, const StatusCallback &callback) { @@ -301,6 +303,31 @@ Status RedisTaskInfoAccessor::AsyncUnsubscribe(const TaskID &task_id, return task_sub_executor_.AsyncUnsubscribe(subscribe_id_, task_id, done); } +Status RedisTaskInfoAccessor::AsyncAddTaskLease( + const std::shared_ptr &data_ptr, const StatusCallback &callback) { + TaskLeaseTable::WriteCallback on_done = nullptr; + if (callback != nullptr) { + on_done = [callback](RedisGcsClient *client, const TaskID &id, + const TaskLeaseData &data) { callback(Status::OK()); }; + } + TaskID task_id = TaskID::FromBinary(data_ptr->task_id()); + TaskLeaseTable &task_lease_table = client_impl_->task_lease_table(); + return task_lease_table.Add(JobID::Nil(), task_id, data_ptr, on_done); +} + +Status RedisTaskInfoAccessor::AsyncSubscribeTaskLease( + const TaskID &task_id, + const SubscribeCallback> &subscribe, + const StatusCallback &done) { + RAY_CHECK(subscribe != nullptr); + return task_lease_sub_executor_.AsyncSubscribe(subscribe_id_, task_id, subscribe, done); +} + +Status RedisTaskInfoAccessor::AsyncUnsubscribeTaskLease(const TaskID &task_id, + const StatusCallback &done) { + return task_lease_sub_executor_.AsyncUnsubscribe(subscribe_id_, task_id, done); +} + 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 d82259e96..4a61b7c2c 100644 --- a/src/ray/gcs/redis_accessor.h +++ b/src/ray/gcs/redis_accessor.h @@ -122,21 +122,33 @@ class RedisTaskInfoAccessor : public TaskInfoAccessor { public: explicit RedisTaskInfoAccessor(RedisGcsClient *client_impl); - ~RedisTaskInfoAccessor() {} + virtual ~RedisTaskInfoAccessor() {} Status AsyncAdd(const std::shared_ptr &data_ptr, - const StatusCallback &callback); + const StatusCallback &callback) override; Status AsyncGet(const TaskID &task_id, - const OptionalItemCallback &callback); + const OptionalItemCallback &callback) override; - Status AsyncDelete(const std::vector &task_ids, const StatusCallback &callback); + Status AsyncDelete(const std::vector &task_ids, + const StatusCallback &callback) override; Status AsyncSubscribe(const TaskID &task_id, const SubscribeCallback &subscribe, - const StatusCallback &done); + const StatusCallback &done) override; - Status AsyncUnsubscribe(const TaskID &task_id, const StatusCallback &done); + Status AsyncUnsubscribe(const TaskID &task_id, const StatusCallback &done) override; + + Status AsyncAddTaskLease(const std::shared_ptr &data_ptr, + const StatusCallback &callback) override; + + Status AsyncSubscribeTaskLease( + const TaskID &task_id, + const SubscribeCallback> &subscribe, + const StatusCallback &done) override; + + Status AsyncUnsubscribeTaskLease(const TaskID &task_id, + const StatusCallback &done) override; private: RedisGcsClient *client_impl_{nullptr}; @@ -151,6 +163,10 @@ class RedisTaskInfoAccessor : public TaskInfoAccessor { typedef SubscriptionExecutor TaskSubscriptionExecutor; TaskSubscriptionExecutor task_sub_executor_; + + typedef SubscriptionExecutor, TaskLeaseTable> + TaskLeaseSubscriptionExecutor; + TaskLeaseSubscriptionExecutor task_lease_sub_executor_; }; /// \class RedisObjectInfoAccessor diff --git a/src/ray/gcs/redis_gcs_client.h b/src/ray/gcs/redis_gcs_client.h index c0a135350..f5d48fda6 100644 --- a/src/ray/gcs/redis_gcs_client.h +++ b/src/ray/gcs/redis_gcs_client.h @@ -63,7 +63,6 @@ class RAY_EXPORT RedisGcsClient : public GcsClient { // TODO: Some API for getting the error on the driver TaskReconstructionLog &task_reconstruction_log(); - TaskLeaseTable &task_lease_table(); ErrorTable &error_table(); ProfileTable &profile_table(); @@ -101,8 +100,9 @@ class RAY_EXPORT RedisGcsClient : public GcsClient { HeartbeatTable &heartbeat_table(); HeartbeatBatchTable &heartbeat_batch_table(); DynamicResourceTable &resource_table(); - /// This method will be deprecated, use method Tasks() instead. + /// The following two methods will be deprecated, use method Tasks() instead. raylet::TaskTable &raylet_task_table(); + TaskLeaseTable &task_lease_table(); // GCS command type. If CommandType::kChain, chain-replicated versions of the tables // might be used, if available. diff --git a/src/ray/gcs/subscription_executor.cc b/src/ray/gcs/subscription_executor.cc index c6e0f4e4e..862326f04 100644 --- a/src/ray/gcs/subscription_executor.cc +++ b/src/ray/gcs/subscription_executor.cc @@ -188,6 +188,8 @@ template class SubscriptionExecutor; template class SubscriptionExecutor; template class SubscriptionExecutor; template class SubscriptionExecutor; +template class SubscriptionExecutor, + TaskLeaseTable>; template class SubscriptionExecutor; template class SubscriptionExecutor; diff --git a/src/ray/gcs/subscription_executor.h b/src/ray/gcs/subscription_executor.h index 83921d52b..47f5fabd9 100644 --- a/src/ray/gcs/subscription_executor.h +++ b/src/ray/gcs/subscription_executor.h @@ -18,7 +18,7 @@ namespace gcs { template class SubscriptionExecutor { public: - SubscriptionExecutor(Table &table) : table_(table) {} + explicit SubscriptionExecutor(Table &table) : table_(table) {} ~SubscriptionExecutor() {} diff --git a/src/ray/gcs/tables.cc b/src/ray/gcs/tables.cc index 6abf60deb..b28731114 100644 --- a/src/ray/gcs/tables.cc +++ b/src/ray/gcs/tables.cc @@ -730,6 +730,25 @@ std::string ClientTable::DebugString() const { return result.str(); } +Status TaskLeaseTable::Subscribe(const JobID &job_id, const ClientID &client_id, + const Callback &subscribe, + const SubscriptionCallback &done) { + auto on_subscribe = [subscribe](RedisGcsClient *client, const TaskID &task_id, + const std::vector &data) { + std::vector> result; + for (const auto &item : data) { + boost::optional optional_item(item); + result.emplace_back(std::move(optional_item)); + } + if (result.empty()) { + boost::optional optional_item; + result.emplace_back(std::move(optional_item)); + } + subscribe(client, task_id, result); + }; + return Table::Subscribe(job_id, client_id, on_subscribe, done); +} + Status ActorCheckpointIdTable::AddCheckpointId(const JobID &job_id, const ActorID &actor_id, const ActorCheckpointID &checkpoint_id, diff --git a/src/ray/gcs/tables.h b/src/ray/gcs/tables.h index b9d9fdac7..62c722ef3 100644 --- a/src/ray/gcs/tables.h +++ b/src/ray/gcs/tables.h @@ -727,6 +727,12 @@ class TaskReconstructionLog : public Log { class TaskLeaseTable : public Table { public: + /// Use boost::optional to represent subscription results, so that we can + /// notify raylet whether the entry of task lease is empty. + using Callback = + std::function> &data)>; + TaskLeaseTable(const std::vector> &contexts, RedisGcsClient *client) : Table(contexts, client) { @@ -749,6 +755,11 @@ class TaskLeaseTable : public Table { return GetRedisContext(id)->RunArgvAsync(args); } + + /// Implement this method for the subscription tools class SubscriptionExecutor. + /// In this way TaskLeaseTable() can also reuse class SubscriptionExecutor. + Status Subscribe(const JobID &job_id, const ClientID &client_id, + const Callback &subscribe, const SubscriptionCallback &done); }; class ActorCheckpointTable : public Table { diff --git a/src/ray/protobuf/gcs.proto b/src/ray/protobuf/gcs.proto index 280aab56f..cc436dfd3 100644 --- a/src/ray/protobuf/gcs.proto +++ b/src/ray/protobuf/gcs.proto @@ -213,14 +213,16 @@ message HeartbeatBatchTableData { // Data for a lease on task execution. message TaskLeaseData { + // The task ID. + bytes task_id = 1; // Node manager client ID. - bytes node_manager_id = 1; + bytes node_manager_id = 2; // The time that the lease was last acquired at. NOTE(swang): This is the // system clock time according to the node that added the entry and is not // synchronized with other nodes. - uint64 acquired_at = 2; + uint64 acquired_at = 3; // The period that the lease is active for. - uint64 timeout = 3; + uint64 timeout = 4; } message JobTableData { diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index c9adfbdba..736fb4b0e 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -96,12 +96,12 @@ 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_->task_lease_table(), object_directory_, + self_node_id_, gcs_client_, object_directory_, gcs_client_->task_reconstruction_log()), task_dependency_manager_( object_manager, reconstruction_policy_, io_service, self_node_id_, RayConfig::instance().initial_reconstruction_timeout_milliseconds(), - gcs_client_->task_lease_table()), + gcs_client_), lineage_cache_(self_node_id_, gcs_client_, config.max_lineage_size), actor_registry_(), node_manager_server_("NodeManager", config.node_manager_port), @@ -136,30 +136,7 @@ NodeManager::NodeManager(boost::asio::io_service &io_service, } ray::Status NodeManager::RegisterGcs() { - const auto task_lease_notification_callback = [this](gcs::RedisGcsClient *client, - const TaskID &task_id, - const TaskLeaseData &task_lease) { - const ClientID node_manager_id = ClientID::FromBinary(task_lease.node_manager_id()); - if (gcs_client_->Nodes().IsRemoved(node_manager_id)) { - // The node manager that added the task lease is already removed. The - // lease is considered inactive. - reconstruction_policy_.HandleTaskLeaseNotification(task_id, 0); - } else { - // NOTE(swang): The task_lease.timeout is an overestimate of the lease's - // expiration period since the entry may have been in the GCS for some - // time already. For a more accurate estimate, the age of the entry in - // the GCS should be subtracted from task_lease.timeout. - reconstruction_policy_.HandleTaskLeaseNotification(task_id, task_lease.timeout()); - } - }; - const auto task_lease_empty_callback = [this](gcs::RedisGcsClient *client, - const TaskID &task_id) { - reconstruction_policy_.HandleTaskLeaseNotification(task_id, 0); - }; - RAY_RETURN_NOT_OK(gcs_client_->task_lease_table().Subscribe( - JobID::Nil(), self_node_id_, task_lease_notification_callback, - task_lease_empty_callback, nullptr)); - + // The TaskLease subscription is done on demand in reconstruction policy. // Register a callback to handle actor notifications. auto actor_notification_callback = [this](const ActorID &actor_id, const ActorTableData &data) { diff --git a/src/ray/raylet/reconstruction_policy.cc b/src/ray/raylet/reconstruction_policy.cc index 020f41437..6302bf4c9 100644 --- a/src/ray/raylet/reconstruction_policy.cc +++ b/src/ray/raylet/reconstruction_policy.cc @@ -10,14 +10,14 @@ ReconstructionPolicy::ReconstructionPolicy( boost::asio::io_service &io_service, std::function reconstruction_handler, int64_t initial_reconstruction_timeout_ms, const ClientID &client_id, - gcs::PubsubInterface &task_lease_pubsub, + std::shared_ptr gcs_client, std::shared_ptr object_directory, gcs::LogInterface &task_reconstruction_log) : io_service_(io_service), reconstruction_handler_(reconstruction_handler), initial_reconstruction_timeout_ms_(initial_reconstruction_timeout_ms), client_id_(client_id), - task_lease_pubsub_(task_lease_pubsub), + gcs_client_(gcs_client), object_directory_(std::move(object_directory)), task_reconstruction_log_(task_reconstruction_log) {} @@ -43,6 +43,11 @@ void ReconstructionPolicy::SetTaskTimeout( // received. The current lease is now considered expired. HandleTaskLeaseExpired(task_id); } else { + const auto task_lease_notification_callback = + [this](const TaskID &task_id, + const boost::optional &task_lease) { + OnTaskLeaseNotification(task_id, task_lease); + }; // This task is still required, so subscribe to task lease // notifications. Reconstruction will be triggered if the current // task lease expires, or if no one has acquired the task lease. @@ -52,9 +57,8 @@ void ReconstructionPolicy::SetTaskTimeout( // required by the task are no longer needed soon after. If the // task is still required after this initial period, then we now // subscribe to task lease notifications. - RAY_CHECK_OK(task_lease_pubsub_.RequestNotifications(JobID::Nil(), task_id, - client_id_, - /*done*/ nullptr)); + RAY_CHECK_OK(gcs_client_->Tasks().AsyncSubscribeTaskLease( + task_id, task_lease_notification_callback, /*done*/ nullptr)); it->second.subscribed = true; } } else { @@ -64,6 +68,28 @@ void ReconstructionPolicy::SetTaskTimeout( }); } +void ReconstructionPolicy::OnTaskLeaseNotification( + const TaskID &task_id, const boost::optional &task_lease) { + if (!task_lease) { + // Task lease not exist. + HandleTaskLeaseNotification(task_id, 0); + return; + } + + const ClientID node_manager_id = ClientID::FromBinary(task_lease->node_manager_id()); + if (gcs_client_->Nodes().IsRemoved(node_manager_id)) { + // The node manager that added the task lease is already removed. The + // lease is considered inactive. + HandleTaskLeaseNotification(task_id, 0); + } else { + // NOTE(swang): The task_lease.timeout is an overestimate of the + // lease's expiration period since the entry may have been in the GCS + // for some time already. For a more accurate estimate, the age of the + // entry in the GCS should be subtracted from task_lease.timeout. + HandleTaskLeaseNotification(task_id, task_lease->timeout()); + } +} + void ReconstructionPolicy::HandleReconstructionLogAppend( const TaskID &task_id, const ObjectID &required_object_id, bool success) { auto it = listening_tasks_.find(task_id); @@ -201,9 +227,8 @@ void ReconstructionPolicy::Cancel(const ObjectID &object_id) { if (it->second.created_objects.empty()) { // Cancel notifications for the task lease if we were subscribed to them. if (it->second.subscribed) { - RAY_CHECK_OK(task_lease_pubsub_.CancelNotifications(JobID::Nil(), task_id, - client_id_, - /*done*/ nullptr)); + RAY_CHECK_OK( + gcs_client_->Tasks().AsyncUnsubscribeTaskLease(task_id, /*done*/ nullptr)); } listening_tasks_.erase(it); } diff --git a/src/ray/raylet/reconstruction_policy.h b/src/ray/raylet/reconstruction_policy.h index dd9d23cb5..be11ff1d0 100644 --- a/src/ray/raylet/reconstruction_policy.h +++ b/src/ray/raylet/reconstruction_policy.h @@ -37,13 +37,13 @@ class ReconstructionPolicy : public ReconstructionPolicyInterface { /// will be triggered. /// \param client_id The client ID to use when requesting notifications from /// the GCS. - /// \param task_lease_pubsub The GCS pub-sub storage system to request task + /// \param gcs_client The Client of GCS. /// lease notifications from. ReconstructionPolicy( boost::asio::io_service &io_service, std::function reconstruction_handler, int64_t initial_reconstruction_timeout_ms, const ClientID &client_id, - gcs::PubsubInterface &task_lease_pubsub, + std::shared_ptr gcs_client, std::shared_ptr object_directory, gcs::LogInterface &task_reconstruction_log); @@ -109,6 +109,10 @@ class ReconstructionPolicy : public ReconstructionPolicyInterface { void SetTaskTimeout(std::unordered_map::iterator task_it, int64_t timeout_ms); + /// Handle task lease notification from GCS. + void OnTaskLeaseNotification(const TaskID &task_id, + const boost::optional &task_lease); + /// Attempt to re-execute a task to reconstruct the required object. /// /// \param task_id The task to attempt to re-execute. @@ -138,8 +142,8 @@ class ReconstructionPolicy : public ReconstructionPolicyInterface { const int64_t initial_reconstruction_timeout_ms_; /// The client ID to use when requesting notifications from the GCS. const ClientID client_id_; - /// The GCS pub-sub storage system to request task lease notifications from. - gcs::PubsubInterface &task_lease_pubsub_; + /// A client connection to the GCS. + std::shared_ptr gcs_client_; /// The object directory used to lookup object locations. std::shared_ptr object_directory_; gcs::LogInterface &task_reconstruction_log_; diff --git a/src/ray/raylet/reconstruction_policy_test.cc b/src/ray/raylet/reconstruction_policy_test.cc index 8556d51b4..829ad11a6 100644 --- a/src/ray/raylet/reconstruction_policy_test.cc +++ b/src/ray/raylet/reconstruction_policy_test.cc @@ -7,6 +7,7 @@ #include #include "ray/gcs/callback.h" +#include "ray/gcs/redis_accessor.h" #include "ray/raylet/format/node_manager_generated.h" #include "ray/raylet/reconstruction_policy.h" @@ -84,44 +85,68 @@ class MockObjectDirectory : public ObjectDirectoryInterface { std::unordered_map> locations_; }; -class MockGcs : public gcs::PubsubInterface, - public ray::gcs::LogInterface { +class MockNodeInfoAccessor : public gcs::RedisNodeInfoAccessor { public: - MockGcs() : notification_callback_(nullptr), failure_callback_(nullptr){}; + MockNodeInfoAccessor(gcs::RedisGcsClient *client) + : gcs::RedisNodeInfoAccessor(client) {} - void Subscribe(const gcs::TaskLeaseTable::WriteCallback ¬ification_callback, - const gcs::TaskLeaseTable::FailureCallback &failure_callback) { - notification_callback_ = notification_callback; - failure_callback_ = failure_callback; - } + bool IsRemoved(const ClientID &node_id) const override { return false; } +}; - void Add(const JobID &job_id, const TaskID &task_id, - const std::shared_ptr &task_lease_data) { - task_lease_table_[task_id] = task_lease_data; - if (subscribed_tasks_.count(task_id) == 1) { - notification_callback_(nullptr, task_id, *task_lease_data); - } - } +class MockTaskInfoAccessor : public gcs::RedisTaskInfoAccessor { + public: + MockTaskInfoAccessor(gcs::RedisGcsClient *client) : RedisTaskInfoAccessor(client) {} - Status RequestNotifications(const JobID &job_id, const TaskID &task_id, - const ClientID &client_id, - const gcs::StatusCallback &done) { + Status AsyncSubscribeTaskLease( + const TaskID &task_id, + const gcs::SubscribeCallback> &subscribe, + const gcs::StatusCallback &done) override { + subscribe_callback_ = subscribe; subscribed_tasks_.insert(task_id); auto entry = task_lease_table_.find(task_id); if (entry == task_lease_table_.end()) { - failure_callback_(nullptr, task_id); + boost::optional result; + subscribe(task_id, result); } else { - notification_callback_(nullptr, task_id, *entry->second); + boost::optional result(*entry->second); + subscribe(task_id, result); } return ray::Status::OK(); } - Status CancelNotifications(const JobID &job_id, const TaskID &task_id, - const ClientID &client_id, const gcs::StatusCallback &done) { + Status AsyncUnsubscribeTaskLease(const TaskID &task_id, + const gcs::StatusCallback &done) override { subscribed_tasks_.erase(task_id); return ray::Status::OK(); } + Status AsyncAddTaskLease(const std::shared_ptr &task_lease_data, + const gcs::StatusCallback &done) override { + TaskID task_id = TaskID::FromBinary(task_lease_data->task_id()); + task_lease_table_[task_id] = task_lease_data; + if (subscribed_tasks_.count(task_id) == 1) { + boost::optional result(*task_lease_data); + subscribe_callback_(task_id, result); + } + return Status::OK(); + } + + private: + gcs::SubscribeCallback> subscribe_callback_; + std::unordered_map> task_lease_table_; + std::unordered_set subscribed_tasks_; +}; + +class MockGcs : public gcs::RedisGcsClient, + public ray::gcs::LogInterface { + public: + MockGcs() : gcs::RedisGcsClient(gcs::GcsClientOptions("", 0, "")){}; + + void Init(gcs::TaskInfoAccessor *task_accessor, gcs::NodeInfoAccessor *node_accessor) { + 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, @@ -150,10 +175,6 @@ class MockGcs : public gcs::PubsubInterface, const ray::gcs::LogInterface::WriteCallback &)); private: - gcs::TaskLeaseTable::WriteCallback notification_callback_; - gcs::TaskLeaseTable::FailureCallback failure_callback_; - std::unordered_map> task_lease_table_; - std::unordered_set subscribed_tasks_; std::unordered_map> task_reconstruction_log_; }; @@ -162,7 +183,9 @@ class ReconstructionPolicyTest : public ::testing::Test { public: ReconstructionPolicyTest() : io_service_(), - mock_gcs_(), + mock_gcs_(new MockGcs()), + task_accessor_(new MockTaskInfoAccessor(mock_gcs_.get())), + node_accessor_(new MockNodeInfoAccessor(mock_gcs_.get())), mock_object_directory_(std::make_shared()), reconstruction_timeout_ms_(50), reconstruction_policy_(std::make_shared( @@ -171,17 +194,19 @@ class ReconstructionPolicyTest : public ::testing::Test { TriggerReconstruction(task_id); }, reconstruction_timeout_ms_, ClientID::FromRandom(), mock_gcs_, - mock_object_directory_, mock_gcs_)), + mock_object_directory_, *mock_gcs_)), timer_canceled_(false) { - mock_gcs_.Subscribe( - [this](gcs::RedisGcsClient *client, const TaskID &task_id, - const TaskLeaseData &task_lease) { - reconstruction_policy_->HandleTaskLeaseNotification(task_id, - task_lease.timeout()); - }, - [this](gcs::RedisGcsClient *client, const TaskID &task_id) { - reconstruction_policy_->HandleTaskLeaseNotification(task_id, 0); - }); + subscribe_callback_ = [this](const TaskID &task_id, + const boost::optional &task_lease) { + if (task_lease) { + reconstruction_policy_->HandleTaskLeaseNotification(task_id, + task_lease->timeout()); + } else { + reconstruction_policy_->HandleTaskLeaseNotification(task_id, 0); + } + }; + + mock_gcs_->Init(task_accessor_, node_accessor_); } void TriggerReconstruction(const TaskID &task_id) { reconstructed_tasks_[task_id]++; } @@ -230,7 +255,10 @@ class ReconstructionPolicyTest : public ::testing::Test { protected: boost::asio::io_service io_service_; - MockGcs mock_gcs_; + std::shared_ptr mock_gcs_; + MockTaskInfoAccessor *task_accessor_; + MockNodeInfoAccessor *node_accessor_; + gcs::SubscribeCallback> subscribe_callback_; std::shared_ptr mock_object_directory_; uint64_t reconstruction_timeout_ms_; std::shared_ptr reconstruction_policy_; @@ -340,7 +368,8 @@ TEST_F(ReconstructionPolicyTest, TestReconstructionSuppressed) { task_lease_data->set_node_manager_id(ClientID::FromRandom().Binary()); task_lease_data->set_acquired_at(absl::GetCurrentTimeNanos() / 1000000); task_lease_data->set_timeout(2 * test_period); - mock_gcs_.Add(JobID::Nil(), task_id, task_lease_data); + task_lease_data->set_task_id(task_id.Binary()); + RAY_CHECK_OK(mock_gcs_->Tasks().AsyncAddTaskLease(task_lease_data, nullptr)); // Listen for an object. reconstruction_policy_->ListenAndMaybeReconstruct(object_id); @@ -368,7 +397,8 @@ TEST_F(ReconstructionPolicyTest, TestReconstructionContinuallySuppressed) { task_lease_data->set_node_manager_id(ClientID::FromRandom().Binary()); task_lease_data->set_acquired_at(absl::GetCurrentTimeNanos() / 1000000); task_lease_data->set_timeout(reconstruction_timeout_ms_); - mock_gcs_.Add(JobID::Nil(), task_id, task_lease_data); + task_lease_data->set_task_id(task_id.Binary()); + RAY_CHECK_OK(mock_gcs_->Tasks().AsyncAddTaskLease(task_lease_data, nullptr)); }); // Run the test for much longer than the reconstruction timeout. Run(reconstruction_timeout_ms_ * 2); @@ -422,11 +452,11 @@ TEST_F(ReconstructionPolicyTest, TestSimultaneousReconstructionSuppressed) { 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)); + 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)); // Listen for an object. reconstruction_policy_->ListenAndMaybeReconstruct(object_id); diff --git a/src/ray/raylet/task_dependency_manager.cc b/src/ray/raylet/task_dependency_manager.cc index 258e4f5db..b6ea1575e 100644 --- a/src/ray/raylet/task_dependency_manager.cc +++ b/src/ray/raylet/task_dependency_manager.cc @@ -12,14 +12,13 @@ TaskDependencyManager::TaskDependencyManager( ObjectManagerInterface &object_manager, ReconstructionPolicyInterface &reconstruction_policy, boost::asio::io_service &io_service, const ClientID &client_id, - int64_t initial_lease_period_ms, - gcs::TableInterface &task_lease_table) + int64_t initial_lease_period_ms, std::shared_ptr gcs_client) : object_manager_(object_manager), reconstruction_policy_(reconstruction_policy), io_service_(io_service), client_id_(client_id), initial_lease_period_ms_(initial_lease_period_ms), - task_lease_table_(task_lease_table) {} + gcs_client_(gcs_client) {} bool TaskDependencyManager::CheckObjectLocal(const ObjectID &object_id) const { return local_objects_.count(object_id) == 1; @@ -383,10 +382,11 @@ void TaskDependencyManager::AcquireTaskLease(const TaskID &task_id) { } auto task_lease_data = std::make_shared(); + task_lease_data->set_task_id(task_id.Binary()); task_lease_data->set_node_manager_id(client_id_.Hex()); task_lease_data->set_acquired_at(absl::GetCurrentTimeNanos() / 1000000); task_lease_data->set_timeout(it->second.lease_period); - RAY_CHECK_OK(task_lease_table_.Add(JobID::Nil(), task_id, task_lease_data, nullptr)); + RAY_CHECK_OK(gcs_client_->Tasks().AsyncAddTaskLease(task_lease_data, nullptr)); auto period = boost::posix_time::milliseconds(it->second.lease_period / 2); it->second.lease_timer->expires_from_now(period); diff --git a/src/ray/raylet/task_dependency_manager.h b/src/ray/raylet/task_dependency_manager.h index c714ac586..9403ba64e 100644 --- a/src/ray/raylet/task_dependency_manager.h +++ b/src/ray/raylet/task_dependency_manager.h @@ -4,6 +4,7 @@ // clang-format off #include "ray/common/id.h" #include "ray/common/task/task.h" +#include "ray/gcs/redis_gcs_client.h" #include "ray/object_manager/object_manager.h" #include "ray/raylet/reconstruction_policy.h" // clang-format on @@ -33,7 +34,7 @@ class TaskDependencyManager { ReconstructionPolicyInterface &reconstruction_policy, boost::asio::io_service &io_service, const ClientID &client_id, int64_t initial_lease_period_ms, - gcs::TableInterface &task_lease_table); + std::shared_ptr gcs_client); /// Check whether an object is locally available. /// @@ -226,8 +227,8 @@ class TaskDependencyManager { /// added to the GCS. The lease expiration period is doubled every time the /// lease is renewed. const int64_t initial_lease_period_ms_; - /// The storage system for the task lease table. - gcs::TableInterface &task_lease_table_; + /// A client connection to the GCS. + std::shared_ptr gcs_client_; /// A mapping from task ID of each subscribed task to its list of object /// dependencies, either task arguments or objects passed into `ray.get`. std::unordered_map task_dependencies_; diff --git a/src/ray/raylet/task_dependency_manager_test.cc b/src/ray/raylet/task_dependency_manager_test.cc index 5431f5d0a..d7cff6c4f 100644 --- a/src/ray/raylet/task_dependency_manager_test.cc +++ b/src/ray/raylet/task_dependency_manager_test.cc @@ -6,6 +6,8 @@ #include #include "ray/common/task/task_util.h" +#include "ray/gcs/redis_accessor.h" +#include "ray/gcs/redis_gcs_client.h" #include "ray/raylet/task_dependency_manager.h" #include "ray/util/test_util.h" @@ -31,13 +33,23 @@ class MockReconstructionPolicy : public ReconstructionPolicyInterface { MOCK_METHOD1(Cancel, void(const ObjectID &object_id)); }; -class MockGcs : public gcs::TableInterface { +class MockTaskInfoAccessor : public gcs::RedisTaskInfoAccessor { public: - MOCK_METHOD4( - Add, - ray::Status(const JobID &job_id, const TaskID &task_id, - const std::shared_ptr &task_data, - const gcs::TableInterface::WriteCallback &done)); + MockTaskInfoAccessor(gcs::RedisGcsClient *client) + : gcs::RedisTaskInfoAccessor(client) {} + + MOCK_METHOD2(AsyncAddTaskLease, + ray::Status(const std::shared_ptr &data_ptr, + const gcs::StatusCallback &callback)); +}; + +class MockGcsClient : public gcs::RedisGcsClient { + public: + MockGcsClient(const gcs::GcsClientOptions &options) : gcs::RedisGcsClient(options) {} + + void Init(MockTaskInfoAccessor *task_accessor_mock) { + task_accessor_.reset(task_accessor_mock); + } }; class TaskDependencyManagerTest : public ::testing::Test { @@ -46,11 +58,15 @@ class TaskDependencyManagerTest : public ::testing::Test { : object_manager_mock_(), reconstruction_policy_mock_(), io_service_(), - gcs_mock_(), + options_("", 1, ""), + gcs_client_mock_(new MockGcsClient(options_)), + task_accessor_mock_(new MockTaskInfoAccessor(gcs_client_mock_.get())), initial_lease_period_ms_(100), task_dependency_manager_(object_manager_mock_, reconstruction_policy_mock_, io_service_, ClientID::Nil(), initial_lease_period_ms_, - gcs_mock_) {} + gcs_client_mock_) { + gcs_client_mock_->Init(task_accessor_mock_); + } void Run(uint64_t timeout_ms) { auto timer_period = boost::posix_time::milliseconds(timeout_ms); @@ -67,7 +83,9 @@ class TaskDependencyManagerTest : public ::testing::Test { MockObjectManager object_manager_mock_; MockReconstructionPolicy reconstruction_policy_mock_; boost::asio::io_service io_service_; - MockGcs gcs_mock_; + gcs::GcsClientOptions options_; + std::shared_ptr gcs_client_mock_; + MockTaskInfoAccessor *task_accessor_mock_; int64_t initial_lease_period_ms_; TaskDependencyManager task_dependency_manager_; }; @@ -235,7 +253,7 @@ TEST_F(TaskDependencyManagerTest, TestTaskChain) { // Mark each task as pending. A lease entry should be added to the GCS for // each task. - EXPECT_CALL(gcs_mock_, Add(_, task.GetTaskSpecification().TaskId(), _, _)); + EXPECT_CALL(*task_accessor_mock_, AsyncAddTaskLease(_, _)); task_dependency_manager_.TaskPending(task); i++; @@ -287,7 +305,7 @@ TEST_F(TaskDependencyManagerTest, TestDependentPut) { // it is pending execution. EXPECT_CALL(object_manager_mock_, CancelPull(put_id)); EXPECT_CALL(reconstruction_policy_mock_, Cancel(put_id)); - EXPECT_CALL(gcs_mock_, Add(_, task1.GetTaskSpecification().TaskId(), _, _)); + EXPECT_CALL(*task_accessor_mock_, AsyncAddTaskLease(_, _)); task_dependency_manager_.TaskPending(task1); } @@ -300,7 +318,7 @@ TEST_F(TaskDependencyManagerTest, TestTaskForwarding) { const auto &arguments = task.GetDependencies(); static_cast(task_dependency_manager_.SubscribeGetDependencies( task.GetTaskSpecification().TaskId(), arguments)); - EXPECT_CALL(gcs_mock_, Add(_, task.GetTaskSpecification().TaskId(), _, _)); + EXPECT_CALL(*task_accessor_mock_, AsyncAddTaskLease(_, _)); task_dependency_manager_.TaskPending(task); } @@ -404,7 +422,8 @@ TEST_F(TaskDependencyManagerTest, TestTaskLeaseRenewal) { // Mark a task as pending. auto task = ExampleTask({}, 0); // We expect an initial call to acquire the lease. - EXPECT_CALL(gcs_mock_, Add(_, task.GetTaskSpecification().TaskId(), _, _)); + EXPECT_CALL(*task_accessor_mock_, AsyncAddTaskLease(_, _)); + task_dependency_manager_.TaskPending(task); // Check that while the task is still pending, there is one call to renew the @@ -415,8 +434,7 @@ TEST_F(TaskDependencyManagerTest, TestTaskLeaseRenewal) { for (int i = 1; i <= num_expected_calls; i++) { sleep_time += i * initial_lease_period_ms_; } - EXPECT_CALL(gcs_mock_, Add(_, task.GetTaskSpecification().TaskId(), _, _)) - .Times(num_expected_calls); + EXPECT_CALL(*task_accessor_mock_, AsyncAddTaskLease(_, _)).Times(num_expected_calls); Run(sleep_time); } @@ -438,7 +456,7 @@ TEST_F(TaskDependencyManagerTest, TestRemoveTasksAndRelatedObjects) { task.GetTaskSpecification().TaskId(), arguments); // Mark each task as pending. A lease entry should be added to the GCS for // each task. - EXPECT_CALL(gcs_mock_, Add(_, task.GetTaskSpecification().TaskId(), _, _)); + EXPECT_CALL(*task_accessor_mock_, AsyncAddTaskLease(_, _)); task_dependency_manager_.TaskPending(task); }