[xray] Lineage cache requests notifications from the GCS about remote tasks (#1834)

* Add PubsubInterface to GCS tables

* Add task table PubsubInterface to lineage cache and tests

* Request notifications for remote tasks in the lineage cache

* Add RegisterGCS method to node manager

* Fix NodeManager member initialization order, subscribe to task table notifications

* Comments

* Use returned statuses.

* Fix double commit bug in lineage cache

* lint

* More linting.

* Fix pure virtual method declarations

src/ray/gcs/tables.cc

@@ -347,6 +347,7 @@ template class Table<TaskID, TaskTableData>;
 template class Log<ActorID, ActorTableData>;
 template class Log<TaskID, TaskReconstructionData>;
 template class Table<ClientID, HeartbeatTableData>;
+template class Log<UniqueID, ClientTableData>;
 
 }  // namespace gcs
 

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 <typename ID>
+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 <typename ID, typename Data>
-class Log {
+class Log : virtual public PubsubInterface<ID> {
  public:
   using DataT = typename Data::NativeTableType;
   using Callback = std::function<void(AsyncGcsClient *client, const ID &id,
@@ -183,7 +198,9 @@ class TableInterface {
 /// Example tables backed by Log:
 ///   TaskTable: Stores Task metadata needed for executing the task.
 template <typename ID, typename Data>
-class Table : private Log<ID, Data>, public TableInterface<ID, Data> {
+class Table : private Log<ID, Data>,
+              public TableInterface<ID, Data>,
+              virtual public PubsubInterface<ID> {
  public:
   using DataT = typename Log<ID, Data>::DataT;
   using Callback =

src/ray/raylet/lineage_cache.cc

@@ -123,8 +123,10 @@ flatbuffers::Offset<protocol::ForwardTaskRequest> Lineage::ToFlatbuffer(
   return request;
 }
 
-LineageCache::LineageCache(gcs::TableInterface<TaskID, protocol::Task> &task_storage)
-    : task_storage_(task_storage) {}
+LineageCache::LineageCache(const ClientID &client_id,
+                           gcs::TableInterface<TaskID, protocol::Task> &task_storage,
+                           gcs::PubsubInterface<TaskID> &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

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<TaskID, protocol::Task> &task_storage);
+  LineageCache(const ClientID &client_id,
+               gcs::TableInterface<TaskID, protocol::Task> &task_storage,
+               gcs::PubsubInterface<TaskID> &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<TaskID, protocol::Task> &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<TaskID> &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<TaskID, UniqueIDHasher> subscribed_tasks_;
 };
 
 }  // namespace raylet

src/ray/raylet/lineage_cache_test.cc

@@ -13,9 +13,15 @@ namespace ray {
 
 namespace raylet {
 
-class MockGcs : virtual public gcs::TableInterface<TaskID, protocol::Task> {
+class MockGcs : public gcs::TableInterface<TaskID, protocol::Task>,
+                public gcs::PubsubInterface<TaskID> {
  public:
-  MockGcs(){};
+  MockGcs() {}
+
+  void Subscribe(const gcs::raylet::TaskTable::WriteCallback &notification_callback) {
+    notification_callback_ = notification_callback;
+  }
+
   Status Add(const JobID &job_id, const TaskID &task_id,
              std::shared_ptr<protocol::TaskT> task_data,
              const gcs::TableInterface<TaskID, protocol::Task>::WriteCallback &done) {
@@ -23,29 +29,71 @@ class MockGcs : virtual public gcs::TableInterface<TaskID, protocol::Task> {
     callbacks_.push_back(
         std::pair<gcs::raylet::TaskTable::WriteCallback, TaskID>(done, task_id));
     return ray::Status::OK();
-  };
+  }
+
+  Status RemoteAdd(const TaskID &task_id, std::shared_ptr<protocol::TaskT> 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<protocol::TaskT> 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<TaskID, std::shared_ptr<protocol::TaskT>, UniqueIDHasher>
       &TaskTable() const {
     return task_table_;
   }
 
+  const std::unordered_set<TaskID, UniqueIDHasher> &SubscribedTasks() const {
+    return subscribed_tasks_;
+  }
+
  private:
   std::unordered_map<TaskID, std::shared_ptr<protocol::TaskT>, UniqueIDHasher>
       task_table_;
   std::vector<std::pair<gcs::raylet::TaskTable::WriteCallback, TaskID>> callbacks_;
+  gcs::raylet::TaskTable::WriteCallback notification_callback_;
+  std::unordered_set<TaskID, UniqueIDHasher> 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<ray::protocol::TaskT> 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<Task> tasks;
   auto return_values1 =
       InsertTaskChain(lineage_cache_, tasks, 3, std::vector<ObjectID>(), 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<protocol::ForwardTaskRequest>(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<Task> tasks;
+  auto return_values1 =
+      InsertTaskChain(lineage_cache_, tasks, 3, std::vector<ObjectID>(), 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<protocol::TaskT>();
+  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

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::AsyncGcsClient> 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,

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::AsyncGcsClient> 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::AsyncGcsClient> gcs_client_;
   std::vector<ClientID> remote_clients_;
   std::unordered_map<ClientID, TcpServerConnection, UniqueIDHasher>
       remote_server_connections_;
-  ObjectManager &object_manager_;
   std::unordered_map<ActorID, ActorRegistration, UniqueIDHasher> actor_registry_;
 };