From 136c8ff19e75b55e6bcc122cf033577dd9014707 Mon Sep 17 00:00:00 2001 From: Alex Wu Date: Fri, 21 Aug 2020 15:00:08 -0700 Subject: [PATCH] [NewScheduler] Pass test_basic.py (#10059) * . * . * Cleanup * . * whoops * Update src/ray/raylet/scheduling/cluster_task_manager.h Co-authored-by: Edward Oakes * Update src/ray/raylet/scheduling/cluster_task_manager.h Co-authored-by: Stephanie Wang * CR * . * . * done * . * Unit tests Co-authored-by: Edward Oakes Co-authored-by: Stephanie Wang --- src/ray/raylet/node_manager.cc | 44 ++++++++----- .../raylet/scheduling/cluster_task_manager.cc | 59 ++++++++++++++++-- .../raylet/scheduling/cluster_task_manager.h | 22 ++++++- .../scheduling/cluster_task_manager_test.cc | 62 ++++++++++++++++--- 4 files changed, 158 insertions(+), 29 deletions(-) diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index e5b2b659b..6867b5fd2 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -1911,10 +1911,7 @@ void NodeManager::HandleReturnWorker(const rpc::ReturnWorkerRequest &request, HandleDirectCallTaskUnblocked(worker); } if (new_scheduler_enabled_) { - new_resource_scheduler_->SubtractCPUResourceInstances( - worker->GetBorrowedCPUInstances()); - new_resource_scheduler_->FreeLocalTaskResources(worker->GetAllocatedInstances()); - worker->ClearAllocatedInstances(); + cluster_task_manager_->HandleTaskFinished(worker); } HandleWorkerAvailable(worker); } @@ -1955,22 +1952,37 @@ void NodeManager::HandleCancelWorkerLease(const rpc::CancelWorkerLeaseRequest &r const TaskID task_id = TaskID::FromBinary(request.task_id()); Task removed_task; TaskState removed_task_state; - const auto canceled = - local_queues_.RemoveTask(task_id, &removed_task, &removed_task_state); - if (!canceled) { - // We do not have the task. This could be because we haven't received the - // lease request yet, or because we already granted the lease request and - // it has already been returned. - } else { - if (removed_task.OnDispatch()) { + bool canceled; + if (new_scheduler_enabled_) { + canceled = cluster_task_manager_->CancelTask(task_id); + if (canceled) { // We have not yet granted the worker lease. Cancel it now. - removed_task.OnCancellation()(); task_dependency_manager_.TaskCanceled(task_id); task_dependency_manager_.UnsubscribeGetDependencies(task_id); } else { - // We already granted the worker lease and sent the reply. Re-queue the - // task and wait for the requester to return the leased worker. - local_queues_.QueueTasks({removed_task}, removed_task_state); + // There are 2 cases here. + // 1. We haven't received the lease request yet. It's the caller's job to + // retry the cancellation once we've received the request. + // 2. We have already granted the lease. The caller is now responsible + // for returning the lease, not cancelling it. + } + } else { + canceled = local_queues_.RemoveTask(task_id, &removed_task, &removed_task_state); + if (!canceled) { + // We do not have the task. This could be because we haven't received the + // lease request yet, or because we already granted the lease request and + // it has already been returned. + } else { + if (removed_task.OnDispatch()) { + // We have not yet granted the worker lease. Cancel it now. + removed_task.OnCancellation()(); + task_dependency_manager_.TaskCanceled(task_id); + task_dependency_manager_.UnsubscribeGetDependencies(task_id); + } else { + // We already granted the worker lease and sent the reply. Re-queue the + // task and wait for the requester to return the leased worker. + local_queues_.QueueTasks({removed_task}, removed_task_state); + } } } // The task cancellation failed if we did not have the task queued, since diff --git a/src/ray/raylet/scheduling/cluster_task_manager.cc b/src/ray/raylet/scheduling/cluster_task_manager.cc index baa467ae8..68ce8d8f7 100644 --- a/src/ray/raylet/scheduling/cluster_task_manager.cc +++ b/src/ray/raylet/scheduling/cluster_task_manager.cc @@ -41,7 +41,10 @@ bool ClusterTaskManager::SchedulePendingTasks() { continue; } else { if (node_id_string == self_node_id_.Binary()) { - did_schedule = did_schedule || WaitForTaskArgsRequests(work); + // Warning: WaitForTaskArgsRequests must execute (do not let it short + // circuit if did_schedule is true). + bool task_scheduled = WaitForTaskArgsRequests(work); + did_schedule = task_scheduled || did_schedule; } else { // Should spill over to a different node. cluster_resource_scheduler_->AllocateRemoteTaskResources(node_id_string, @@ -151,17 +154,63 @@ void ClusterTaskManager::TasksUnblocked(const std::vector ready_ids) { } } +void ClusterTaskManager::HandleTaskFinished(std::shared_ptr worker) { + cluster_resource_scheduler_->SubtractCPUResourceInstances( + worker->GetBorrowedCPUInstances()); + cluster_resource_scheduler_->FreeLocalTaskResources(worker->GetAllocatedInstances()); + worker->ClearAllocatedInstances(); +} + +bool ClusterTaskManager::CancelTask(const TaskID &task_id) { + for (auto iter = tasks_to_schedule_.begin(); iter != tasks_to_schedule_.end(); iter++) { + if (std::get<0>(*iter).GetTaskSpecification().TaskId() == task_id) { + tasks_to_schedule_.erase(iter); + return true; + } + } + for (auto iter = tasks_to_dispatch_.begin(); iter != tasks_to_dispatch_.end(); iter++) { + if (std::get<0>(*iter).GetTaskSpecification().TaskId() == task_id) { + tasks_to_dispatch_.erase(iter); + return true; + } + } + + auto iter = waiting_tasks_.find(task_id); + if (iter != waiting_tasks_.end()) { + waiting_tasks_.erase(iter); + return true; + } + + return false; +} + +std::string ClusterTaskManager::DebugString() { + std::stringstream buffer; + buffer << "========== Node: " << self_node_id_ << " =================\n"; + buffer << "Schedule queue length: " << tasks_to_schedule_.size() << "\n"; + buffer << "Dispatch queue length: " << tasks_to_dispatch_.size() << "\n"; + buffer << "Waiting tasks size: " << waiting_tasks_.size() << "\n"; + buffer << "cluster_resource_scheduler state: " + << cluster_resource_scheduler_->DebugString() << "\n"; + buffer << "=================================================="; + return buffer.str(); +} + void ClusterTaskManager::Dispatch( std::shared_ptr worker, - std::unordered_map> &leased_workers_, + std::unordered_map> &leased_workers, const TaskSpecification &task_spec, rpc::RequestWorkerLeaseReply *reply, std::function send_reply_callback) { + // Pass the contact info of the worker to use. reply->mutable_worker_address()->set_ip_address(worker->IpAddress()); reply->mutable_worker_address()->set_port(worker->Port()); reply->mutable_worker_address()->set_worker_id(worker->WorkerId().Binary()); reply->mutable_worker_address()->set_raylet_id(self_node_id_.Binary()); - RAY_CHECK(leased_workers_.find(worker->WorkerId()) == leased_workers_.end()); - leased_workers_[worker->WorkerId()] = worker; + + RAY_CHECK(leased_workers.find(worker->WorkerId()) == leased_workers.end()); + leased_workers[worker->WorkerId()] = worker; + + // Update our internal view of the cluster state. std::shared_ptr allocated_resources; if (task_spec.IsActorCreationTask()) { allocated_resources = worker->GetLifetimeAllocatedInstances(); @@ -206,6 +255,8 @@ void ClusterTaskManager::Dispatch( } } } + + // Send the result back. send_reply_callback(); } diff --git a/src/ray/raylet/scheduling/cluster_task_manager.h b/src/ray/raylet/scheduling/cluster_task_manager.h index a173de753..a53764f6d 100644 --- a/src/ray/raylet/scheduling/cluster_task_manager.h +++ b/src/ray/raylet/scheduling/cluster_task_manager.h @@ -1,5 +1,8 @@ #pragma once +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" + #include "ray/common/task/task.h" #include "ray/common/task/task_common.h" #include "ray/raylet/scheduling/cluster_resource_scheduler.h" @@ -29,9 +32,11 @@ typedef std::function(const ClientID &node_id) /// 3. If a task has unresolved dependencies, set it aside to wait for /// dependencies to be resolved. /// 4. When a task is ready to be dispatched, ensure that the local node is -/// still capable of running the task. +/// still capable of running the task, then dispatch it. /// * Step 4 should be run any time there is a new task to dispatch *or* /// there is a new worker which can dispatch the tasks. +/// 5. When a worker finishes executing its task(s), the requester will return +/// it and we should release the resources in our view of the node's state. class ClusterTaskManager { public: /// fullfills_dependencies_func Should return if all dependencies are @@ -79,6 +84,21 @@ class ClusterTaskManager { /// \param readyIds: The tasks which are now ready to be dispatched. void TasksUnblocked(const std::vector ready_ids); + /// (Step 5) Call once a task finishes (i.e. a worker is returned). + /// + /// \param worker: The worker which was running the task. + void HandleTaskFinished(std::shared_ptr worker); + + /// Attempt to cancel an already queued task. + /// + /// \param task_id: The id of the task to remove. + /// + /// \return True if task was successfully removed. This function will return + /// false if the task is already running. + bool CancelTask(const TaskID &task_id); + + std::string DebugString(); + private: const ClientID &self_node_id_; std::shared_ptr cluster_resource_scheduler_; diff --git a/src/ray/raylet/scheduling/cluster_task_manager_test.cc b/src/ray/raylet/scheduling/cluster_task_manager_test.cc index 031cfa335..1a9e39a9b 100644 --- a/src/ray/raylet/scheduling/cluster_task_manager_test.cc +++ b/src/ray/raylet/scheduling/cluster_task_manager_test.cc @@ -199,12 +199,14 @@ class MockWorker : public WorkerInterface { RAY_CHECK(false) << "Method unused"; } - void ClearAllocatedInstances() { RAY_CHECK(false) << "Method unused"; } + void ClearAllocatedInstances() { + allocated_instances_ = nullptr; + } void ClearLifetimeAllocatedInstances() { RAY_CHECK(false) << "Method unused"; } void SetBorrowedCPUInstances(std::vector &cpu_instances) { - RAY_CHECK(false) << "Method unused"; + borrowed_cpu_instances_ = cpu_instances; } const PlacementGroupID &GetPlacementGroupId() const { @@ -217,9 +219,7 @@ class MockWorker : public WorkerInterface { } std::vector &GetBorrowedCPUInstances() { - RAY_CHECK(false) << "Method unused"; - auto *t = new std::vector(); - return *t; + return borrowed_cpu_instances_; } void ClearBorrowedCPUInstances() { RAY_CHECK(false) << "Method unused"; } @@ -246,6 +246,7 @@ class MockWorker : public WorkerInterface { rpc::Address address_; std::shared_ptr allocated_instances_; std::shared_ptr lifetime_allocated_instances_; + std::vector borrowed_cpu_instances_; }; std::shared_ptr CreateSingleNodeScheduler( @@ -428,10 +429,8 @@ TEST_F(ClusterTaskManagerTest, ResourceTakenWhileResolving) { ASSERT_EQ(pool_.workers.size(), 1); /* Second task finishes, making space for the original task */ - single_node_resource_scheduler_->FreeLocalTaskResources( - worker->GetAllocatedInstances()); - // single_node_resource_scheduler_->UpdateLocalAvailableResourcesFromResourceInstances(); leased_workers_.clear(); + task_manager_.HandleTaskFinished(worker); task_manager_.DispatchScheduledTasksToWorkers(pool_, leased_workers_); @@ -441,6 +440,53 @@ TEST_F(ClusterTaskManagerTest, ResourceTakenWhileResolving) { ASSERT_EQ(pool_.workers.size(), 0); } +TEST_F(ClusterTaskManagerTest, TaskCancellationTest) { + std::shared_ptr worker = + std::make_shared(WorkerID::FromRandom(), 1234); + pool_.PushWorker(std::dynamic_pointer_cast(worker)); + + Task task = CreateTask({{ray::kCPU_ResourceLabel, 1}}); + rpc::RequestWorkerLeaseReply reply; + + bool callback_called = false; + bool *callback_called_ptr = &callback_called; + auto callback = [callback_called_ptr]() { *callback_called_ptr = true; }; + + // Task not queued so we can't cancel it. + ASSERT_FALSE(task_manager_.CancelTask(task.GetTaskSpecification().TaskId())); + + task_manager_.QueueTask(task, &reply, callback); + + // Task is now queued so cancellation works. + ASSERT_TRUE(task_manager_.CancelTask(task.GetTaskSpecification().TaskId())); + task_manager_.DispatchScheduledTasksToWorkers(pool_, leased_workers_); + // Task will not execute. + ASSERT_FALSE(callback_called); + ASSERT_EQ(leased_workers_.size(), 0); + ASSERT_EQ(pool_.workers.size(), 1); + + task_manager_.QueueTask(task, &reply, callback); + task_manager_.SchedulePendingTasks(); + + // We can still cancel the task if it's on the dispatch queue. + ASSERT_TRUE(task_manager_.CancelTask(task.GetTaskSpecification().TaskId())); + // Task will not execute. + ASSERT_FALSE(callback_called); + ASSERT_EQ(leased_workers_.size(), 0); + ASSERT_EQ(pool_.workers.size(), 1); + + task_manager_.QueueTask(task, &reply, callback); + task_manager_.SchedulePendingTasks(); + task_manager_.DispatchScheduledTasksToWorkers(pool_, leased_workers_); + + // Task is now running so we can't cancel it. + ASSERT_FALSE(task_manager_.CancelTask(task.GetTaskSpecification().TaskId())); + // Task will not execute. + ASSERT_TRUE(callback_called); + ASSERT_EQ(leased_workers_.size(), 1); + ASSERT_EQ(pool_.workers.size(), 0); +} + } // namespace raylet } // namespace ray