diff --git a/src/ray/raylet/lineage_cache.cc b/src/ray/raylet/lineage_cache.cc index 910c3481b..795f2b54a 100644 --- a/src/ray/raylet/lineage_cache.cc +++ b/src/ray/raylet/lineage_cache.cc @@ -68,7 +68,7 @@ Lineage::Lineage(const protocol::ForwardTaskRequest &task_request) { auto tasks = task_request.uncommitted_tasks(); for (auto it = tasks->begin(); it != tasks->end(); it++) { const auto &task = **it; - RAY_CHECK(SetEntry(task, GcsStatus::UNCOMMITTED_REMOTE)); + RAY_CHECK(SetEntry(task, GcsStatus::UNCOMMITTED)); } } @@ -108,38 +108,23 @@ bool Lineage::SetEntry(const Task &task, GcsStatus status) { auto task_id = task.GetTaskSpecification().TaskId(); auto it = entries_.find(task_id); bool updated = false; - std::unordered_set old_parents; if (it != entries_.end()) { if (it->second.SetStatus(status)) { - // The task's spec may have changed, so record its old dependencies. - old_parents = it->second.GetParentTaskIds(); - // SetStatus() would check if the new status is greater, - // if it succeeds, go ahead to update the task field. - it->second.UpdateTaskData(task); + // We assume here that the new `task` has the same fields as the task + // already in the lineage cache. If this is not true, then it is + // necessary to update the task data of the existing lineage cache entry + // with LineageEntry::UpdateTaskData. updated = true; } } else { LineageEntry new_entry(task, status); it = entries_.emplace(std::make_pair(task_id, std::move(new_entry))).first; updated = true; - } - // If the task data was updated, then record which tasks it depends on. Add - // all new tasks that it depends on and remove any old tasks that it no - // longer depends on. - // TODO(swang): Updating the task data every time could be inefficient for - // tasks that have lots of dependencies and/or large specs. A flag could be - // passed in for tasks whose data has not changed. - if (updated) { + // New task data was added to the local cache, so record which tasks it + // depends on. Add all new tasks that it depends on. for (const auto &parent_id : it->second.GetParentTaskIds()) { - if (old_parents.count(parent_id) == 0) { - AddChild(parent_id, task_id); - } else { - old_parents.erase(parent_id); - } - } - for (const auto &old_parent_id : old_parents) { - RemoveChild(old_parent_id, task_id); + AddChild(parent_id, task_id); } } return updated; @@ -198,15 +183,15 @@ LineageCache::LineageCache(const ClientID &client_id, /// A helper function to add some uncommitted lineage to the local cache. void LineageCache::AddUncommittedLineage(const TaskID &task_id, - const Lineage &uncommitted_lineage, - std::unordered_set &subscribe_tasks) { + const Lineage &uncommitted_lineage) { + RAY_LOG(DEBUG) << "Adding uncommitted task " << task_id << " on " << client_id_; // If the entry is not found in the lineage to merge, then we stop since // there is nothing to copy into the merged lineage. auto entry = uncommitted_lineage.GetEntry(task_id); if (!entry) { return; } - RAY_CHECK(entry->GetStatus() == GcsStatus::UNCOMMITTED_REMOTE); + RAY_CHECK(entry->GetStatus() == GcsStatus::UNCOMMITTED); // Insert a copy of the entry into our cache. const auto &parent_ids = entry->GetParentTaskIds(); @@ -214,90 +199,34 @@ void LineageCache::AddUncommittedLineage(const TaskID &task_id, // if the new entry has an equal or lower GCS status than the current entry // in our cache. This also prevents us from traversing the same node twice. if (lineage_.SetEntry(entry->TaskData(), entry->GetStatus())) { - subscribe_tasks.insert(task_id); + RAY_CHECK(SubscribeTask(task_id)); for (const auto &parent_id : parent_ids) { - AddUncommittedLineage(parent_id, uncommitted_lineage, subscribe_tasks); + AddUncommittedLineage(parent_id, uncommitted_lineage); } } } -bool LineageCache::AddWaitingTask(const Task &task, const Lineage &uncommitted_lineage) { - auto task_id = task.GetTaskSpecification().TaskId(); - RAY_LOG(DEBUG) << "Add waiting task " << task_id << " on " << client_id_; - - // Merge the uncommitted lineage into the lineage cache. Collect the IDs of - // tasks that we should subscribe to. These are all of the tasks that were - // included in the uncommitted lineage that we did not already have in our - // stash. - std::unordered_set subscribe_tasks; - AddUncommittedLineage(task_id, uncommitted_lineage, subscribe_tasks); - // Add the submitted task to the lineage cache as UNCOMMITTED_WAITING. It - // should be marked as UNCOMMITTED_READY once the task starts execution. - auto added = lineage_.SetEntry(task, GcsStatus::UNCOMMITTED_WAITING); - - // Do not subscribe to the waiting task itself. We just added it as - // UNCOMMITTED_WAITING, so the task is local. - subscribe_tasks.erase(task_id); - // Unsubscribe to the waiting task since we may have previously been - // subscribed to it. - UnsubscribeTask(task_id); - // Subscribe to all other tasks that were included in the uncommitted lineage - // and that were not already in the local stash. These tasks haven't been - // committed yet and will be committed by a different node, so we will not - // evict them until a notification for their commit is received. - for (const auto &task_id : subscribe_tasks) { - RAY_CHECK(SubscribeTask(task_id)); - } - - return added; -} - -bool LineageCache::AddReadyTask(const Task &task) { +bool LineageCache::CommitTask(const Task &task) { const TaskID task_id = task.GetTaskSpecification().TaskId(); - RAY_LOG(DEBUG) << "Add ready task " << task_id << " on " << client_id_; + RAY_LOG(DEBUG) << "Committing task " << task_id << " on " << client_id_; - // Set the task to READY. - if (lineage_.SetEntry(task, GcsStatus::UNCOMMITTED_READY)) { - // Attempt to flush the task. + if (lineage_.SetEntry(task, GcsStatus::UNCOMMITTED) || + lineage_.GetEntry(task_id)->GetStatus() == GcsStatus::UNCOMMITTED) { + // Attempt to flush the task if the task is uncommitted. FlushTask(task_id); return true; } else { - // The task was already ready to be committed (UNCOMMITTED_READY) or - // committing (COMMITTING). + // The task was already committing (COMMITTING). return false; } } -bool LineageCache::RemoveWaitingTask(const TaskID &task_id) { - RAY_LOG(DEBUG) << "Remove waiting task " << task_id << " on " << client_id_; - auto entry = lineage_.GetEntryMutable(task_id); - if (!entry) { - // The task was already evicted. - return false; - } - - // If the task is already not in WAITING status, then exit. This should only - // happen when there are two copies of the task executing at the node, due to - // a spurious reconstruction. Then, either the task is already past WAITING - // status, in which case it will be committed, or it is in - // UNCOMMITTED_REMOTE, in which case it was already removed. - if (entry->GetStatus() != GcsStatus::UNCOMMITTED_WAITING) { - return false; - } - - // Reset the status to REMOTE. We keep the task instead of removing it - // completely in case another task is submitted locally that depends on this - // one. - entry->ResetStatus(GcsStatus::UNCOMMITTED_REMOTE); - // The task is now remote, so subscribe to the task to make sure that we'll - // eventually clean it up. - RAY_CHECK(SubscribeTask(task_id)); - return true; -} - void LineageCache::MarkTaskAsForwarded(const TaskID &task_id, const ClientID &node_id) { RAY_CHECK(!node_id.IsNil()); - lineage_.GetEntryMutable(task_id)->MarkExplicitlyForwarded(node_id); + auto entry = lineage_.GetEntryMutable(task_id); + if (entry) { + entry->MarkExplicitlyForwarded(node_id); + } } /// A helper function to get the uncommitted lineage of a task. @@ -345,7 +274,7 @@ Lineage LineageCache::GetUncommittedLineageOrDie(const TaskID &task_id, void LineageCache::FlushTask(const TaskID &task_id) { auto entry = lineage_.GetEntryMutable(task_id); RAY_CHECK(entry); - RAY_CHECK(entry->GetStatus() == GcsStatus::UNCOMMITTED_READY); + RAY_CHECK(entry->GetStatus() < GcsStatus::COMMITTING); gcs::raylet::TaskTable::WriteCallback task_callback = [this]( ray::gcs::AsyncGcsClient *client, const TaskID &id, const protocol::TaskT &data) { @@ -406,11 +335,6 @@ void LineageCache::EvictTask(const TaskID &task_id) { if (!entry) { return; } - // Only evict tasks that we were subscribed to or that we were committing. - if (!(entry->GetStatus() == GcsStatus::UNCOMMITTED_REMOTE || - entry->GetStatus() == GcsStatus::COMMITTING)) { - return; - } // Entries cannot be safely evicted until their parents are all evicted. for (const auto &parent_id : entry->GetParentTaskIds()) { if (ContainsTask(parent_id)) { diff --git a/src/ray/raylet/lineage_cache.h b/src/ray/raylet/lineage_cache.h index 02d98b8cf..2dff0e94a 100644 --- a/src/ray/raylet/lineage_cache.h +++ b/src/ray/raylet/lineage_cache.h @@ -17,19 +17,23 @@ namespace ray { namespace raylet { /// The status of a lineage cache entry according to its status in the GCS. +/// Tasks can only transition to a higher GcsStatus (e.g., an UNCOMMITTED state +/// can become COMMITTING but not vice versa). If a task is evicted from the +/// local cache, it implicitly goes back to state `NONE`, after which it may be +/// added to the local cache again (e.g., if it is forwarded to us again). enum class GcsStatus { /// The task is not in the lineage cache. NONE = 0, - /// The task is being executed or created on a remote node. - UNCOMMITTED_REMOTE, - /// The task is waiting to be executed or created locally. - UNCOMMITTED_WAITING, - /// The task has started execution, but the entry has not been written to the - /// GCS yet. - UNCOMMITTED_READY, - /// The task has been written to the GCS and we are waiting for an - /// acknowledgement of the commit. + /// The task is uncommitted. Unless there is a failure, we will expect a + /// different node to commit this task. + UNCOMMITTED, + /// We flushed this task and are waiting for the commit acknowledgement. COMMITTING, + // TODO(swang): Add a COMMITTED state for tasks for which we received a + // commit acknowledgement, but which we cannot evict yet (due to an ancestor + // that has not been evicted). This is to allow a performance optimization + // that avoids unnecessary subscribes when we receive tasks that were + // already COMMITTED at the sender. }; /// \class LineageEntry @@ -220,37 +224,23 @@ class LineageCache { gcs::TableInterface &task_storage, gcs::PubsubInterface &task_pubsub, uint64_t max_lineage_size); - /// 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. + /// Asynchronously commit a task to the GCS. /// - /// \param task The waiting task to add. + /// \param task The task to commit. It will be moved to the COMMITTING state. + /// \return Whether the task was successfully committed. This can fail if the + /// task was already in the COMMITTING state. + bool CommitTask(const Task &task); + + /// Add a task and its (estimated) uncommitted lineage to the local cache. We + /// will subscribe to commit notifications for all uncommitted tasks to + /// determine when it is safe to evict the lineage from the local cache. + /// + /// \param task_id The ID of the uncommitted task to add. /// \param uncommitted_lineage The task's uncommitted lineage. These are the /// tasks that the given task is data-dependent on, but that have not - /// been made durable in the GCS, as far the task's submitter knows. - /// \return Whether the task was successfully marked as waiting to be - /// committed. This will return false if the task is already waiting to be - /// committed (UNCOMMITTED_WAITING), ready to be committed - /// (UNCOMMITTED_READY), or committing (COMMITTING). - bool AddWaitingTask(const Task &task, const Lineage &uncommitted_lineage); - - /// Add a task that is ready for GCS writeback. This overwrites the task’s - /// mutable fields in the execution specification. - /// - /// \param task The task to set as ready. - /// \return Whether the task was successfully marked as ready to be - /// committed. This will return false if the task is already ready to be - /// committed (UNCOMMITTED_READY) or committing (COMMITTING). - bool AddReadyTask(const Task &task); - - /// 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. - /// \return Whether the task was successfully removed. This will return false - /// if the task is not waiting to be committed. Then, the waiting task has - /// already been removed (UNCOMMITTED_REMOTE), or if it's ready to be - /// committed (UNCOMMITTED_READY) or committing (COMMITTING). - bool RemoveWaitingTask(const TaskID &task_id); + /// been committed to the GCS. This must contain the given task ID. + /// \return Void. + void AddUncommittedLineage(const TaskID &task_id, const Lineage &uncommitted_lineage); /// Mark a task as having been explicitly forwarded to a node. /// The lineage of the task is implicitly assumed to have also been forwarded. @@ -317,9 +307,6 @@ class LineageCache { /// Unsubscribe from notifications for a task. Returns whether the operation /// was successful (whether we were subscribed). bool UnsubscribeTask(const TaskID &task_id); - /// Add a task and its uncommitted lineage to the local stash. - void AddUncommittedLineage(const TaskID &task_id, const Lineage &uncommitted_lineage, - std::unordered_set &subscribe_tasks); /// The client ID, used to request notifications for specific tasks. /// TODO(swang): Move the ClientID into the generic Table implementation. diff --git a/src/ray/raylet/lineage_cache_test.cc b/src/ray/raylet/lineage_cache_test.cc index a61ae846a..e5c126bcf 100644 --- a/src/ray/raylet/lineage_cache_test.cc +++ b/src/ray/raylet/lineage_cache_test.cc @@ -122,15 +122,22 @@ static inline Task ExampleTask(const std::vector &arguments, return task; } +/// Helper method to create a Lineage object with a single task. +Lineage CreateSingletonLineage(const Task &task) { + Lineage singleton_lineage; + singleton_lineage.SetEntry(task, GcsStatus::UNCOMMITTED); + return singleton_lineage; +} + std::vector InsertTaskChain(LineageCache &lineage_cache, std::vector &inserted_tasks, int chain_size, const std::vector &initial_arguments, int64_t num_returns) { - Lineage empty_lineage; std::vector arguments = initial_arguments; for (int i = 0; i < chain_size; i++) { auto task = ExampleTask(arguments, num_returns); - RAY_CHECK(lineage_cache.AddWaitingTask(task, empty_lineage)); + Lineage lineage = CreateSingletonLineage(task); + lineage_cache.AddUncommittedLineage(task.GetTaskSpecification().TaskId(), lineage); inserted_tasks.push_back(task); arguments.clear(); for (int j = 0; j < task.GetTaskSpecification().NumReturns(); j++) { @@ -190,6 +197,34 @@ TEST_F(LineageCacheTest, TestGetUncommittedLineageOrDie) { } } +TEST_F(LineageCacheTest, TestDuplicateUncommittedLineage) { + // Insert a chain of tasks. + std::vector tasks; + auto return_values = + InsertTaskChain(lineage_cache_, tasks, 3, std::vector(), 1); + std::vector task_ids; + for (const auto &task : tasks) { + task_ids.push_back(task.GetTaskSpecification().TaskId()); + } + // Check that we subscribed to each of the uncommitted tasks. + ASSERT_EQ(mock_gcs_.NumRequestedNotifications(), task_ids.size()); + + // Check that if we add the same tasks as UNCOMMITTED again, we do not issue + // duplicate subscribe requests. + Lineage duplicate_lineage; + for (const auto &task : tasks) { + duplicate_lineage.SetEntry(task, GcsStatus::UNCOMMITTED); + } + lineage_cache_.AddUncommittedLineage(task_ids.back(), duplicate_lineage); + ASSERT_EQ(mock_gcs_.NumRequestedNotifications(), task_ids.size()); + + // Check that if we commit one of the tasks, we still do not issue any + // duplicate subscribe requests. + lineage_cache_.CommitTask(tasks.front()); + lineage_cache_.AddUncommittedLineage(task_ids.back(), duplicate_lineage); + ASSERT_EQ(mock_gcs_.NumRequestedNotifications(), task_ids.size()); +} + TEST_F(LineageCacheTest, TestMarkTaskAsForwarded) { // Insert chain of tasks. std::vector tasks; @@ -222,7 +257,7 @@ TEST_F(LineageCacheTest, TestMarkTaskAsForwarded) { ASSERT_EQ(1, uncommitted_lineage_forwarded.GetEntries().size()); } -TEST_F(LineageCacheTest, TestWritebackNoneReady) { +TEST_F(LineageCacheTest, TestWritebackReady) { // Insert a chain of dependent tasks. size_t num_tasks_flushed = 0; std::vector tasks; @@ -231,16 +266,9 @@ TEST_F(LineageCacheTest, TestWritebackNoneReady) { // Check that when no tasks have been marked as ready, we do not flush any // entries. ASSERT_EQ(mock_gcs_.TaskTable().size(), num_tasks_flushed); -} - -TEST_F(LineageCacheTest, TestWritebackReady) { - // Insert a chain of dependent tasks. - size_t num_tasks_flushed = 0; - std::vector tasks; - InsertTaskChain(lineage_cache_, tasks, 3, std::vector(), 1); // Check that after marking the first task as ready, we flush only that task. - ASSERT_TRUE(lineage_cache_.AddReadyTask(tasks.front())); + ASSERT_TRUE(lineage_cache_.CommitTask(tasks.front())); num_tasks_flushed++; ASSERT_EQ(mock_gcs_.TaskTable().size(), num_tasks_flushed); } @@ -253,7 +281,7 @@ TEST_F(LineageCacheTest, TestWritebackOrder) { // Mark all tasks as ready. All tasks should be flushed. for (const auto &task : tasks) { - ASSERT_TRUE(lineage_cache_.AddReadyTask(task)); + ASSERT_TRUE(lineage_cache_.CommitTask(task)); } ASSERT_EQ(mock_gcs_.TaskTable().size(), num_tasks_flushed); @@ -272,12 +300,13 @@ TEST_F(LineageCacheTest, TestEvictChain) { Lineage uncommitted_lineage; for (const auto &task : tasks) { - uncommitted_lineage.SetEntry(task, GcsStatus::UNCOMMITTED_REMOTE); + uncommitted_lineage.SetEntry(task, GcsStatus::UNCOMMITTED); } // Mark the last task as ready to flush. - ASSERT_TRUE(lineage_cache_.AddWaitingTask(tasks.back(), uncommitted_lineage)); + lineage_cache_.AddUncommittedLineage(tasks.back().GetTaskSpecification().TaskId(), + uncommitted_lineage); ASSERT_EQ(lineage_cache_.GetLineage().GetEntries().size(), tasks.size()); - ASSERT_TRUE(lineage_cache_.AddReadyTask(tasks.back())); + ASSERT_TRUE(lineage_cache_.CommitTask(tasks.back())); num_tasks_flushed++; ASSERT_EQ(mock_gcs_.TaskTable().size(), num_tasks_flushed); // Flush acknowledgements. The lineage cache should receive the commit for @@ -320,17 +349,20 @@ TEST_F(LineageCacheTest, TestEvictManyParents) { auto task = ExampleTask({}, 1); parent_tasks.push_back(task); arguments.push_back(task.GetTaskSpecification().ReturnId(0)); - ASSERT_TRUE(lineage_cache_.AddWaitingTask(task, Lineage())); + auto lineage = CreateSingletonLineage(task); + lineage_cache_.AddUncommittedLineage(task.GetTaskSpecification().TaskId(), lineage); } // Create a child task that is dependent on all of the previous tasks. auto child_task = ExampleTask(arguments, 1); - ASSERT_TRUE(lineage_cache_.AddWaitingTask(child_task, Lineage())); + auto lineage = CreateSingletonLineage(child_task); + lineage_cache_.AddUncommittedLineage(child_task.GetTaskSpecification().TaskId(), + lineage); // Flush the child task. Make sure that it remains in the cache, since none // of its parents have been committed yet, and that the uncommitted lineage // still includes all of the parent tasks. size_t total_tasks = parent_tasks.size() + 1; - lineage_cache_.AddReadyTask(child_task); + lineage_cache_.CommitTask(child_task); mock_gcs_.Flush(); ASSERT_EQ(lineage_cache_.GetLineage().GetEntries().size(), total_tasks); ASSERT_EQ(lineage_cache_ @@ -342,7 +374,7 @@ TEST_F(LineageCacheTest, TestEvictManyParents) { // Flush each parent task and check for eviction safety. for (const auto &parent_task : parent_tasks) { - lineage_cache_.AddReadyTask(parent_task); + lineage_cache_.CommitTask(parent_task); mock_gcs_.Flush(); total_tasks--; if (total_tasks > 1) { @@ -364,75 +396,6 @@ TEST_F(LineageCacheTest, TestEvictManyParents) { ASSERT_EQ(lineage_cache_.GetLineage().GetChildrenSize(), 0); } -TEST_F(LineageCacheTest, TestForwardTasksRoundTrip) { - // Insert a chain of dependent tasks. - uint64_t lineage_size = max_lineage_size_ + 1; - std::vector tasks; - InsertTaskChain(lineage_cache_, tasks, lineage_size, std::vector(), 1); - - // Simulate removing each task, forwarding it to another node, then - // receiving the task back again. - for (auto it = tasks.begin(); it != tasks.end(); it++) { - const auto task_id = it->GetTaskSpecification().TaskId(); - // Simulate removing the task and forwarding it to another node. - auto uncommitted_lineage = - lineage_cache_.GetUncommittedLineageOrDie(task_id, ClientID::Nil()); - ASSERT_TRUE(lineage_cache_.RemoveWaitingTask(task_id)); - // Simulate receiving the task again. Make sure we can add the task back. - flatbuffers::FlatBufferBuilder fbb; - auto uncommitted_lineage_message = uncommitted_lineage.ToFlatbuffer(fbb, task_id); - fbb.Finish(uncommitted_lineage_message); - uncommitted_lineage = Lineage( - *flatbuffers::GetRoot(fbb.GetBufferPointer())); - ASSERT_TRUE(lineage_cache_.AddWaitingTask(*it, uncommitted_lineage)); - } -} - -TEST_F(LineageCacheTest, TestForwardTask) { - // Insert a chain of dependent tasks. - size_t num_tasks_flushed = 0; - std::vector tasks; - 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_.GetUncommittedLineageOrDie(task_id_to_remove, ClientID::Nil()); - ASSERT_TRUE(lineage_cache_.RemoveWaitingTask(task_id_to_remove)); - ASSERT_EQ(lineage_cache_.GetLineage().GetEntries().size(), 3); - - // Simulate executing the remaining tasks. - for (const auto &task : tasks) { - ASSERT_TRUE(lineage_cache_.AddReadyTask(task)); - num_tasks_flushed++; - } - // 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. - ASSERT_EQ(mock_gcs_.TaskTable().size(), num_tasks_flushed); - mock_gcs_.Flush(); - ASSERT_EQ(lineage_cache_.GetLineage().GetEntries().size(), 2); - - // 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++; - ASSERT_EQ(mock_gcs_.TaskTable().size(), 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(); - ASSERT_EQ(mock_gcs_.SubscribedTasks().size(), 0); - ASSERT_EQ(lineage_cache_.GetLineage().GetEntries().size(), 0); - ASSERT_EQ(lineage_cache_.GetLineage().GetChildrenSize(), 0); -} - TEST_F(LineageCacheTest, TestEviction) { // Insert a chain of dependent tasks. uint64_t lineage_size = max_lineage_size_ + 1; @@ -440,12 +403,6 @@ TEST_F(LineageCacheTest, TestEviction) { std::vector tasks; InsertTaskChain(lineage_cache_, tasks, lineage_size, std::vector(), 1); - // Simulate forwarding the chain of tasks to a remote node. - for (const auto &task : tasks) { - auto task_id = task.GetTaskSpecification().TaskId(); - ASSERT_TRUE(lineage_cache_.RemoveWaitingTask(task_id)); - } - // Check that the last task in the chain still has all tasks in its // uncommitted lineage. const auto last_task_id = tasks.back().GetTaskSpecification().TaskId(); @@ -500,12 +457,6 @@ TEST_F(LineageCacheTest, TestOutOfOrderEviction) { std::vector tasks; InsertTaskChain(lineage_cache_, tasks, lineage_size, std::vector(), 1); - // Simulate forwarding the chain of tasks to a remote node. - for (const auto &task : tasks) { - auto task_id = task.GetTaskSpecification().TaskId(); - ASSERT_TRUE(lineage_cache_.RemoveWaitingTask(task_id)); - } - // Check that the last task in the chain still has all tasks in its // uncommitted lineage. const auto last_task_id = tasks.back().GetTaskSpecification().TaskId(); @@ -545,19 +496,15 @@ TEST_F(LineageCacheTest, TestEvictionUncommittedChildren) { std::vector tasks; InsertTaskChain(lineage_cache_, tasks, lineage_size, std::vector(), 1); - // Simulate forwarding the chain of tasks to a remote node. - for (const auto &task : tasks) { - auto task_id = task.GetTaskSpecification().TaskId(); - ASSERT_TRUE(lineage_cache_.RemoveWaitingTask(task_id)); - } - // Add more tasks to the lineage cache that will remain local. Each of these // tasks is dependent one of the tasks that was forwarded above. for (const auto &task : tasks) { auto return_id = task.GetTaskSpecification().ReturnId(0); auto dependent_task = ExampleTask({return_id}, 1); - ASSERT_TRUE(lineage_cache_.AddWaitingTask(dependent_task, Lineage())); - ASSERT_TRUE(lineage_cache_.AddReadyTask(dependent_task)); + auto lineage = CreateSingletonLineage(dependent_task); + lineage_cache_.AddUncommittedLineage(dependent_task.GetTaskSpecification().TaskId(), + lineage); + ASSERT_TRUE(lineage_cache_.CommitTask(dependent_task)); // Once the forwarded tasks are evicted from the lineage cache, we expect // each of these dependent tasks to be flushed, since all of their // dependencies have been committed. diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index e3fd9a0df..07dca3c7a 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -693,11 +693,6 @@ void NodeManager::HandleActorStateTransition(const ActorID &actor_id, // known. auto created_actor_methods = local_queues_.RemoveTasks(created_actor_method_ids); for (const auto &method : created_actor_methods) { - if (!lineage_cache_.RemoveWaitingTask(method.GetTaskSpecification().TaskId())) { - RAY_LOG(WARNING) << "Task " << method.GetTaskSpecification().TaskId() - << " already removed from the lineage cache. This is most " - "likely due to reconstruction."; - } // Maintain the invariant that if a task is in the // MethodsWaitingForActorCreation queue, then it is subscribed to its // respective actor creation task. Since the actor location is now known, @@ -1466,10 +1461,6 @@ void NodeManager::TreatTaskAsFailed(const Task &task, const ErrorType &error_typ current_time_ms())); } } - // A task failing is equivalent to assigning and finishing the task, so clean - // up any leftover state as for any task dispatched and removed from the - // local queue. - lineage_cache_.AddReadyTask(task); task_dependency_manager_.TaskCanceled(spec.TaskId()); // Notify the task dependency manager that we no longer need this task's // object dependencies. TODO(swang): Ideally, we would check the return value @@ -1538,10 +1529,14 @@ void NodeManager::SubmitTask(const Task &task, const Lineage &uncommitted_lineag } // Add the task and its uncommitted lineage to the lineage cache. - if (!lineage_cache_.AddWaitingTask(task, uncommitted_lineage)) { - RAY_LOG(WARNING) - << "Task " << task_id - << " already in lineage cache. This is most likely due to reconstruction."; + if (forwarded) { + lineage_cache_.AddUncommittedLineage(task_id, uncommitted_lineage); + } else { + if (!lineage_cache_.CommitTask(task)) { + RAY_LOG(WARNING) + << "Task " << task_id + << " already committed to the GCS. This is most likely due to reconstruction."; + } } if (spec.IsActorTask()) { @@ -1869,32 +1864,14 @@ bool NodeManager::AssignTask(const Task &task) { actor_entry->second.AddHandle(new_handle_id, execution_dependency); } - // If the task was an actor task, then record this execution to - // guarantee consistency in the case of reconstruction. - auto execution_dependency = actor_entry->second.GetExecutionDependency(); - // The execution dependency is initialized to the actor creation task's - // return value, and is subsequently updated to the assigned tasks' - // return values, so it should never be nil. - RAY_CHECK(!execution_dependency.IsNil()); - // Update the task's execution dependencies to reflect the actual - // execution order, to support deterministic reconstruction. - // NOTE(swang): The update of an actor task's execution dependencies is - // performed asynchronously. This means that if this node manager dies, - // 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. - // (SetExecutionDependencies takes a non-const so copy task in a - // on-const variable.) - assigned_task.SetExecutionDependencies({execution_dependency}); + // TODO(swang): For actors with multiple actor handles, to + // guarantee that tasks are replayed in the same order after a + // failure, we must update the task's execution dependency to be + // the actor's current execution dependency. } else { RAY_CHECK(spec.NewActorHandles().empty()); } - // We started running the task, so the task is ready to write to GCS. - if (!lineage_cache_.AddReadyTask(assigned_task)) { - RAY_LOG(WARNING) << "Task " << spec.TaskId() << " already in lineage cache." - << " This is most likely due to reconstruction."; - } // Mark the task as running. // (See design_docs/task_states.rst for the state transition diagram.) local_queues_.QueueTasks({assigned_task}, TaskState::RUNNING); @@ -2260,9 +2237,6 @@ void NodeManager::ForwardTaskOrResubmit(const Task &task, // Temporarily move the RESUBMITTED task to the SWAP queue while the // timer is active. local_queues_.QueueTasks({task}, TaskState::SWAP); - // Remove the task from the lineage cache. The task will get added back - // once it is resubmitted. - lineage_cache_.RemoveWaitingTask(task_id); } else { // The task is not for an actor and may therefore be placed on another // node immediately. Send it to the scheduling policy to be placed again. @@ -2327,17 +2301,9 @@ void NodeManager::ForwardTask( if (status.ok()) { const auto &spec = task.GetTaskSpecification(); - // If we were able to forward the task, remove the forwarded task from the - // lineage cache since the receiving node is now responsible for writing - // the task to the GCS. - if (!lineage_cache_.RemoveWaitingTask(task_id)) { - RAY_LOG(WARNING) << "Task " << task_id << " already removed from the lineage" - << " cache. This is most likely due to reconstruction."; - } else { - // Mark as forwarded so that the task and its lineage is not - // re-forwarded in the future to the receiving node. - lineage_cache_.MarkTaskAsForwarded(task_id, node_id); - } + // Mark as forwarded so that the task and its lineage are not + // re-forwarded in the future to the receiving node. + lineage_cache_.MarkTaskAsForwarded(task_id, node_id); // Notify the task dependency manager that we are no longer responsible // for executing this task.