diff --git a/src/ray/gcs/tables.cc b/src/ray/gcs/tables.cc index 81b1f0863..be73aedf8 100644 --- a/src/ray/gcs/tables.cc +++ b/src/ray/gcs/tables.cc @@ -347,6 +347,7 @@ template class Table; template class Log; template class Log; template class Table; +template class Log; } // namespace gcs diff --git a/src/ray/gcs/tables.h b/src/ray/gcs/tables.h index 0afa70853..e46ccbddf 100644 --- a/src/ray/gcs/tables.h +++ b/src/ray/gcs/tables.h @@ -27,6 +27,21 @@ class RedisContext; class AsyncGcsClient; +/// \class PubsubInterface +/// +/// The interface for a pubsub storage system. The client of a storage system +/// that implements this interface can request and cancel notifications for +/// specific keys. +template +class PubsubInterface { + public: + virtual Status RequestNotifications(const JobID &job_id, const ID &id, + const ClientID &client_id) = 0; + virtual Status CancelNotifications(const JobID &job_id, const ID &id, + const ClientID &client_id) = 0; + virtual ~PubsubInterface(){}; +}; + /// \class Log /// /// A GCS table where every entry is an append-only log. @@ -36,7 +51,7 @@ class AsyncGcsClient; /// ClientTable: Stores a log of which GCS clients have been added or deleted /// from the system. template -class Log { +class Log : virtual public PubsubInterface { public: using DataT = typename Data::NativeTableType; using Callback = std::function -class Table : private Log, public TableInterface { +class Table : private Log, + public TableInterface, + virtual public PubsubInterface { public: using DataT = typename Log::DataT; using Callback = diff --git a/src/ray/raylet/lineage_cache.cc b/src/ray/raylet/lineage_cache.cc index 0d55b0033..11aef6007 100644 --- a/src/ray/raylet/lineage_cache.cc +++ b/src/ray/raylet/lineage_cache.cc @@ -123,8 +123,10 @@ flatbuffers::Offset Lineage::ToFlatbuffer( return request; } -LineageCache::LineageCache(gcs::TableInterface &task_storage) - : task_storage_(task_storage) {} +LineageCache::LineageCache(const ClientID &client_id, + gcs::TableInterface &task_storage, + gcs::PubsubInterface &task_pubsub) + : client_id_(client_id), task_storage_(task_storage), task_pubsub_(task_pubsub) {} /// A helper function to merge one lineage into another, in DFS order. /// @@ -236,8 +238,15 @@ Status LineageCache::Flush() { // committed yet, then as far as we know, it's still in flight to the // GCS. Skip this task for now. if (parent && parent->GetStatus() != GcsStatus_COMMITTED) { - // TODO(swang): Once GCS notifications for the task table are ready, - // request notification for commit of the parent task here. + // Request notifications about the parent entry's commit in the GCS. + // Once we receive a notification about the task's commit via + // HandleEntryCommitted, then this task will be ready to write on the + // next call to Flush(). + auto inserted = subscribed_tasks_.insert(parent_id); + if (inserted.second) { + RAY_CHECK_OK( + task_pubsub_.RequestNotifications(JobID::nil(), parent_id, client_id_)); + } all_arguments_committed = false; break; } @@ -290,12 +299,28 @@ void PopAncestorTasks(const UniqueID &task_id, Lineage &lineage) { } void LineageCache::HandleEntryCommitted(const UniqueID &task_id) { + RAY_LOG(DEBUG) << "task committed: " << task_id; auto entry = lineage_.PopEntry(task_id); for (const auto &parent_id : entry->GetParentTaskIds()) { PopAncestorTasks(parent_id, lineage_); } - RAY_CHECK(entry->SetStatus(GcsStatus_COMMITTED)); + // Mark this task as COMMITTED. Any tasks that were dependent on it and are + // ready to be written may now be flushed to the GCS. + bool committed = entry->SetStatus(GcsStatus_COMMITTED); + if (!committed) { + // If we failed to mark the task as committed, check that it's because it + // was committed before. This means that we already received a notification + // about the commit. + RAY_CHECK(entry->GetStatus() == GcsStatus_COMMITTED); + } RAY_CHECK(lineage_.SetEntry(std::move(*entry))); + + // Stop listening for notifications about this task. + auto it = subscribed_tasks_.find(task_id); + if (it != subscribed_tasks_.end()) { + RAY_CHECK_OK(task_pubsub_.CancelNotifications(JobID::nil(), task_id, client_id_)); + subscribed_tasks_.erase(it); + } } } // namespace raylet diff --git a/src/ray/raylet/lineage_cache.h b/src/ray/raylet/lineage_cache.h index e9904277f..3b40f8d88 100644 --- a/src/ray/raylet/lineage_cache.h +++ b/src/ray/raylet/lineage_cache.h @@ -164,7 +164,9 @@ class LineageCache { public: /// Create a lineage cache for the given task storage system. /// TODO(swang): Pass in the policy (interface?). - LineageCache(gcs::TableInterface &task_storage); + LineageCache(const ClientID &client_id, + gcs::TableInterface &task_storage, + gcs::PubsubInterface &task_pubsub); /// Add a task that is waiting for execution and its uncommitted lineage. /// These entries will not be written to the GCS until set to ready. @@ -181,7 +183,11 @@ class LineageCache { /// \param task The task to set as ready. void AddReadyTask(const Task &task); - void RemoveWaitingTask(const TaskID &entry_id); + /// Remove a task that was waiting for execution. Its uncommitted lineage + /// will remain unchanged. + /// + /// \param task_id The ID of the waiting task to remove. + void RemoveWaitingTask(const TaskID &task_id); /// Get the uncommitted lineage of a task. The uncommitted lineage consists /// of all tasks in the given task's lineage that have not been committed in @@ -199,11 +205,21 @@ class LineageCache { /// \return Status. Status Flush(); - private: - void HandleEntryCommitted(const TaskID &unique_id); + /// Handle the commit of a task entry in the GCS. This sets the task to + /// COMMITTED and cleans up any ancestor tasks that are in the cache. + /// + /// \param task_id The ID of the task entry that was committed. + void HandleEntryCommitted(const TaskID &task_id); + private: + /// The client ID, used to request notifications for specific tasks. + /// TODO(swang): Move the ClientID into the generic Table implementation. + ClientID client_id_; /// The durable storage system for task information. gcs::TableInterface &task_storage_; + /// The pubsub storage system for task information. This can be used to + /// request notifications for the commit of a task entry. + gcs::PubsubInterface &task_pubsub_; /// The set of tasks that are in UNCOMMITTED_READY state. This is a cache of /// the tasks that may be flushable. // TODO(swang): As an optimization, we may also want to further distinguish @@ -214,6 +230,9 @@ class LineageCache { /// All tasks and objects that we are responsible for writing back to the /// GCS, and the tasks and objects in their lineage. Lineage lineage_; + /// The tasks that we've subscribed to notifications for from the pubsub + /// storage system. We will receive a notification for these tasks on commit. + std::unordered_set subscribed_tasks_; }; } // namespace raylet diff --git a/src/ray/raylet/lineage_cache_test.cc b/src/ray/raylet/lineage_cache_test.cc index 0610997ff..de89a5b57 100644 --- a/src/ray/raylet/lineage_cache_test.cc +++ b/src/ray/raylet/lineage_cache_test.cc @@ -13,9 +13,15 @@ namespace ray { namespace raylet { -class MockGcs : virtual public gcs::TableInterface { +class MockGcs : public gcs::TableInterface, + public gcs::PubsubInterface { public: - MockGcs(){}; + MockGcs() {} + + void Subscribe(const gcs::raylet::TaskTable::WriteCallback ¬ification_callback) { + notification_callback_ = notification_callback; + } + Status Add(const JobID &job_id, const TaskID &task_id, std::shared_ptr task_data, const gcs::TableInterface::WriteCallback &done) { @@ -23,29 +29,71 @@ class MockGcs : virtual public gcs::TableInterface { callbacks_.push_back( std::pair(done, task_id)); return ray::Status::OK(); - }; + } + + Status RemoteAdd(const TaskID &task_id, std::shared_ptr task_data) { + task_table_[task_id] = task_data; + // Send a notification after the add if the lineage cache requested + // notifications for this key. + bool send_notification = (subscribed_tasks_.count(task_id) == 1); + auto callback = [this, send_notification](ray::gcs::AsyncGcsClient *client, + const TaskID &task_id, + std::shared_ptr data) { + if (send_notification) { + notification_callback_(client, task_id, data); + } + }; + return Add(JobID::nil(), task_id, task_data, callback); + } + + Status RequestNotifications(const JobID &job_id, const TaskID &task_id, + const ClientID &client_id) { + subscribed_tasks_.insert(task_id); + if (task_table_.count(task_id) == 1) { + callbacks_.push_back({notification_callback_, task_id}); + } + return ray::Status::OK(); + } + + Status CancelNotifications(const JobID &job_id, const TaskID &task_id, + const ClientID &client_id) { + subscribed_tasks_.erase(task_id); + return ray::Status::OK(); + } void Flush() { for (const auto &callback : callbacks_) { callback.first(NULL, callback.second, task_table_[callback.second]); } callbacks_.clear(); - }; + } const std::unordered_map, UniqueIDHasher> &TaskTable() const { return task_table_; } + const std::unordered_set &SubscribedTasks() const { + return subscribed_tasks_; + } + private: std::unordered_map, UniqueIDHasher> task_table_; std::vector> callbacks_; + gcs::raylet::TaskTable::WriteCallback notification_callback_; + std::unordered_set subscribed_tasks_; }; class LineageCacheTest : public ::testing::Test { public: - LineageCacheTest() : mock_gcs_(), lineage_cache_(mock_gcs_) {} + LineageCacheTest() + : mock_gcs_(), lineage_cache_(ClientID::from_random(), mock_gcs_, mock_gcs_) { + mock_gcs_.Subscribe([this](ray::gcs::AsyncGcsClient *client, const TaskID &task_id, + std::shared_ptr data) { + lineage_cache_.HandleEntryCommitted(task_id); + }); + } protected: MockGcs mock_gcs_; @@ -234,30 +282,68 @@ TEST_F(LineageCacheTest, TestWritebackPartiallyReady) { CheckFlush(lineage_cache_, mock_gcs_, num_tasks_flushed); } -TEST_F(LineageCacheTest, TestRemoveWaitingTask) { +TEST_F(LineageCacheTest, TestForwardTaskRoundTrip) { // Insert a chain of dependent tasks. std::vector tasks; auto return_values1 = InsertTaskChain(lineage_cache_, tasks, 3, std::vector(), 1); - auto task_to_remove = tasks[1]; - auto task_id_to_remove = task_to_remove.GetTaskSpecification().TaskId(); + // Simulate removing the task and forwarding it to another node. + auto forwarded_task = tasks[1]; + auto task_id_to_remove = forwarded_task.GetTaskSpecification().TaskId(); auto uncommitted_lineage = lineage_cache_.GetUncommittedLineage(task_id_to_remove); + lineage_cache_.RemoveWaitingTask(task_id_to_remove); + + // Simulate receiving the task again. flatbuffers::FlatBufferBuilder fbb; auto uncommitted_lineage_message = uncommitted_lineage.ToFlatbuffer(fbb, task_id_to_remove); fbb.Finish(uncommitted_lineage_message); uncommitted_lineage = Lineage( *flatbuffers::GetRoot(fbb.GetBufferPointer())); + lineage_cache_.AddWaitingTask(forwarded_task, uncommitted_lineage); +} - const Task &task = uncommitted_lineage.GetEntry(task_id_to_remove)->TaskData(); - RAY_LOG(INFO) << "removing task " << task.GetTaskSpecification().TaskId() - << "with numforwards=" - << task.GetTaskExecutionSpecReadonly().NumForwards(); - ASSERT_EQ(task.GetTaskExecutionSpecReadonly().NumForwards(), 1); +TEST_F(LineageCacheTest, TestForwardTask) { + // Insert a chain of dependent tasks. + size_t num_tasks_flushed = 0; + std::vector tasks; + auto return_values1 = + InsertTaskChain(lineage_cache_, tasks, 3, std::vector(), 1); + // Simulate removing the task and forwarding it to another node. + auto it = tasks.begin() + 1; + auto forwarded_task = *it; + tasks.erase(it); + auto task_id_to_remove = forwarded_task.GetTaskSpecification().TaskId(); + auto uncommitted_lineage = lineage_cache_.GetUncommittedLineage(task_id_to_remove); lineage_cache_.RemoveWaitingTask(task_id_to_remove); - lineage_cache_.AddWaitingTask(task_to_remove, uncommitted_lineage); + + // Simulate executing the remaining tasks. + for (const auto &task : tasks) { + lineage_cache_.AddReadyTask(task); + } + // Check that the first task, which has no dependencies can be flushed. The + // last task cannot be flushed since one of its dependencies has not been + // added by the remote node yet. + num_tasks_flushed++; + CheckFlush(lineage_cache_, mock_gcs_, num_tasks_flushed); + + // Simulate executing the task on a remote node and adding it to the GCS. + auto task_data = std::make_shared(); + RAY_CHECK_OK( + mock_gcs_.RemoteAdd(forwarded_task.GetTaskSpecification().TaskId(), task_data)); + // Check that the remote task is flushed. + num_tasks_flushed++; + CheckFlush(lineage_cache_, mock_gcs_, num_tasks_flushed); + ASSERT_EQ(mock_gcs_.SubscribedTasks().size(), 1); + + // Check that once we receive the callback for the remote task, we can now + // flush the last task. + mock_gcs_.Flush(); + num_tasks_flushed++; + CheckFlush(lineage_cache_, mock_gcs_, num_tasks_flushed); + ASSERT_EQ(mock_gcs_.SubscribedTasks().size(), 0); } } // namespace raylet diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index 08d7d0983..58ded7caf 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -42,6 +42,8 @@ NodeManager::NodeManager(boost::asio::io_service &io_service, const NodeManagerConfig &config, ObjectManager &object_manager, std::shared_ptr gcs_client) : io_service_(io_service), + object_manager_(object_manager), + gcs_client_(gcs_client), heartbeat_timer_(io_service), heartbeat_period_ms_(config.heartbeat_period_ms), local_resources_(config.resource_config), @@ -53,11 +55,10 @@ NodeManager::NodeManager(boost::asio::io_service &io_service, object_manager, // reconstruction_policy_, [this](const TaskID &task_id) { HandleWaitingTaskReady(task_id); }), - lineage_cache_(gcs_client->raylet_task_table()), - gcs_client_(gcs_client), + lineage_cache_(gcs_client_->client_table().GetLocalClientId(), + gcs_client->raylet_task_table(), gcs_client->raylet_task_table()), remote_clients_(), remote_server_connections_(), - object_manager_(object_manager), actor_registry_() { RAY_CHECK(heartbeat_period_ms_ > 0); // Initialize the resource map with own cluster resource configuration. @@ -67,6 +68,18 @@ NodeManager::NodeManager(boost::asio::io_service &io_service, } ray::Status NodeManager::RegisterGcs() { + // Subscribe to task entry commits in the GCS. These notifications are + // forwarded to the lineage cache, which requests notifications about tasks + // that were executed remotely. + const auto task_committed_callback = [this](gcs::AsyncGcsClient *client, + const TaskID &task_id, + const ray::protocol::TaskT &task_data) { + lineage_cache_.HandleEntryCommitted(task_id); + }; + RAY_RETURN_NOT_OK(gcs_client_->raylet_task_table().Subscribe( + JobID::nil(), gcs_client_->client_table().GetLocalClientId(), + task_committed_callback, nullptr)); + // Register a callback for actor creation notifications. auto actor_creation_callback = [this]( gcs::AsyncGcsClient *client, const ActorID &actor_id, diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h index bfee64c6b..de9d65a0b 100644 --- a/src/ray/raylet/node_manager.h +++ b/src/ray/raylet/node_manager.h @@ -56,9 +56,6 @@ class NodeManager { ray::Status RegisterGcs(); - void HeartbeatAdded(gcs::AsyncGcsClient *client, const ClientID &id, - const HeartbeatTableDataT &data); - private: // Handler for the addition of a new GCS client. void ClientAdded(const ClientTableDataT &data); @@ -84,8 +81,17 @@ class NodeManager { ray::Status ForwardTask(const Task &task, const ClientID &node_id); /// Send heartbeats to the GCS. void Heartbeat(); + /// Handler for a notification about a new client from the GCS. + void ClientAdded(gcs::AsyncGcsClient *client, const UniqueID &id, + const ClientTableDataT &data); + /// Handler for a heartbeat notification from the GCS. + void HeartbeatAdded(gcs::AsyncGcsClient *client, const ClientID &id, + const HeartbeatTableDataT &data); boost::asio::io_service &io_service_; + ObjectManager &object_manager_; + /// A client connection to the GCS. + std::shared_ptr gcs_client_; boost::asio::deadline_timer heartbeat_timer_; uint64_t heartbeat_period_ms_; /// The resources local to this node. @@ -104,12 +110,9 @@ class NodeManager { TaskDependencyManager task_dependency_manager_; /// The lineage cache for the GCS object and task tables. LineageCache lineage_cache_; - /// A client connection to the GCS. - std::shared_ptr gcs_client_; std::vector remote_clients_; std::unordered_map remote_server_connections_; - ObjectManager &object_manager_; std::unordered_map actor_registry_; };