Don't add assigned tasks to SWAP queue (#6325)

This commit is contained in:
Edward Oakes
2019-12-02 16:39:02 -08:00
committed by Ed Oakes
parent b17b455bed
commit dbfe752c39
4 changed files with 99 additions and 123 deletions
+85 -101
View File
@@ -813,8 +813,6 @@ std::unordered_map<SchedulingClass, ordered_set<TaskID>> MakeTasksByClass(
void NodeManager::DispatchTasks(
const std::unordered_map<SchedulingClass, ordered_set<TaskID>> &tasks_by_class) {
std::unordered_set<TaskID> 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 =
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> &worker, const Task &task,
std::vector<std::function<void()>> *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 = 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=<max_memory_bytes>).";
}
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=<max_memory_bytes>).";
}
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> &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.
+10 -8
View File
@@ -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> &worker, const Task &task,
std::vector<std::function<void()>> *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> &worker, const TaskID &task_id,
bool success);
/// Handle a `WorkerLease` request.
void HandleWorkerLeaseRequest(const rpc::WorkerLeaseRequest &request,
+3 -12
View File
@@ -128,8 +128,7 @@ void Worker::SetActiveObjectIds(const std::unordered_set<ObjectID> &&object_ids)
active_object_ids_ = object_ids;
}
void Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set,
const std::function<void(Status)> 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) {
+1 -2
View File
@@ -64,8 +64,7 @@ class Worker {
const std::unordered_set<ObjectID> &GetActiveObjectIds() const;
void SetActiveObjectIds(const std::unordered_set<ObjectID> &&object_ids);
void AssignTask(const Task &task, const ResourceIdSet &resource_id_set,
const std::function<void(Status)> 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);