Cache resources in SchedulingQueue (#3232)

* cache resources

* fix

* documentation and remove old code

* fix PR

* update documentation

* linting

src/ray/raylet/scheduling_queue.cc

@@ -10,7 +10,7 @@ void RemoveTasksFromQueue(ray::raylet::SchedulingQueue::TaskQueue &queue,
                           std::unordered_set<ray::TaskID> &task_ids,
                           std::vector<ray::raylet::Task> &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<Task> &removed_tasks) {
+                                            std::vector<Task> *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<Task> &SchedulingQueue::TaskQueue::GetTasks() const { return task_list_; }
 
+const ResourceSet &SchedulingQueue::TaskQueue::GetCurrentResourceLoad() const {
+  return current_resource_load_;
+}
+
 const std::list<Task> &SchedulingQueue::GetMethodsWaitingForActorCreation() const {
   return this->methods_waiting_for_actor_creation_.GetTasks();
 }
@@ -135,25 +134,13 @@ const std::list<Task> &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<Task> &SchedulingQueue::GetRunningTasks() const {

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<Task> &removed_tasks);
+    bool RemoveTask(const TaskID &task_id, std::vector<Task> *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<Task> &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> task_list_;
-    // A hash to speed up looking up a task.
+    /// A hash to speed up looking up a task.
     std::unordered_map<TaskID, std::list<Task>::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<TaskID> 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