From 4182b856118a38295e76db37c752e0a5c68352d8 Mon Sep 17 00:00:00 2001 From: Philipp Moritz Date: Tue, 6 Nov 2018 21:23:31 -0800 Subject: [PATCH] Cache resources in SchedulingQueue (#3232) * cache resources * fix * documentation and remove old code * fix PR * update documentation * linting --- src/ray/raylet/scheduling_queue.cc | 45 +++++++++++------------------- src/ray/raylet/scheduling_queue.h | 21 +++++++------- 2 files changed, 27 insertions(+), 39 deletions(-) diff --git a/src/ray/raylet/scheduling_queue.cc b/src/ray/raylet/scheduling_queue.cc index 7943fc2e3..049f772a4 100644 --- a/src/ray/raylet/scheduling_queue.cc +++ b/src/ray/raylet/scheduling_queue.cc @@ -10,7 +10,7 @@ void RemoveTasksFromQueue(ray::raylet::SchedulingQueue::TaskQueue &queue, std::unordered_set &task_ids, std::vector &removed_tasks) { for (auto it = task_ids.begin(); it != task_ids.end();) { - if (queue.RemoveTask(*it, removed_tasks)) { + if (queue.RemoveTask(*it, &removed_tasks)) { it = task_ids.erase(it); } else { it++; @@ -80,30 +80,25 @@ bool SchedulingQueue::TaskQueue::AppendTask(const TaskID &task_id, const Task &t RAY_CHECK(task_map_.find(task_id) == task_map_.end()); auto list_iterator = task_list_.insert(task_list_.end(), task); task_map_[task_id] = list_iterator; - return true; -} - -bool SchedulingQueue::TaskQueue::RemoveTask(const TaskID &task_id) { - auto task_found_iterator = task_map_.find(task_id); - if (task_found_iterator == task_map_.end()) { - return false; - } - - auto list_iterator = task_found_iterator->second; - task_map_.erase(task_found_iterator); - task_list_.erase(list_iterator); + // Resource bookkeeping + current_resource_load_.AddResources(task.GetTaskSpecification().GetRequiredResources()); return true; } bool SchedulingQueue::TaskQueue::RemoveTask(const TaskID &task_id, - std::vector &removed_tasks) { + std::vector *removed_tasks) { auto task_found_iterator = task_map_.find(task_id); if (task_found_iterator == task_map_.end()) { return false; } auto list_iterator = task_found_iterator->second; - removed_tasks.push_back(std::move(*list_iterator)); + // Resource bookkeeping + current_resource_load_.SubtractResourcesStrict( + list_iterator->GetTaskSpecification().GetRequiredResources()); + if (removed_tasks) { + removed_tasks->push_back(std::move(*list_iterator)); + } task_map_.erase(task_found_iterator); task_list_.erase(list_iterator); return true; @@ -115,6 +110,10 @@ bool SchedulingQueue::TaskQueue::HasTask(const TaskID &task_id) const { const std::list &SchedulingQueue::TaskQueue::GetTasks() const { return task_list_; } +const ResourceSet &SchedulingQueue::TaskQueue::GetCurrentResourceLoad() const { + return current_resource_load_; +} + const std::list &SchedulingQueue::GetMethodsWaitingForActorCreation() const { return this->methods_waiting_for_actor_creation_.GetTasks(); } @@ -135,25 +134,13 @@ const std::list &SchedulingQueue::GetInfeasibleTasks() const { return this->infeasible_tasks_.GetTasks(); } -ResourceSet SchedulingQueue::GetQueueResources(const TaskQueue &task_queue) const { - // Iterate over all tasks of the specified queue and aggregate total resource - // demand in a resource set. - ResourceSet queue_resources; - for (const auto &task : task_queue.GetTasks()) { - queue_resources.AddResources(task.GetTaskSpecification().GetRequiredResources()); - } - return queue_resources; -} - ResourceSet SchedulingQueue::GetReadyQueueResources() const { - return GetQueueResources(ready_tasks_); + return ready_tasks_.GetCurrentResourceLoad(); } ResourceSet SchedulingQueue::GetResourceLoad() const { - ResourceSet load_resource_set; - load_resource_set.AddResources(GetReadyQueueResources()); // TODO(atumanov): consider other types of tasks as part of load. - return load_resource_set; + return ready_tasks_.GetCurrentResourceLoad(); } const std::list &SchedulingQueue::GetRunningTasks() const { diff --git a/src/ray/raylet/scheduling_queue.h b/src/ray/raylet/scheduling_queue.h index d8ecf6ae9..14030426a 100644 --- a/src/ray/raylet/scheduling_queue.h +++ b/src/ray/raylet/scheduling_queue.h @@ -223,9 +223,10 @@ class SchedulingQueue { /// /// \param task_id The task ID for the task to remove from the queue. /// \param removed_tasks If the task specified by task_id is successfully - // removed from the queue, the task data is appended to the vector. + /// removed from the queue, the task data is appended to the vector. Can + /// be a nullptr, in which case nothing is appended. /// \return Whether the removal succeeds. - bool RemoveTask(const TaskID &task_id, std::vector &removed_tasks); + bool RemoveTask(const TaskID &task_id, std::vector *removed_tasks = nullptr); /// \brief Check if the queue contains a specific task id. /// @@ -237,11 +238,17 @@ class SchedulingQueue { /// \return A list of tasks contained in this queue. const std::list &GetTasks() const; + /// \brief Get the total resources required by the tasks in the queue. + /// \return Total resources required by the tasks in the queue. + const ResourceSet &GetCurrentResourceLoad() const; + private: - // A list of tasks. + /// A list of tasks. std::list task_list_; - // A hash to speed up looking up a task. + /// A hash to speed up looking up a task. std::unordered_map::iterator> task_map_; + /// Aggregate resources of all the tasks in this queue. + ResourceSet current_resource_load_; }; private: @@ -265,12 +272,6 @@ class SchedulingQueue { /// The set of currently running driver tasks. These are empty tasks that are /// started by a driver process on initialization. std::unordered_set driver_task_ids_; - - /// \brief Return all resource demand associated with the specified task queue. - /// - /// \param task_queue The task queue for which aggregate resource demand is calculated. - /// \return Aggregate resource demand. - ResourceSet GetQueueResources(const TaskQueue &task_queue) const; }; } // namespace raylet