From e6a81d40a56554a6b707e8775d0ceb81c12d697f Mon Sep 17 00:00:00 2001 From: Philipp Moritz Date: Wed, 10 Jul 2019 13:33:41 -0700 Subject: [PATCH] [stability] Make task result for RemoveTask optional (#5146) * make task result for RemoveTask optional * lint * update * update * update * rename * lint --- src/ray/raylet/node_manager.cc | 53 +++++++++++++++++++--------- src/ray/raylet/scheduling_queue.cc | 16 ++++++--- src/ray/raylet/scheduling_queue.h | 6 ++-- src/ray/raylet/task.h | 4 +++ src/ray/raylet/task_execution_spec.h | 4 +++ src/ray/raylet/task_spec.h | 3 ++ src/ray/rpc/message_wrapper.h | 3 ++ 7 files changed, 66 insertions(+), 23 deletions(-) diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index f9b5864ae..da275f5a6 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -549,8 +549,11 @@ void NodeManager::HeartbeatAdded(const ClientID &client_id, std::unordered_set local_task_ids; for (const auto &task_id : decision) { // (See design_docs/task_states.rst for the state transition diagram.) + Task task; TaskState state; - const auto task = local_queues_.RemoveTask(task_id, &state); + if (!local_queues_.RemoveTask(task_id, &task, &state)) { + return; + } // Since we are spilling back from the ready and waiting queues, we need // to unsubscribe the dependencies. if (state != TaskState::INFEASIBLE) { @@ -980,8 +983,10 @@ void NodeManager::ProcessDisconnectClientMessage( // If the worker was an actor, the task was already cleaned up in // `HandleDisconnectedActor`. if (actor_id.IsNil()) { - const Task &task = local_queues_.RemoveTask(task_id); - TreatTaskAsFailed(task, ErrorType::WORKER_DIED); + Task task; + if (local_queues_.RemoveTask(task_id, &task)) { + TreatTaskAsFailed(task, ErrorType::WORKER_DIED); + } } if (!intentional_disconnect) { @@ -1035,6 +1040,8 @@ void NodeManager::ProcessDisconnectClientMessage( << "job_id: " << worker->GetAssignedJobId(); } + client->Close(); + // TODO(rkn): Tell the object manager that this client has disconnected so // that it can clean up the wait requests for this client. Currently I think // these can be leaked. @@ -1315,10 +1322,12 @@ void NodeManager::ScheduleTasks( } else { // TODO(atumanov): need a better interface for task exit on forward. // (See design_docs/task_states.rst for the state transition diagram.) - const auto task = local_queues_.RemoveTask(task_id); - // Attempt to forward the task. If this fails to forward the task, - // the task will be resubmit locally. - ForwardTaskOrResubmit(task, client_id); + Task task; + if (local_queues_.RemoveTask(task_id, &task)) { + // Attempt to forward the task. If this fails to forward the task, + // the task will be resubmit locally. + ForwardTaskOrResubmit(task, client_id); + } } } @@ -1604,7 +1613,8 @@ void NodeManager::HandleTaskBlocked(const std::shared_ptr // worker as blocked. This temporarily releases any resources that the // worker holds while it is blocked. if (!worker->IsBlocked() && current_task_id == worker->GetAssignedTaskId()) { - const auto task = local_queues_.RemoveTask(current_task_id); + Task task; + RAY_CHECK(local_queues_.RemoveTask(current_task_id, &task)); local_queues_.QueueTasks({task}, TaskState::RUNNING); // Get the CPU resources required by the running task. const auto required_resources = task.GetTaskSpecification().GetRequiredResources(); @@ -1653,7 +1663,8 @@ void NodeManager::HandleTaskUnblocked( // the worker. if (worker->IsBlocked() && current_task_id == worker->GetAssignedTaskId()) { // (See design_docs/task_states.rst for the state transition diagram.) - const auto task = local_queues_.RemoveTask(current_task_id); + Task task; + RAY_CHECK(local_queues_.RemoveTask(current_task_id, &task)); local_queues_.QueueTasks({task}, TaskState::RUNNING); // Get the CPU resources required by the running task. const auto required_resources = task.GetTaskSpecification().GetRequiredResources(); @@ -1771,8 +1782,12 @@ bool NodeManager::AssignTask(const Task &task) { static_cast(protocol::MessageType::ExecuteTask), fbb.GetSize(), fbb.GetBufferPointer(), [this, worker, task_id](ray::Status status) { // Remove the ASSIGNED task from the SWAP queue. + Task assigned_task; TaskState state; - auto assigned_task = local_queues_.RemoveTask(task_id, &state); + if (!local_queues_.RemoveTask(task_id, &assigned_task, &state)) { + return; + } + RAY_CHECK(state == TaskState::SWAP); if (status.ok()) { @@ -1844,7 +1859,8 @@ void NodeManager::FinishAssignedTask(Worker &worker) { RAY_LOG(DEBUG) << "Finished task " << task_id; // (See design_docs/task_states.rst for the state transition diagram.) - const auto task = local_queues_.RemoveTask(task_id); + Task task; + RAY_CHECK(local_queues_.RemoveTask(task_id, &task)); // Release task's resources. The worker's lifetime resources are still held. auto const &task_resources = worker.GetTaskResourceIds(); @@ -2171,11 +2187,13 @@ void NodeManager::ForwardTaskOrResubmit(const Task &task, RAY_LOG(INFO) << "Resubmitting task " << task_id << " because ForwardTask failed."; // Remove the RESUBMITTED task from the SWAP queue. + Task task; TaskState state; - const auto task = local_queues_.RemoveTask(task_id, &state); - RAY_CHECK(state == TaskState::SWAP); - // Submit the task again. - SubmitTask(task, Lineage()); + if (local_queues_.RemoveTask(task_id, &task, &state)) { + RAY_CHECK(state == TaskState::SWAP); + // Submit the task again. + SubmitTask(task, Lineage()); + } }); // Temporarily move the RESUBMITTED task to the SWAP queue while the // timer is active. @@ -2246,8 +2264,11 @@ void NodeManager::ForwardTask( client->ForwardTask(request, [this, on_error, task_id, node_id]( Status status, const rpc::ForwardTaskReply &reply) { // Remove the FORWARDING task from the SWAP queue. + Task task; TaskState state; - const auto task = local_queues_.RemoveTask(task_id, &state); + if (!local_queues_.RemoveTask(task_id, &task, &state)) { + return; + } RAY_CHECK(state == TaskState::SWAP); if (status.ok()) { diff --git a/src/ray/raylet/scheduling_queue.cc b/src/ray/raylet/scheduling_queue.cc index 701b0c06b..9270b4798 100644 --- a/src/ray/raylet/scheduling_queue.cc +++ b/src/ray/raylet/scheduling_queue.cc @@ -247,7 +247,8 @@ std::vector SchedulingQueue::RemoveTasks(std::unordered_set &task_ return removed_tasks; } -Task SchedulingQueue::RemoveTask(const TaskID &task_id, TaskState *removed_task_state) { +bool SchedulingQueue::RemoveTask(const TaskID &task_id, Task *removed_task, + TaskState *removed_task_state) { std::vector removed_tasks; std::unordered_set task_id_set = {task_id}; // Try to find the task to remove in the queues. @@ -273,10 +274,15 @@ Task SchedulingQueue::RemoveTask(const TaskID &task_id, TaskState *removed_task_ } // Make sure we got the removed task. - RAY_CHECK(removed_tasks.size() == 1) << task_id; - const auto &task = removed_tasks.front(); - RAY_CHECK(task.GetTaskSpecification().TaskId() == task_id); - return task; + if (removed_tasks.size() == 1) { + *removed_task = removed_tasks.front(); + RAY_CHECK(removed_task->GetTaskSpecification().TaskId() == task_id); + return true; + } + RAY_LOG(DEBUG) << "Task " << task_id + << " that is to be removed could not be found any more." + << " Probably its driver was removed."; + return false; } void SchedulingQueue::MoveTasks(std::unordered_set &task_ids, TaskState src_state, diff --git a/src/ray/raylet/scheduling_queue.h b/src/ray/raylet/scheduling_queue.h index 2cb0ca7de..4a6b4f50b 100644 --- a/src/ray/raylet/scheduling_queue.h +++ b/src/ray/raylet/scheduling_queue.h @@ -227,10 +227,12 @@ class SchedulingQueue { /// /// \param task_id The task ID to remove from the queue. The corresponding /// task must be contained in the queue. + /// \param task The removed task will be written here, if any. /// \param task_state If this is not nullptr, then the state of the removed /// task will be written here. - /// \return The task that was removed. - Task RemoveTask(const TaskID &task_id, TaskState *task_state = nullptr); + /// \return true if the task was removed, false if it is not in the queue. + bool RemoveTask(const TaskID &task_id, Task *removed_task, + TaskState *removed_task_state = nullptr); /// Remove a driver task ID. This is an empty task used to represent a driver. /// diff --git a/src/ray/raylet/task.h b/src/ray/raylet/task.h index e1a8a878b..0f80060fb 100644 --- a/src/ray/raylet/task.h +++ b/src/ray/raylet/task.h @@ -21,6 +21,10 @@ namespace raylet { /// time. class Task { public: + /// Construct an empty task. This should only be used to pass a task + /// as an out parameter to a function or method. + Task() {} + /// Construct a `Task` object from a protobuf message. /// /// \param message The protobuf message. diff --git a/src/ray/raylet/task_execution_spec.h b/src/ray/raylet/task_execution_spec.h index bdc16c8b4..e78d064f5 100644 --- a/src/ray/raylet/task_execution_spec.h +++ b/src/ray/raylet/task_execution_spec.h @@ -17,6 +17,10 @@ using rpc::MessageWrapper; /// Wrapper class of protobuf `TaskExecutionSpec`, see `common.proto` for details. class TaskExecutionSpecification : public MessageWrapper { public: + /// Construct an emtpy task execution specification. This should not be used + /// directly. + TaskExecutionSpecification() {} + /// Construct from a protobuf message object. /// The input message will be **copied** into this object. /// diff --git a/src/ray/raylet/task_spec.h b/src/ray/raylet/task_spec.h index ce87d2580..ea72efc15 100644 --- a/src/ray/raylet/task_spec.h +++ b/src/ray/raylet/task_spec.h @@ -26,6 +26,9 @@ using rpc::TaskType; /// Wrapper class of protobuf `TaskSpec`, see `common.proto` for details. class TaskSpecification : public MessageWrapper { public: + /// Construct an empty task specification. This should not be used directly. + TaskSpecification() {} + /// Construct from a protobuf message object. /// The input message will be **copied** into this object. /// diff --git a/src/ray/rpc/message_wrapper.h b/src/ray/rpc/message_wrapper.h index 452ad24e5..56fb5c53a 100644 --- a/src/ray/rpc/message_wrapper.h +++ b/src/ray/rpc/message_wrapper.h @@ -11,6 +11,9 @@ namespace rpc { template class MessageWrapper { public: + /// Construct an empty message wrapper. This should not be used directly. + MessageWrapper() {} + /// Construct from a protobuf message object. /// The input message will be **copied** into this object. ///