[xray] Cache a task's object dependencies (#2623)

* Cache a Task's object dependencies

* Cache the parent task IDs for lineage cache entries

* Cache the parent task IDs in lineage cache entries

* revert

* Fix test

* remove unused line

* Fix test

src/ray/raylet/lineage_cache.cc

@@ -5,7 +5,9 @@ namespace ray {
 namespace raylet {
 
 LineageEntry::LineageEntry(const Task &task, GcsStatus status)
-    : status_(status), task_(task) {}
+    : status_(status), task_(task) {
+  ComputeParentTaskIds();
+}
 
 GcsStatus LineageEntry::GetStatus() const { return status_; }
 
@@ -35,20 +37,27 @@ const TaskID LineageEntry::GetEntryId() const {
   return task_.GetTaskSpecification().TaskId();
 }
 
-const std::unordered_set<TaskID> LineageEntry::GetParentTaskIds() const {
-  std::unordered_set<TaskID> parent_ids;
+const std::unordered_set<TaskID> &LineageEntry::GetParentTaskIds() const {
+  return parent_task_ids_;
+}
+
+void LineageEntry::ComputeParentTaskIds() {
+  parent_task_ids_.clear();
   // A task's parents are the tasks that created its arguments.
-  auto dependencies = task_.GetDependencies();
-  for (auto &dependency : dependencies) {
-    parent_ids.insert(ComputeTaskId(dependency));
+  for (const auto &dependency : task_.GetDependencies()) {
+    parent_task_ids_.insert(ComputeTaskId(dependency));
   }
-  return parent_ids;
 }
 
 const Task &LineageEntry::TaskData() const { return task_; }
 
 Task &LineageEntry::TaskDataMutable() { return task_; }
 
+void LineageEntry::UpdateTaskData(const Task &task) {
+  task_.CopyTaskExecutionSpec(task);
+  ComputeParentTaskIds();
+}
+
 Lineage::Lineage() {}
 
 Lineage::Lineage(const protocol::ForwardTaskRequest &task_request) {
@@ -86,7 +95,7 @@ bool Lineage::SetEntry(const Task &task, GcsStatus status) {
     if (current_entry->SetStatus(status)) {
       // SetStatus() would check if the new status is greater,
       // if it succeeds, go ahead to update the task field.
-      current_entry->TaskDataMutable().CopyTaskExecutionSpec(task);
+      current_entry->UpdateTaskData(task);
       return true;
     }
     return false;
@@ -160,7 +169,7 @@ void MergeLineageHelper(const TaskID &task_id, const Lineage &lineage_from,
   }
 
   // Insert a copy of the entry into lineage_to.
-  auto parent_ids = entry->GetParentTaskIds();
+  const auto &parent_ids = entry->GetParentTaskIds();
   // If the insert is successful, then continue the DFS. The insert will fail
   // if the new entry has an equal or lower GCS status than the current entry
   // in lineage_to. This also prevents us from traversing the same node twice.
@@ -312,9 +321,16 @@ Lineage LineageCache::GetUncommittedLineage(const TaskID &task_id,
       task_id, lineage_, uncommitted_lineage, [&](const LineageEntry &entry) {
         // The stopping condition for recursion is that the entry has
         // been committed to the GCS or has already been forwarded.
-        // The lineage always includes the requested task id.
-        return entry.WasExplicitlyForwarded(node_id) && !(entry.GetEntryId() == task_id);
+        return entry.WasExplicitlyForwarded(node_id);
       });
+  // The lineage always includes the requested task id, so add the task if it
+  // wasn't already added. The requested task may not have been added if it was
+  // already explicitly forwarded to this node before.
+  if (uncommitted_lineage.GetEntries().empty()) {
+    auto entry = lineage_.GetEntry(task_id);
+    RAY_CHECK(entry);
+    RAY_CHECK(uncommitted_lineage.SetEntry(entry->TaskData(), entry->GetStatus()));
+  }
   return uncommitted_lineage;
 }
 

src/ray/raylet/lineage_cache.h

@@ -84,22 +84,38 @@ class LineageEntry {
   /// that created its arguments.
   ///
   /// \return The IDs of the parent entries.
-  const std::unordered_set<TaskID> GetParentTaskIds() const;
+  const std::unordered_set<TaskID> &GetParentTaskIds() const;
 
   /// Get the task data.
   ///
   /// \return The task data.
   const Task &TaskData() const;
 
+  /// Get a mutable version of the task data.
+  ///
+  /// \return The task data.
+  /// TODO(swang): This is pretty ugly.
   Task &TaskDataMutable();
 
+  /// Update the task data with a new task.
+  ///
+  /// \return Void.
+  void UpdateTaskData(const Task &task);
+
+ private:
+  /// Compute cached parent task IDs. This task is dependent on values returned
+  /// by these tasks.
+  void ComputeParentTaskIds();
+
   /// The current state of this entry according to its status in the GCS.
   GcsStatus status_;
   /// The task data to be written to the GCS. This is nullptr if the entry is
   /// an object.
   //  const Task task_;
   Task task_;
-
+  /// A cached copy of the parent task IDs. This task is dependent on values
+  /// returned by these tasks.
+  std::unordered_set<TaskID> parent_task_ids_;
   /// IDs of node managers that this task has been explicitly forwarded to.
   std::unordered_set<ClientID> forwarded_to_;
 };

src/ray/raylet/lineage_cache_test.cc

@@ -217,7 +217,7 @@ TEST_F(LineageCacheTest, TestMarkTaskAsForwarded) {
   // Check that lineage of requested task includes itself, regardless of whether
   // it has been forwarded before.
   auto uncommitted_lineage_forwarded =
-      lineage_cache_.GetUncommittedLineage(remaining_task_id, node_id);
+      lineage_cache_.GetUncommittedLineage(forwarded_task_id, node_id);
   ASSERT_EQ(1, uncommitted_lineage_forwarded.GetEntries().size());
 }
 
@@ -284,7 +284,6 @@ TEST_F(LineageCacheTest, TestWritebackPartiallyReady) {
     returns.push_back(task2.GetTaskSpecification().ReturnId(i));
   }
   auto dependent_task = ExampleTask(returns, 1);
-  auto dependencies = dependent_task.GetDependencies();
 
   // Insert all tasks as waiting for execution.
   ASSERT_TRUE(lineage_cache_.AddWaitingTask(task1, Lineage()));
@@ -413,8 +412,9 @@ TEST_F(LineageCacheTest, TestEviction) {
   // uncommitted lineage.
   ASSERT_EQ(uncommitted_lineage.GetEntries().size(), lineage_size);
 
-  // Simulate executing the rest of the tasks on a remote node and adding them
-  // to the GCS.
+  // Simulate executing all the rest of the tasks except the last one on a
+  // remote node and adding them to the GCS.
+  tasks.pop_back();
   for (; it != tasks.end(); it++) {
     RAY_CHECK_OK(mock_gcs_.RemoteAdd(it->GetTaskSpecification().TaskId(), task_data));
     // Check that the remote task is flushed.
@@ -426,7 +426,9 @@ TEST_F(LineageCacheTest, TestEviction) {
   // evicted that the uncommitted lineage is now less than the maximum size.
   uncommitted_lineage =
       lineage_cache_.GetUncommittedLineage(last_task_id, ClientID::nil());
-  ASSERT_TRUE(uncommitted_lineage.GetEntries().size() <= max_lineage_size_);
+  ASSERT_TRUE(uncommitted_lineage.GetEntries().size() < max_lineage_size_);
+  // The remaining task should have no uncommitted lineage.
+  ASSERT_EQ(uncommitted_lineage.GetEntries().size(), 1);
 }
 
 TEST_F(LineageCacheTest, TestOutOfOrderEviction) {
@@ -453,8 +455,10 @@ TEST_F(LineageCacheTest, TestOutOfOrderEviction) {
       lineage_cache_.GetUncommittedLineage(last_task_id, ClientID::nil());
   ASSERT_EQ(uncommitted_lineage.GetEntries().size(), lineage_size);
 
-  // Simulate executing the tasks at the remote node and receiving the
-  // notifications from the GCS in reverse order of execution.
+  // Simulate executing all the rest of the tasks except the last one at the
+  // remote node. Simulate receiving the notifications from the GCS in reverse
+  // order of execution.
+  tasks.pop_back();
   auto task_data = std::make_shared<protocol::TaskT>();
   auto it = tasks.rbegin();
   RAY_CHECK_OK(mock_gcs_.RemoteAdd(it->GetTaskSpecification().TaskId(), task_data));
@@ -480,7 +484,7 @@ TEST_F(LineageCacheTest, TestOutOfOrderEviction) {
   // evicted that the uncommitted lineage is now less than the maximum size.
   uncommitted_lineage =
       lineage_cache_.GetUncommittedLineage(last_task_id, ClientID::nil());
-  ASSERT_TRUE(uncommitted_lineage.GetEntries().size() <= max_lineage_size_);
+  ASSERT_TRUE(uncommitted_lineage.GetEntries().size() < max_lineage_size_);
 }
 
 TEST_F(LineageCacheTest, TestEvictionUncommittedChildren) {

src/ray/raylet/node_manager.cc

@@ -736,7 +736,7 @@ void NodeManager::ProcessNodeManagerMessage(TcpClientConnection &node_manager_cl
     Lineage uncommitted_lineage(*message);
     const Task &task = uncommitted_lineage.GetEntry(task_id)->TaskData();
     RAY_LOG(DEBUG) << "got task " << task.GetTaskSpecification().TaskId()
-                   << " spillback=" << task.GetTaskExecutionSpecReadonly().NumForwards();
+                   << " spillback=" << task.GetTaskExecutionSpec().NumForwards();
     SubmitTask(task, uncommitted_lineage, /* forwarded = */ true);
   } break;
   case protocol::MessageType::DisconnectClient: {
@@ -1085,8 +1085,7 @@ void NodeManager::AssignTask(Task &task) {
       // we may lose updates that are in flight to the task table. We only
       // guarantee deterministic reconstruction ordering for tasks whose
       // updates are reflected in the task table.
-      TaskExecutionSpecification &mutable_spec = task.GetTaskExecutionSpec();
-      mutable_spec.SetExecutionDependencies({execution_dependency});
+      task.SetExecutionDependencies({execution_dependency});
       // Extend the frontier to include the executing task.
       actor_entry->second.ExtendFrontier(spec.ActorHandleId(), spec.ActorDummyObject());
     }
@@ -1323,7 +1322,7 @@ ray::Status NodeManager::ForwardTask(const Task &task, const ClientID &node_id)
       uncommitted_lineage.GetEntryMutable(task_id)->TaskDataMutable();
 
   // Increment forward count for the forwarded task.
-  lineage_cache_entry_task.GetTaskExecutionSpec().IncrementNumForwards();
+  lineage_cache_entry_task.IncrementNumForwards();
 
   flatbuffers::FlatBufferBuilder fbb;
   auto request = uncommitted_lineage.ToFlatbuffer(fbb, task_id);

src/ray/raylet/scheduling_policy.cc

@@ -33,12 +33,12 @@ std::unordered_map<TaskID, ClientID> SchedulingPolicy::Schedule(
     const auto &resource_demand = t.GetTaskSpecification().GetRequiredResources();
     const TaskID &task_id = t.GetTaskSpecification().TaskId();
     RAY_LOG(DEBUG) << "[SchedulingPolicy]: task=" << task_id
-                   << " numforwards=" << t.GetTaskExecutionSpecReadonly().NumForwards()
+                   << " numforwards=" << t.GetTaskExecutionSpec().NumForwards()
                    << " resources="
                    << t.GetTaskSpecification().GetRequiredResources().ToString();
     // TODO(atumanov): replace the simple spillback policy with exponential backoff based
     // policy.
-    if (t.GetTaskExecutionSpecReadonly().NumForwards() >= 1) {
+    if (t.GetTaskExecutionSpec().NumForwards() >= 1) {
       decision[task_id] = local_client_id;
       continue;
     }

src/ray/raylet/task.cc

@@ -11,55 +11,39 @@ flatbuffers::Offset<protocol::Task> Task::ToFlatbuffer(
   return task;
 }
 
-TaskExecutionSpecification &Task::GetTaskExecutionSpec() { return task_execution_spec_; }
-
-const TaskExecutionSpecification &Task::GetTaskExecutionSpecReadonly() const {
+const TaskExecutionSpecification &Task::GetTaskExecutionSpec() const {
   return task_execution_spec_;
 }
 
 const TaskSpecification &Task::GetTaskSpecification() const { return task_spec_; }
 
-const std::vector<ObjectID> Task::GetDependencies() const {
-  std::vector<ObjectID> dependencies;
+void Task::SetExecutionDependencies(const std::vector<ObjectID> &dependencies) {
+  task_execution_spec_.SetExecutionDependencies(dependencies);
+  ComputeDependencies();
+}
+
+void Task::IncrementNumForwards() { task_execution_spec_.IncrementNumForwards(); }
+
+const std::vector<ObjectID> &Task::GetDependencies() const { return dependencies_; }
+
+void Task::ComputeDependencies() {
+  dependencies_.clear();
   for (int i = 0; i < task_spec_.NumArgs(); ++i) {
     int count = task_spec_.ArgIdCount(i);
     for (int j = 0; j < count; j++) {
-      dependencies.push_back(task_spec_.ArgId(i, j));
+      dependencies_.push_back(task_spec_.ArgId(i, j));
     }
   }
   // TODO(atumanov): why not just return a const reference to ExecutionDependencies() and
   // avoid a copy.
   auto execution_dependencies = task_execution_spec_.ExecutionDependencies();
-  dependencies.insert(dependencies.end(), execution_dependencies.begin(),
-                      execution_dependencies.end());
-  return dependencies;
-}
-
-bool Task::DependsOn(const ObjectID &object_id) const {
-  // Iterate through the task arguments to see if it contains object_id.
-  int64_t num_args = task_spec_.NumArgs();
-  for (int i = 0; i < num_args; ++i) {
-    int count = task_spec_.ArgIdCount(i);
-    for (int j = 0; j < count; j++) {
-      ObjectID arg_id = task_spec_.ArgId(i, j);
-      if (arg_id == object_id) {
-        return true;
-      }
-    }
-  }
-  // Iterate through the execution dependencies to see if it contains object_id.
-  for (const auto &dependency_id : task_execution_spec_.ExecutionDependencies()) {
-    if (dependency_id == object_id) {
-      return true;
-    }
-  }
-  // The requested object ID was not a task argument or an execution dependency.
-  // This task is not dependent on it.
-  return false;
+  dependencies_.insert(dependencies_.end(), execution_dependencies.begin(),
+                       execution_dependencies.end());
 }
 
 void Task::CopyTaskExecutionSpec(const Task &task) {
-  task_execution_spec_ = task.GetTaskExecutionSpecReadonly();
+  task_execution_spec_ = task.GetTaskExecutionSpec();
+  ComputeDependencies();
 }
 
 }  // namespace raylet

src/ray/raylet/task.h

@@ -28,21 +28,22 @@ class Task {
   /// are determined at task submission time.
   Task(const TaskExecutionSpecification &execution_spec,
        const TaskSpecification &task_spec)
-      : task_execution_spec_(execution_spec), task_spec_(task_spec) {}
+      : task_execution_spec_(execution_spec), task_spec_(task_spec) {
+    ComputeDependencies();
+  }
 
   /// Create a task from a serialized flatbuffer.
   ///
   /// \param task_flatbuffer The serialized task.
   Task(const protocol::Task &task_flatbuffer)
-      : task_execution_spec_(*task_flatbuffer.task_execution_spec()),
-        task_spec_(*task_flatbuffer.task_specification()) {}
+      : Task(*task_flatbuffer.task_execution_spec(),
+             *task_flatbuffer.task_specification()) {}
 
   /// Create a task from a flatbuffer object.
   ///
   /// \param task_data The task flatbuffer object.
   Task(const protocol::TaskT &task_data)
-      : task_execution_spec_(*task_data.task_execution_spec),
-        task_spec_(task_data.task_specification) {}
+      : Task(*task_data.task_execution_spec, task_data.task_specification) {}
 
   /// Destroy the task.
   virtual ~Task() {}
@@ -54,37 +55,38 @@ class Task {
   flatbuffers::Offset<protocol::Task> ToFlatbuffer(
       flatbuffers::FlatBufferBuilder &fbb) const;
 
-  /// Get the execution specification for the task.
+  /// Get the mutable specification for the task. This specification may be
+  /// updated at runtime.
   ///
   /// \return The mutable specification for the task.
-  TaskExecutionSpecification &GetTaskExecutionSpec();
-
-  const TaskExecutionSpecification &GetTaskExecutionSpecReadonly() const;
+  const TaskExecutionSpecification &GetTaskExecutionSpec() const;
 
   /// Get the immutable specification for the task.
   ///
   /// \return The immutable specification for the task.
   const TaskSpecification &GetTaskSpecification() const;
 
+  /// Set the task's execution dependencies.
+  ///
+  /// \param dependencies The value to set the execution dependencies to.
+  void SetExecutionDependencies(const std::vector<ObjectID> &dependencies);
+
+  /// Increment the number of times this task has been forwarded.
+  void IncrementNumForwards();
+
   /// Get the task's object dependencies. This comprises the immutable task
   /// arguments and the mutable execution dependencies.
   ///
   /// \return The object dependencies.
-  /// TODO(atumanov): consider returning a constant reference.
-  const std::vector<ObjectID> GetDependencies() const;
-
-  /// Compute whether the task is dependent on an object ID.
-  ///
-  /// \param object_id The object ID that the task may be dependent on.
-  /// \return Returns true if the task is dependent on the given object ID and
-  /// false otherwise.
-  bool DependsOn(const ObjectID &object_id) const;
+  const std::vector<ObjectID> &GetDependencies() const;
 
   /// Update the dynamic/mutable information for this task.
   /// \param task Task structure with updated dynamic information.
   void CopyTaskExecutionSpec(const Task &task);
 
  private:
+  void ComputeDependencies();
+
   /// Task execution specification, consisting of all dynamic/mutable
   /// information about this task determined at execution time..
   TaskExecutionSpecification task_execution_spec_;
@@ -92,6 +94,10 @@ class Task {
   /// task determined at submission time. Includes resource demand, object
   /// dependencies, etc.
   TaskSpecification task_spec_;
+  /// A cached copy of the task's object dependencies, including arguments from
+  /// the TaskSpecification and execution dependencies from the
+  /// TaskExecutionSpecification.
+  std::vector<ObjectID> dependencies_;
 };
 
 }  // namespace raylet

src/ray/raylet/task_dependency_manager_test.cc

@@ -217,7 +217,7 @@ TEST_F(TaskDependencyManagerTest, TestTaskChain) {
   EXPECT_CALL(reconstruction_policy_mock_, Cancel(_)).Times(0);
   for (const auto &task : tasks) {
     // Subscribe to each of the tasks' arguments.
-    auto arguments = task.GetDependencies();
+    const auto &arguments = task.GetDependencies();
     bool ready = task_dependency_manager_.SubscribeDependencies(
         task.GetTaskSpecification().TaskId(), arguments);
     if (i < num_ready_tasks) {
@@ -291,7 +291,7 @@ TEST_F(TaskDependencyManagerTest, TestTaskForwarding) {
   auto tasks = MakeTaskChain(num_tasks, {}, 1);
   for (const auto &task : tasks) {
     // Subscribe to each of the tasks' arguments.
-    auto arguments = task.GetDependencies();
+    const auto &arguments = task.GetDependencies();
     static_cast<void>(task_dependency_manager_.SubscribeDependencies(
         task.GetTaskSpecification().TaskId(), arguments));
     EXPECT_CALL(gcs_mock_, Add(_, task.GetTaskSpecification().TaskId(), _, _));