From dbfe752c3974ce3d864cfd08d5e52c53c04b0d12 Mon Sep 17 00:00:00 2001 From: Edward Oakes Date: Mon, 2 Dec 2019 16:39:02 -0800 Subject: [PATCH] Don't add assigned tasks to SWAP queue (#6325) --- src/ray/raylet/node_manager.cc | 186 +++++++++++++++------------------ src/ray/raylet/node_manager.h | 18 ++-- src/ray/raylet/worker.cc | 15 +-- src/ray/raylet/worker.h | 3 +- 4 files changed, 99 insertions(+), 123 deletions(-) diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index f833924c4..ddc4b9800 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -813,8 +813,6 @@ std::unordered_map> MakeTasksByClass( void NodeManager::DispatchTasks( const std::unordered_map> &tasks_by_class) { - std::unordered_set removed_task_ids; - // Dispatch tasks in priority order by class. This avoids starvation problems where // one class of tasks become stuck behind others in the queue, causing Ray to start // many workers. See #3644 for a more detailed description of this issue. @@ -846,16 +844,20 @@ void NodeManager::DispatchTasks( // once the first task is not feasible, we can break out of this loop break; } - if (AssignTask(task, &post_assign_callbacks)) { - removed_task_ids.insert(task_id); + + // Try to get an idle worker to execute this task. If nullptr, there + // aren't any available workers so we can't assign the task. + std::shared_ptr worker = + worker_pool_.PopWorker(task.GetTaskSpecification()); + if (worker != nullptr) { + AssignTask(worker, task, &post_assign_callbacks); } } } - // Move the ASSIGNED task to the SWAP queue so that we remember that we have - // it queued locally. Once the GetTaskReply has been sent, the task will get - // re-queued, depending on whether the message succeeded or not. - local_queues_.MoveTasks(removed_task_ids, TaskState::READY, TaskState::SWAP); - for (auto func : post_assign_callbacks) { + // Call the callbacks from the AssignTask calls above. These need to be called + // after the above loop, as they may alter the scheduling queues and invalidate + // the loop iterator. + for (auto &func : post_assign_callbacks) { func(); } } @@ -2042,9 +2044,10 @@ void NodeManager::EnqueuePlaceableTask(const Task &task) { task_dependency_manager_.TaskPending(task); } -bool NodeManager::AssignTask(const Task &task, +void NodeManager::AssignTask(const std::shared_ptr &worker, const Task &task, std::vector> *post_assign_callbacks) { const TaskSpecification &spec = task.GetTaskSpecification(); + RAY_CHECK(post_assign_callbacks); // If this is an actor task, check that the new task has the correct counter. if (spec.IsActorTask()) { @@ -2057,14 +2060,6 @@ bool NodeManager::AssignTask(const Task &task, << spec.TaskId() << " has: " << spec.ActorCounter(); } - // Try to get an idle worker that can execute this task. - std::shared_ptr worker = worker_pool_.PopWorker(spec); - if (worker == nullptr) { - // There are no workers that can execute this task. - // We couldn't assign this task, as no worker available. - return false; - } - RAY_LOG(DEBUG) << "Assigning task " << spec.TaskId() << " to worker with pid " << worker->Pid(); flatbuffers::FlatBufferBuilder fbb; @@ -2085,40 +2080,28 @@ bool NodeManager::AssignTask(const Task &task, } auto task_id = spec.TaskId(); - auto finish_assign_task_callback = [this, worker, task_id](Status status) { - // NOTE: we cannot directly call `FinishAssignTask` here because - // it assumes the task is in SWAP queue, thus we need to delay invoking this - // function after the assigned tasks are moved from READY queue to SWAP queue - // in `DispatchTasks`. - // Another option is to move the tasks to SWAP queue here just before calling - // `FinishAssignTask` so we can save an io_service post, at the - // expense of calling `MoveTask` for each of the assigned tasks. - // TODO(zhijunfu): after all workers are fully migrated to push mode, the - // `post` below and swap queue can be removed. - io_service_.post([this, status, worker, task_id]() { - FinishAssignTask(task_id, *worker, status.ok()); - }); - }; - - ResourceIdSet resource_id_set = - worker->GetTaskResourceIds().Plus(worker->GetLifetimeResourceIds()); - if (task.OnDispatch() != nullptr) { task.OnDispatch()(worker, initial_config_.node_manager_address, worker->Port()); - if (post_assign_callbacks != nullptr) { - // Moves the tasks from SWAP to RUNNING state atomically. This avoids race - // conditions with ReturnLease requests. - post_assign_callbacks->push_back( - [this, worker, task_id]() { FinishAssignTask(task_id, *worker, true); }); - } + post_assign_callbacks->push_back([this, worker, task_id]() { + FinishAssignTask(worker, task_id, /*success=*/true); + }); } else { - worker->AssignTask(task, resource_id_set, finish_assign_task_callback); + ResourceIdSet resource_id_set = + worker->GetTaskResourceIds().Plus(worker->GetLifetimeResourceIds()); + if (worker->AssignTask(task, resource_id_set).ok()) { + RAY_LOG(DEBUG) << "Assigned task " << task_id << " to worker " + << worker->WorkerId(); + post_assign_callbacks->push_back([this, worker, task_id]() { + FinishAssignTask(worker, task_id, /*success=*/true); + }); + } else { + RAY_LOG(ERROR) << "Failed to assign task " << task_id << " to worker " + << worker->WorkerId() << ", disconnecting client"; + post_assign_callbacks->push_back([this, worker, task_id]() { + FinishAssignTask(worker, task_id, /*success=*/false); + }); + } } - - // We assigned this task to a worker. - // (Note this means that we sent the task to the worker. The assignment - // might still fail if the worker fails in the meantime, for instance.) - return true; } void NodeManager::FinishAssignedTask(Worker &worker) { @@ -2243,51 +2226,52 @@ void NodeManager::FinishAssignedActorTask(Worker &worker, const Task &task) { // Lookup the parent actor id. auto parent_task_id = task_spec.ParentTaskId(); int port = worker.Port(); - RAY_CHECK_OK(gcs_client_->raylet_task_table().Lookup( - JobID::Nil(), parent_task_id, - /*success_callback=*/ - [this, task_spec, resumed_from_checkpoint, port]( - ray::gcs::RedisGcsClient *client, const TaskID &parent_task_id, - const TaskTableData &parent_task_data) { - // The task was in the GCS task table. Use the stored task spec to - // get the parent actor id. - Task parent_task(parent_task_data.task()); - ActorID parent_actor_id = ActorID::Nil(); - if (parent_task.GetTaskSpecification().IsActorCreationTask()) { - parent_actor_id = parent_task.GetTaskSpecification().ActorCreationId(); - } else if (parent_task.GetTaskSpecification().IsActorTask()) { - parent_actor_id = parent_task.GetTaskSpecification().ActorId(); - } - FinishAssignedActorCreationTask(parent_actor_id, task_spec, - resumed_from_checkpoint, port); - }, - /*failure_callback=*/ - [this, task_spec, resumed_from_checkpoint, port](ray::gcs::RedisGcsClient *client, - const TaskID &parent_task_id) { - // The parent task was not in the GCS task table. It should most likely be in - // the - // lineage cache. - ActorID parent_actor_id = ActorID::Nil(); - if (lineage_cache_.ContainsTask(parent_task_id)) { - // Use a copy of the cached task spec to get the parent actor id. - Task parent_task = lineage_cache_.GetTaskOrDie(parent_task_id); - if (parent_task.GetTaskSpecification().IsActorCreationTask()) { - parent_actor_id = parent_task.GetTaskSpecification().ActorCreationId(); - } else if (parent_task.GetTaskSpecification().IsActorTask()) { - parent_actor_id = parent_task.GetTaskSpecification().ActorId(); - } - } else { - RAY_LOG(WARNING) - << "Task metadata not found in either GCS or lineage cache. It may have " - "been " - "evicted " - << "by the redis LRU configuration. Consider increasing the memory " - "allocation via " - << "ray.init(redis_max_memory=)."; - } - FinishAssignedActorCreationTask(parent_actor_id, task_spec, - resumed_from_checkpoint, port); - })); + RAY_CHECK_OK( + gcs_client_->raylet_task_table().Lookup( + JobID::Nil(), parent_task_id, + /*success_callback=*/ + [this, task_spec, resumed_from_checkpoint, port]( + ray::gcs::RedisGcsClient *client, const TaskID &parent_task_id, + const TaskTableData &parent_task_data) { + // The task was in the GCS task table. Use the stored task spec to + // get the parent actor id. + Task parent_task(parent_task_data.task()); + ActorID parent_actor_id = ActorID::Nil(); + if (parent_task.GetTaskSpecification().IsActorCreationTask()) { + parent_actor_id = parent_task.GetTaskSpecification().ActorCreationId(); + } else if (parent_task.GetTaskSpecification().IsActorTask()) { + parent_actor_id = parent_task.GetTaskSpecification().ActorId(); + } + FinishAssignedActorCreationTask(parent_actor_id, task_spec, + resumed_from_checkpoint, port); + }, + /*failure_callback=*/ + [this, task_spec, resumed_from_checkpoint, port]( + ray::gcs::RedisGcsClient *client, const TaskID &parent_task_id) { + // The parent task was not in the GCS task table. It should most likely be + // in the lineage cache. + ActorID parent_actor_id = ActorID::Nil(); + if (lineage_cache_.ContainsTask(parent_task_id)) { + // Use a copy of the cached task spec to get the parent actor id. + Task parent_task = lineage_cache_.GetTaskOrDie(parent_task_id); + if (parent_task.GetTaskSpecification().IsActorCreationTask()) { + parent_actor_id = parent_task.GetTaskSpecification().ActorCreationId(); + } else if (parent_task.GetTaskSpecification().IsActorTask()) { + parent_actor_id = parent_task.GetTaskSpecification().ActorId(); + } + } else { + RAY_LOG(WARNING) + << "Task metadata not found in either GCS or lineage cache. It may " + "have " + "been " + "evicted " + << "by the redis LRU configuration. Consider increasing the memory " + "allocation via " + << "ray.init(redis_max_memory=)."; + } + FinishAssignedActorCreationTask(parent_actor_id, task_spec, + resumed_from_checkpoint, port); + })); } else { auto actor_entry = actor_registry_.find(actor_id); RAY_CHECK(actor_entry != actor_registry_.end()); @@ -2635,7 +2619,6 @@ void NodeManager::ForwardTask( client->ForwardTask(request, [this, on_error, task, task_id, node_id]( Status status, const rpc::ForwardTaskReply &reply) { - // Remove the FORWARDING task from the SWAP queue. if (status.ok()) { const auto &spec = task.GetTaskSpecification(); // Mark as forwarded so that the task and its lineage are not @@ -2669,21 +2652,22 @@ void NodeManager::ForwardTask( }); } -void NodeManager::FinishAssignTask(const TaskID &task_id, Worker &worker, bool success) { - // Remove the ASSIGNED task from the SWAP queue. +void NodeManager::FinishAssignTask(const std::shared_ptr &worker, + const TaskID &task_id, bool success) { + // Remove the ASSIGNED task from the READY queue. Task assigned_task; TaskState state; if (!local_queues_.RemoveTask(task_id, &assigned_task, &state)) { + // TODO(edoakes): should we be failing silently here? return; } - - RAY_CHECK(state == TaskState::SWAP); + RAY_CHECK(state == TaskState::READY); if (success) { auto spec = assigned_task.GetTaskSpecification(); // We successfully assigned the task to the worker. - worker.AssignTaskId(spec.TaskId()); - worker.AssignJobId(spec.JobId()); + worker->AssignTaskId(spec.TaskId()); + worker->AssignJobId(spec.JobId()); // 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 @@ -2698,7 +2682,7 @@ void NodeManager::FinishAssignTask(const TaskID &task_id, Worker &worker, bool s } else { RAY_LOG(WARNING) << "Failed to send task to worker, disconnecting client"; // We failed to send the task to the worker, so disconnect the worker. - ProcessDisconnectClientMessage(worker.Connection()); + ProcessDisconnectClientMessage(worker->Connection()); // Queue this task for future assignment. We need to do this since // DispatchTasks() removed it from the ready queue. The task will be // assigned to a worker once one becomes available. diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h index ad4416c1f..7e443d97a 100644 --- a/src/ray/raylet/node_manager.h +++ b/src/ray/raylet/node_manager.h @@ -208,12 +208,13 @@ class NodeManager : public rpc::NodeManagerServiceHandler { /// \return Void. void SubmitTask(const Task &task, const Lineage &uncommitted_lineage, bool forwarded = false); - /// Assign a task. The task is assumed to not be queued in local_queues_. + /// Assign a task to a worker. The task is assumed to not be queued in local_queues_. /// - /// \param task The task in question. - /// \param post_assign_callbacks Set of functions to run after assignments finish. - /// \return true, if tasks was assigned to a worker, false otherwise. - bool AssignTask(const Task &task, + /// \param[in] worker The worker to assign the task to. + /// \param[in] task The task in question. + /// \param[out] post_assign_callbacks Vector of callbacks that will be appended + /// to with any logic that should run after the DispatchTasks loop runs. + void AssignTask(const std::shared_ptr &worker, const Task &task, std::vector> *post_assign_callbacks); /// Handle a worker finishing its assigned task. /// @@ -510,11 +511,12 @@ class NodeManager : public rpc::NodeManagerServiceHandler { /// Finish assigning a task to a worker. /// + /// \param worker Worker that the task is assigned to. /// \param task_id Id of the task. - /// \param worker Worker which the task is assigned to. - /// \param success Whether the task is successfully assigned to the worker. + /// \param success Whether or not assigning the task was successful. /// \return void. - void FinishAssignTask(const TaskID &task_id, Worker &worker, bool success); + void FinishAssignTask(const std::shared_ptr &worker, const TaskID &task_id, + bool success); /// Handle a `WorkerLease` request. void HandleWorkerLeaseRequest(const rpc::WorkerLeaseRequest &request, diff --git a/src/ray/raylet/worker.cc b/src/ray/raylet/worker.cc index 9c46180a7..89288366a 100644 --- a/src/ray/raylet/worker.cc +++ b/src/ray/raylet/worker.cc @@ -128,8 +128,7 @@ void Worker::SetActiveObjectIds(const std::unordered_set &&object_ids) active_object_ids_ = object_ids; } -void Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set, - const std::function finish_assign_callback) { +Status Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set) { RAY_CHECK(port_ > 0); rpc::AssignTaskRequest request; request.mutable_task()->mutable_task_spec()->CopyFrom( @@ -138,8 +137,8 @@ void Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set, task.GetTaskExecutionSpec().GetMessage()); request.set_resource_ids(resource_id_set.Serialize()); - auto status = rpc_client_->AssignTask(request, [](Status status, - const rpc::AssignTaskReply &reply) { + return rpc_client_->AssignTask(request, [](Status status, + const rpc::AssignTaskReply &reply) { if (!status.ok()) { RAY_LOG(DEBUG) << "Worker failed to finish executing task: " << status.ToString(); } @@ -147,14 +146,6 @@ void Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set, // and assigning new task will be done when raylet receives // `TaskDone` message. }); - finish_assign_callback(status); - if (!status.ok()) { - RAY_LOG(ERROR) << "Failed to assign task " << task.GetTaskSpecification().TaskId() - << " to worker " << worker_id_; - } else { - RAY_LOG(DEBUG) << "Assigned task " << task.GetTaskSpecification().TaskId() - << " to worker " << worker_id_; - } } void Worker::DirectActorCallArgWaitComplete(int64_t tag) { diff --git a/src/ray/raylet/worker.h b/src/ray/raylet/worker.h index a162dc710..ac186f990 100644 --- a/src/ray/raylet/worker.h +++ b/src/ray/raylet/worker.h @@ -64,8 +64,7 @@ class Worker { const std::unordered_set &GetActiveObjectIds() const; void SetActiveObjectIds(const std::unordered_set &&object_ids); - void AssignTask(const Task &task, const ResourceIdSet &resource_id_set, - const std::function finish_assign_callback); + Status AssignTask(const Task &task, const ResourceIdSet &resource_id_set); void DirectActorCallArgWaitComplete(int64_t tag); void WorkerLeaseGranted(const std::string &address, int port);