[NewScheduler] Pass test_basic.py (#10059)

* .

* .

* Cleanup

* .

* whoops

* Update src/ray/raylet/scheduling/cluster_task_manager.h

Co-authored-by: Edward Oakes <ed.nmi.oakes@gmail.com>

* Update src/ray/raylet/scheduling/cluster_task_manager.h

Co-authored-by: Stephanie Wang <swang@cs.berkeley.edu>

* CR

* .

* .

* done

* .

* Unit tests

Co-authored-by: Edward Oakes <ed.nmi.oakes@gmail.com>
Co-authored-by: Stephanie Wang <swang@cs.berkeley.edu>

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

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<TaskID> ready_ids) {
   }
 }
 
+void ClusterTaskManager::HandleTaskFinished(std::shared_ptr<WorkerInterface> 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<WorkerInterface> worker,
-    std::unordered_map<WorkerID, std::shared_ptr<WorkerInterface>> &leased_workers_,
+    std::unordered_map<WorkerID, std::shared_ptr<WorkerInterface>> &leased_workers,
     const TaskSpecification &task_spec, rpc::RequestWorkerLeaseReply *reply,
     std::function<void(void)> 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<TaskResourceInstances> allocated_resources;
   if (task_spec.IsActorCreationTask()) {
     allocated_resources = worker->GetLifetimeAllocatedInstances();
@@ -206,6 +255,8 @@ void ClusterTaskManager::Dispatch(
       }
     }
   }
+
+  // Send the result back.
   send_reply_callback();
 }
 

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<boost::optional<rpc::GcsNodeInfo>(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<TaskID> 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<WorkerInterface> 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<ClusterResourceScheduler> cluster_resource_scheduler_;

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<double> &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<double> &GetBorrowedCPUInstances() {
-    RAY_CHECK(false) << "Method unused";
-    auto *t = new std::vector<double>();
-    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<TaskResourceInstances> allocated_instances_;
   std::shared_ptr<TaskResourceInstances> lifetime_allocated_instances_;
+  std::vector<double> borrowed_cpu_instances_;
 };
 
 std::shared_ptr<ClusterResourceScheduler> 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<MockWorker> worker =
+    std::make_shared<MockWorker>(WorkerID::FromRandom(), 1234);
+  pool_.PushWorker(std::dynamic_pointer_cast<WorkerInterface>(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