diff --git a/python/ray/tests/test_actor_advanced.py b/python/ray/tests/test_actor_advanced.py index 579acd0d0..b599ec5c1 100644 --- a/python/ray/tests/test_actor_advanced.py +++ b/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. diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index 60705cef6..960691938 100644 --- a/src/ray/raylet/node_manager.cc +++ b/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( - 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(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 & // 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 &worker, }); } -bool NodeManager::FinishAssignedTask(WorkerInterface &worker) { +bool NodeManager::FinishAssignedTask(const std::shared_ptr &worker_ptr) { + // TODO (Alex): We should standardize to pass + // std::shared_ptr 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.) diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h index 6e7726528..02437de9e 100644 --- a/src/ray/raylet/node_manager.h +++ b/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 &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 diff --git a/src/ray/raylet/scheduling/cluster_task_manager.cc b/src/ray/raylet/scheduling/cluster_task_manager.cc index b971f1611..c77a89726 100644 --- a/src/ray/raylet/scheduling/cluster_task_manager.cc +++ b/src/ray/raylet/scheduling/cluster_task_manager.cc @@ -1,6 +1,7 @@ +#include "ray/raylet/scheduling/cluster_task_manager.h" + #include -#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 cluster_resource_scheduler, std::function fulfills_dependencies_func, + std::function 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 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 ready_ids) { void ClusterTaskManager::HandleTaskFinished(std::shared_ptr 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"; diff --git a/src/ray/raylet/scheduling/cluster_task_manager.h b/src/ray/raylet/scheduling/cluster_task_manager.h index d94b386a3..b97dc7659 100644 --- a/src/ray/raylet/scheduling/cluster_task_manager.h +++ b/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 cluster_resource_scheduler, std::function fulfills_dependencies_func, + std::function 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 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 cluster_resource_scheduler_; std::function fulfills_dependencies_func_; + std::function is_owner_alive_; NodeInfoGetter get_node_info_; // TODO (Alex): Implement fair queuing for these queues diff --git a/src/ray/raylet/scheduling/cluster_task_manager_test.cc b/src/ray/raylet/scheduling/cluster_task_manager_test.cc index afdf924fc..b46d8ca18 100644 --- a/src/ray/raylet/scheduling/cluster_task_manager_test.cc +++ b/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> 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 worker = + std::make_shared(WorkerID::FromRandom(), 1234); + pool_.PushWorker(std::dynamic_pointer_cast(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(); diff --git a/src/ray/raylet/test/util.h b/src/ray/raylet/test/util.h index 6625043d1..3552d8009 100644 --- a/src/ray/raylet/test/util.h +++ b/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 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 &cpu_instances) { borrowed_cpu_instances_ = cpu_instances; @@ -206,6 +204,7 @@ class MockWorker : public WorkerInterface { std::shared_ptr allocated_instances_; std::shared_ptr lifetime_allocated_instances_; std::vector borrowed_cpu_instances_; + bool is_detached_actor_; }; } // namespace raylet