diff --git a/src/ray/raylet/lineage_cache.cc b/src/ray/raylet/lineage_cache.cc index 72aa7b045..917d925a7 100644 --- a/src/ray/raylet/lineage_cache.cc +++ b/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 LineageEntry::GetParentTaskIds() const { - std::unordered_set parent_ids; +const std::unordered_set &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; } diff --git a/src/ray/raylet/lineage_cache.h b/src/ray/raylet/lineage_cache.h index d5d72027f..14eba7a11 100644 --- a/src/ray/raylet/lineage_cache.h +++ b/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 GetParentTaskIds() const; + const std::unordered_set &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 parent_task_ids_; /// IDs of node managers that this task has been explicitly forwarded to. std::unordered_set forwarded_to_; }; diff --git a/src/ray/raylet/lineage_cache_test.cc b/src/ray/raylet/lineage_cache_test.cc index c1bfe2517..23a399a70 100644 --- a/src/ray/raylet/lineage_cache_test.cc +++ b/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(); 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) { diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index fb118a955..0972a6553 100644 --- a/src/ray/raylet/node_manager.cc +++ b/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); diff --git a/src/ray/raylet/scheduling_policy.cc b/src/ray/raylet/scheduling_policy.cc index 2cf739906..0b0dbf855 100644 --- a/src/ray/raylet/scheduling_policy.cc +++ b/src/ray/raylet/scheduling_policy.cc @@ -33,12 +33,12 @@ std::unordered_map 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; } diff --git a/src/ray/raylet/task.cc b/src/ray/raylet/task.cc index b075dfc03..880956ea3 100644 --- a/src/ray/raylet/task.cc +++ b/src/ray/raylet/task.cc @@ -11,55 +11,39 @@ flatbuffers::Offset 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 Task::GetDependencies() const { - std::vector dependencies; +void Task::SetExecutionDependencies(const std::vector &dependencies) { + task_execution_spec_.SetExecutionDependencies(dependencies); + ComputeDependencies(); +} + +void Task::IncrementNumForwards() { task_execution_spec_.IncrementNumForwards(); } + +const std::vector &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 diff --git a/src/ray/raylet/task.h b/src/ray/raylet/task.h index 472d23296..3663ff213 100644 --- a/src/ray/raylet/task.h +++ b/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 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 &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 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 &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 dependencies_; }; } // namespace raylet diff --git a/src/ray/raylet/task_dependency_manager_test.cc b/src/ray/raylet/task_dependency_manager_test.cc index e62ef1dca..5c79e1ec8 100644 --- a/src/ray/raylet/task_dependency_manager_test.cc +++ b/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(task_dependency_manager_.SubscribeDependencies( task.GetTaskSpecification().TaskId(), arguments)); EXPECT_CALL(gcs_mock_, Add(_, task.GetTaskSpecification().TaskId(), _, _));