diff --git a/src/ray/raylet/lineage_cache.cc b/src/ray/raylet/lineage_cache.cc index d1caa8cd6..7d724cb6d 100644 --- a/src/ray/raylet/lineage_cache.cc +++ b/src/ray/raylet/lineage_cache.cc @@ -47,9 +47,8 @@ Lineage::Lineage(const protocol::ForwardTaskRequest &task_request) { // Deserialize and set entries for the uncommitted tasks. auto tasks = task_request.uncommitted_tasks(); for (auto it = tasks->begin(); it != tasks->end(); it++) { - auto task = Task(**it); - LineageEntry entry(task, GcsStatus::UNCOMMITTED_REMOTE); - RAY_CHECK(SetEntry(std::move(entry))); + const auto &task = **it; + RAY_CHECK(SetEntry(task, GcsStatus::UNCOMMITTED_REMOTE)); } } @@ -71,24 +70,22 @@ boost::optional Lineage::GetEntryMutable(const UniqueID &task_id } } -bool Lineage::SetEntry(LineageEntry &&new_entry) { +bool Lineage::SetEntry(const Task &task, GcsStatus status) { // Get the status of the current entry at the key. - auto task_id = new_entry.GetEntryId(); - GcsStatus current_status = GcsStatus::NONE; - auto current_entry = PopEntry(task_id); + auto task_id = task.GetTaskSpecification().TaskId(); + auto current_entry = GetEntryMutable(task_id); if (current_entry) { - current_status = current_entry->GetStatus(); - } - - if (current_status < new_entry.GetStatus()) { - // If the new status is greater, then overwrite the current entry. + 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); + return true; + } + return false; + } else { + LineageEntry new_entry(task, status); entries_.emplace(std::make_pair(task_id, std::move(new_entry))); return true; - } else { - // If the new status is not greater, then the new entry is invalid. Replace - // the current entry at the key. - entries_.emplace(std::make_pair(task_id, std::move(*current_entry))); - return false; } } @@ -156,12 +153,11 @@ void MergeLineageHelper(const UniqueID &task_id, const Lineage &lineage_from, } // Insert a copy of the entry into lineage_to. - LineageEntry entry_copy = *entry; - auto parent_ids = entry_copy.GetParentTaskIds(); + 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. - if (lineage_to.SetEntry(std::move(entry_copy))) { + if (lineage_to.SetEntry(entry->TaskData(), entry->GetStatus())) { for (const auto &parent_id : parent_ids) { MergeLineageHelper(parent_id, lineage_from, lineage_to, stopping_condition); } @@ -192,20 +188,20 @@ void LineageCache::AddWaitingTask(const Task &task, const Lineage &uncommitted_l // Add the submitted task to the lineage cache as UNCOMMITTED_WAITING. It // should be marked as UNCOMMITTED_READY once the task starts execution. - LineageEntry task_entry(task, GcsStatus::UNCOMMITTED_WAITING); - RAY_CHECK(lineage_.SetEntry(std::move(task_entry))); + RAY_CHECK(lineage_.SetEntry(task, GcsStatus::UNCOMMITTED_WAITING)); } void LineageCache::AddReadyTask(const Task &task) { const TaskID task_id = task.GetTaskSpecification().TaskId(); // Tasks can only become READY if they were in WAITING. - auto entry = lineage_.GetEntry(task_id); + auto entry = lineage_.GetEntryMutable(task_id); RAY_CHECK(entry); RAY_CHECK(entry->GetStatus() == GcsStatus::UNCOMMITTED_WAITING); - auto new_entry = LineageEntry(task, GcsStatus::UNCOMMITTED_READY); - RAY_CHECK(lineage_.SetEntry(std::move(new_entry))); + entry->SetStatus(GcsStatus::UNCOMMITTED_READY); + // TaskSepc is immutable, just update TaskExecSpec. + entry->TaskDataMutable().CopyTaskExecutionSpec(task); // Attempt to flush the task. bool flushed = FlushTask(task_id); if (!flushed) { @@ -231,7 +227,11 @@ uint64_t LineageCache::CountUnsubscribedLineage(const TaskID &task_id) const { } void LineageCache::RemoveWaitingTask(const TaskID &task_id) { - auto entry = lineage_.PopEntry(task_id); + auto entry = lineage_.GetEntryMutable(task_id); + if (!entry) { + return; + } + // It's only okay to remove a task that is waiting for execution. // TODO(swang): Is this necessarily true when there is reconstruction? RAY_CHECK(entry->GetStatus() == GcsStatus::UNCOMMITTED_WAITING); @@ -239,7 +239,6 @@ void LineageCache::RemoveWaitingTask(const TaskID &task_id) { // completely in case another task is submitted locally that depends on this // one. entry->ResetStatus(GcsStatus::UNCOMMITTED_REMOTE); - RAY_CHECK(lineage_.SetEntry(std::move(*entry))); // Request a notification for every max_lineage_size_ tasks, // so that the task and its uncommitted lineage can be evicted @@ -318,9 +317,9 @@ bool LineageCache::FlushTask(const TaskID &task_id) { // We successfully wrote the task, so mark it as committing. // TODO(swang): Use a batched interface and write with all object entries. - auto entry = lineage_.PopEntry(task_id); + auto entry = lineage_.GetEntryMutable(task_id); + RAY_CHECK(entry); RAY_CHECK(entry->SetStatus(GcsStatus::COMMITTING)); - RAY_CHECK(lineage_.SetEntry(std::move(*entry))); } return all_arguments_committed; } diff --git a/src/ray/raylet/lineage_cache.h b/src/ray/raylet/lineage_cache.h index e2467436c..9ce55d033 100644 --- a/src/ray/raylet/lineage_cache.h +++ b/src/ray/raylet/lineage_cache.h @@ -120,10 +120,10 @@ class Lineage { /// GCS status than the current. The current entry's object or task data will /// also be overwritten. /// - /// \param entry The new entry to set in the lineage, if its GCS status is - /// greater than the current entry. + /// \param task The task data to set, if status is greater than the current entry. + /// \param status The GCS status. /// \return Whether the entry was set. - bool SetEntry(LineageEntry &&entry); + bool SetEntry(const Task &task, GcsStatus status); /// Delete and return an entry from the lineage. /// diff --git a/src/ray/raylet/task.cc b/src/ray/raylet/task.cc index 0f209d33b..b075dfc03 100644 --- a/src/ray/raylet/task.cc +++ b/src/ray/raylet/task.cc @@ -58,6 +58,10 @@ bool Task::DependsOn(const ObjectID &object_id) const { return false; } +void Task::CopyTaskExecutionSpec(const Task &task) { + task_execution_spec_ = task.GetTaskExecutionSpecReadonly(); +} + } // namespace raylet } // namespace ray diff --git a/src/ray/raylet/task.h b/src/ray/raylet/task.h index 63bdb35f3..7897d0200 100644 --- a/src/ray/raylet/task.h +++ b/src/ray/raylet/task.h @@ -73,6 +73,10 @@ class Task { /// false otherwise. bool DependsOn(const ObjectID &object_id) const; + /// Update the dynamic/mutable information for this task. + /// \param task Task structure with updated dynamic information. + void CopyTaskExecutionSpec(const Task &task); + private: /// Task execution specification, consisting of all dynamic/mutable /// information about this task determined at execution time..