[New Scheduler] Don't start tasks if the owner is dead (#12050)

python/ray/tests/test_actor_advanced.py

@@ -12,7 +12,7 @@ import ray
 import ray.test_utils
 import ray.cluster_utils
 from ray.test_utils import (run_string_as_driver, get_non_head_nodes,
-                            wait_for_condition, new_scheduler_enabled)
+                            wait_for_condition)
 from ray.experimental.internal_kv import _internal_kv_get, _internal_kv_put
 
 
@@ -943,7 +943,6 @@ def test_actor_creation_task_crash(ray_start_regular):
         }
     }],
     indirect=True)
-@pytest.mark.skipif(new_scheduler_enabled(), reason="todo hangs")
 def test_pending_actor_removed_by_owner(ray_start_regular):
     # Verify when an owner of pending actors is killed, the actor resources
     # are correctly returned.

src/ray/raylet/node_manager.cc

@@ -218,9 +218,14 @@ NodeManager::NodeManager(boost::asio::io_service &io_service, const NodeID &self
     auto get_node_info_func = [this](const NodeID &node_id) {
       return gcs_client_->Nodes().Get(node_id);
     };
-    cluster_task_manager_ = std::shared_ptr<ClusterTaskManager>(
-        new ClusterTaskManager(self_node_id_, new_resource_scheduler_,
-                               fulfills_dependencies_func, get_node_info_func));
+    auto is_owner_alive = [this](const WorkerID &owner_worker_id,
+                                 const NodeID &owner_node_id) {
+      return !(failed_workers_cache_.count(owner_worker_id) > 0 ||
+               failed_nodes_cache_.count(owner_node_id) > 0);
+    };
+    cluster_task_manager_ = std::shared_ptr<ClusterTaskManager>(new ClusterTaskManager(
+        self_node_id_, new_resource_scheduler_, fulfills_dependencies_func,
+        is_owner_alive, get_node_info_func));
   }
 
   RAY_CHECK_OK(store_client_.Connect(config.store_socket_name.c_str()));
@@ -794,6 +799,8 @@ void NodeManager::HandleUnexpectedWorkerFailure(const rpc::Address &address) {
     RAY_LOG(DEBUG) << "Lease " << worker->WorkerId() << " owned by " << owner_worker_id;
     RAY_CHECK(!owner_worker_id.IsNil() && !owner_node_id.IsNil());
     if (!worker->IsDetachedActor()) {
+      // TODO (Alex): Cancel all pending child tasks of the tasks whose owners have failed
+      // because the owner could've submitted lease requests before failing.
       if (!worker_id.IsNil()) {
         // If the failed worker was a leased worker's owner, then kill the leased worker.
         if (owner_worker_id == worker_id) {
@@ -1314,7 +1321,7 @@ void NodeManager::HandleWorkerAvailable(const std::shared_ptr<WorkerInterface> &
 
   // If the worker was assigned a task, mark it as finished.
   if (!worker->GetAssignedTaskId().IsNil()) {
-    worker_idle = FinishAssignedTask(*worker);
+    worker_idle = FinishAssignedTask(worker);
   }
 
   if (worker_idle) {
@@ -2541,7 +2548,10 @@ void NodeManager::AssignTask(const std::shared_ptr<WorkerInterface> &worker,
   });
 }
 
-bool NodeManager::FinishAssignedTask(WorkerInterface &worker) {
+bool NodeManager::FinishAssignedTask(const std::shared_ptr<WorkerInterface> &worker_ptr) {
+  // TODO (Alex): We should standardize to pass
+  // std::shared_ptr<WorkerInterface> instead of refs.
+  auto &worker = *worker_ptr;
   TaskID task_id = worker.GetAssignedTaskId();
   RAY_LOG(DEBUG) << "Finished task " << task_id;
 
@@ -2550,8 +2560,7 @@ bool NodeManager::FinishAssignedTask(WorkerInterface &worker) {
     task = worker.GetAssignedTask();
     // leased_workers_.erase(worker.WorkerId()); // Maybe RAY_CHECK ?
     if (worker.GetAllocatedInstances() != nullptr) {
-      new_resource_scheduler_->FreeLocalTaskResources(worker.GetAllocatedInstances());
-      worker.ClearAllocatedInstances();
+      cluster_task_manager_->HandleTaskFinished(worker_ptr);
     }
   } else {
     // (See design_docs/task_states.rst for the state transition diagram.)

src/ray/raylet/node_manager.h

@@ -284,7 +284,8 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
   /// \return Whether the worker should be returned to the idle pool. This is
   /// only false for direct actor creation calls, which should never be
   /// returned to idle.
-  bool FinishAssignedTask(WorkerInterface &worker);
+  bool FinishAssignedTask(const std::shared_ptr<WorkerInterface> &worker_ptr);
+
   /// Helper function to produce actor table data for a newly created actor.
   ///
   /// \param task_spec Task specification of the actor creation task that created the

src/ray/raylet/scheduling/cluster_task_manager.cc

@@ -1,6 +1,7 @@
+#include "ray/raylet/scheduling/cluster_task_manager.h"
+
 #include <google/protobuf/map.h>
 
-#include "ray/raylet/scheduling/cluster_task_manager.h"
 #include "ray/util/logging.h"
 
 namespace ray {
@@ -10,10 +11,12 @@ ClusterTaskManager::ClusterTaskManager(
     const NodeID &self_node_id,
     std::shared_ptr<ClusterResourceScheduler> cluster_resource_scheduler,
     std::function<bool(const Task &)> fulfills_dependencies_func,
+    std::function<bool(const WorkerID &, const NodeID &)> is_owner_alive,
     NodeInfoGetter get_node_info)
     : self_node_id_(self_node_id),
       cluster_resource_scheduler_(cluster_resource_scheduler),
       fulfills_dependencies_func_(fulfills_dependencies_func),
+      is_owner_alive_(is_owner_alive),
       get_node_info_(get_node_info) {}
 
 bool ClusterTaskManager::SchedulePendingTasks() {
@@ -102,7 +105,9 @@ void ClusterTaskManager::DispatchScheduledTasksToWorkers(
        shapes_it != tasks_to_dispatch_.end();) {
     auto &dispatch_queue = shapes_it->second;
     for (auto work_it = dispatch_queue.begin(); work_it != dispatch_queue.end();) {
-      auto spec = std::get<0>(*work_it).GetTaskSpecification();
+      auto &work = *work_it;
+      auto &task = std::get<0>(work);
+      auto &spec = task.GetTaskSpecification();
 
       std::shared_ptr<WorkerInterface> worker = worker_pool.PopWorker(spec);
       if (!worker) {
@@ -110,20 +115,30 @@ void ClusterTaskManager::DispatchScheduledTasksToWorkers(
         return;
       }
 
-      bool worker_leased;
-      bool remove = AttemptDispatchWork(*work_it, worker, &worker_leased);
-      if (worker_leased) {
-        auto reply = std::get<1>(*work_it);
-        auto callback = std::get<2>(*work_it);
-        Dispatch(worker, leased_workers, spec, reply, callback);
-      } else {
+      const auto owner_worker_id = WorkerID::FromBinary(spec.CallerAddress().worker_id());
+      const auto owner_node_id = NodeID::FromBinary(spec.CallerAddress().raylet_id());
+      // If the owner has died since this task was queued, cancel the task by
+      // killing the worker (unless this task is for a detached actor).
+      if (!spec.IsDetachedActor() && !is_owner_alive_(owner_worker_id, owner_node_id)) {
+        RAY_LOG(WARNING) << "Task: " << task.GetTaskSpecification().TaskId()
+                         << "'s caller is no longer running. Cancelling task.";
         worker_pool.PushWorker(worker);
-      }
-
-      if (remove) {
         work_it = dispatch_queue.erase(work_it);
       } else {
-        break;
+        bool worker_leased;
+        bool remove = AttemptDispatchWork(*work_it, worker, &worker_leased);
+        if (worker_leased) {
+          auto reply = std::get<1>(*work_it);
+          auto callback = std::get<2>(*work_it);
+          Dispatch(worker, leased_workers, spec, reply, callback);
+        } else {
+          worker_pool.PushWorker(worker);
+        }
+        if (remove) {
+          work_it = dispatch_queue.erase(work_it);
+        } else {
+          break;
+        }
       }
     }
     if (dispatch_queue.empty()) {
@@ -203,6 +218,9 @@ void ClusterTaskManager::TasksUnblocked(const std::vector<TaskID> ready_ids) {
 void ClusterTaskManager::HandleTaskFinished(std::shared_ptr<WorkerInterface> worker) {
   cluster_resource_scheduler_->FreeLocalTaskResources(worker->GetAllocatedInstances());
   worker->ClearAllocatedInstances();
+  cluster_resource_scheduler_->FreeLocalTaskResources(
+      worker->GetLifetimeAllocatedInstances());
+  worker->ClearLifetimeAllocatedInstances();
 }
 
 void ReplyCancelled(Work &work) {
@@ -311,7 +329,7 @@ void ClusterTaskManager::Heartbeat(bool light_heartbeat_enabled,
   }
 }
 
-std::string ClusterTaskManager::DebugString() {
+std::string ClusterTaskManager::DebugString() const {
   std::stringstream buffer;
   buffer << "========== Node: " << self_node_id_ << " =================\n";
   buffer << "Schedule queue length: " << tasks_to_schedule_.size() << "\n";

src/ray/raylet/scheduling/cluster_task_manager.h

@@ -48,10 +48,12 @@ class ClusterTaskManager {
   /// the state of the cluster.
   /// \param fulfills_dependencies_func: Returns true if all of a task's
   /// dependencies are fulfilled.
-  /// \param gcs_client: A gcs client.
+  /// \param is_owner_alive: A callback which returns if the owner process is alive
+  /// (according to our ownership model). \param gcs_client: A gcs client.
   ClusterTaskManager(const NodeID &self_node_id,
                      std::shared_ptr<ClusterResourceScheduler> cluster_resource_scheduler,
                      std::function<bool(const Task &)> fulfills_dependencies_func,
+                     std::function<bool(const WorkerID &, const NodeID &)> is_owner_alive,
                      NodeInfoGetter get_node_info);
 
   /// (Step 2) For each task in tasks_to_schedule_, pick a node in the system
@@ -111,7 +113,7 @@ class ClusterTaskManager {
   void Heartbeat(bool light_heartbeat_enabled,
                  std::shared_ptr<HeartbeatTableData> data) const;
 
-  std::string DebugString();
+  std::string DebugString() const;
 
  private:
   /// Helper method to try dispatching a single task from the queue to an
@@ -123,6 +125,7 @@ class ClusterTaskManager {
   const NodeID &self_node_id_;
   std::shared_ptr<ClusterResourceScheduler> cluster_resource_scheduler_;
   std::function<bool(const Task &)> fulfills_dependencies_func_;
+  std::function<bool(const WorkerID &, const NodeID &)> is_owner_alive_;
   NodeInfoGetter get_node_info_;
 
   // TODO (Alex): Implement fair queuing for these queues

src/ray/raylet/scheduling/cluster_task_manager_test.cc

@@ -98,12 +98,16 @@ class ClusterTaskManagerTest : public ::testing::Test {
         scheduler_(CreateSingleNodeScheduler(id_.Binary())),
         fulfills_dependencies_calls_(0),
         dependencies_fulfilled_(true),
+        is_owner_alive_(true),
         node_info_calls_(0),
         task_manager_(id_, scheduler_,
                       [this](const Task &_task) {
                         fulfills_dependencies_calls_++;
                         return dependencies_fulfilled_;
                       },
+                      [this](const WorkerID &worker_id, const NodeID &node_id) {
+                        return is_owner_alive_;
+                      },
                       [this](const NodeID &node_id) {
                         node_info_calls_++;
                         return node_info_[node_id];
@@ -133,6 +137,8 @@ class ClusterTaskManagerTest : public ::testing::Test {
   int fulfills_dependencies_calls_;
   bool dependencies_fulfilled_;
 
+  bool is_owner_alive_;
+
   int node_info_calls_;
   std::unordered_map<NodeID, boost::optional<rpc::GcsNodeInfo>> node_info_;
 
@@ -473,6 +479,39 @@ TEST_F(ClusterTaskManagerTest, HeartbeatTest) {
   }
 }
 
+TEST_F(ClusterTaskManagerTest, OwnerDeadTest) {
+  /*
+    Test the race condition in which the owner of a task dies while the task is pending.
+    This is the essence of test_actor_advanced.py::test_pending_actor_removed_by_owner
+   */
+  Task task = CreateTask({{ray::kCPU_ResourceLabel, 4}});
+  rpc::RequestWorkerLeaseReply reply;
+  bool callback_occurred = false;
+  bool *callback_occurred_ptr = &callback_occurred;
+  auto callback = [callback_occurred_ptr]() { *callback_occurred_ptr = true; };
+
+  std::shared_ptr<MockWorker> worker =
+      std::make_shared<MockWorker>(WorkerID::FromRandom(), 1234);
+  pool_.PushWorker(std::dynamic_pointer_cast<WorkerInterface>(worker));
+
+  task_manager_.QueueTask(task, &reply, callback);
+  task_manager_.SchedulePendingTasks();
+
+  is_owner_alive_ = false;
+  task_manager_.DispatchScheduledTasksToWorkers(pool_, leased_workers_);
+
+  ASSERT_FALSE(callback_occurred);
+  ASSERT_EQ(leased_workers_.size(), 0);
+  ASSERT_EQ(pool_.workers.size(), 1);
+
+  is_owner_alive_ = true;
+  task_manager_.DispatchScheduledTasksToWorkers(pool_, leased_workers_);
+
+  ASSERT_FALSE(callback_occurred);
+  ASSERT_EQ(leased_workers_.size(), 0);
+  ASSERT_EQ(pool_.workers.size(), 1);
+}
+
 int main(int argc, char **argv) {
   ::testing::InitGoogleTest(&argc, argv);
   return RUN_ALL_TESTS();

src/ray/raylet/test/util.h

@@ -20,7 +20,8 @@ namespace raylet {
 
 class MockWorker : public WorkerInterface {
  public:
-  MockWorker(WorkerID worker_id, int port) : worker_id_(worker_id), port_(port) {}
+  MockWorker(WorkerID worker_id, int port)
+      : worker_id_(worker_id), port_(port), is_detached_actor_(false) {}
 
   WorkerID WorkerId() const { return worker_id_; }
 
@@ -111,11 +112,8 @@ class MockWorker : public WorkerInterface {
     RAY_CHECK(false) << "Method unused";
     return ActorID::Nil();
   }
-  void MarkDetachedActor() { RAY_CHECK(false) << "Method unused"; }
-  bool IsDetachedActor() const {
-    RAY_CHECK(false) << "Method unused";
-    return false;
-  }
+  void MarkDetachedActor() { is_detached_actor_ = true; }
+  bool IsDetachedActor() const { return is_detached_actor_; }
   const std::shared_ptr<ClientConnection> Connection() const {
     RAY_CHECK(false) << "Method unused";
     return nullptr;
@@ -164,7 +162,7 @@ class MockWorker : public WorkerInterface {
 
   void ClearAllocatedInstances() { allocated_instances_ = nullptr; }
 
-  void ClearLifetimeAllocatedInstances() { RAY_CHECK(false) << "Method unused"; }
+  void ClearLifetimeAllocatedInstances() { lifetime_allocated_instances_ = nullptr; }
 
   void SetBorrowedCPUInstances(std::vector<double> &cpu_instances) {
     borrowed_cpu_instances_ = cpu_instances;
@@ -206,6 +204,7 @@ class MockWorker : public WorkerInterface {
   std::shared_ptr<TaskResourceInstances> allocated_instances_;
   std::shared_ptr<TaskResourceInstances> lifetime_allocated_instances_;
   std::vector<double> borrowed_cpu_instances_;
+  bool is_detached_actor_;
 };
 
 }  // namespace raylet