diff --git a/python/ray/tests/test_reconstruction.py b/python/ray/tests/test_reconstruction.py index c14d479f8..3e05071d2 100644 --- a/python/ray/tests/test_reconstruction.py +++ b/python/ray/tests/test_reconstruction.py @@ -328,6 +328,69 @@ def test_reconstruction_chain(ray_start_cluster, reconstruction_enabled): raise e.as_instanceof_cause() +def test_reconstruction_stress(ray_start_cluster): + config = json.dumps({ + "num_heartbeats_timeout": 10, + "raylet_heartbeat_timeout_milliseconds": 100, + "lineage_pinning_enabled": 1, + "free_objects_period_milliseconds": -1, + "max_direct_call_object_size": 100, + "task_retry_delay_ms": 100, + }) + cluster = ray_start_cluster + # Head node with no resources. + cluster.add_node(num_cpus=0, _internal_config=config) + ray.init(address=cluster.address) + # Node to place the initial object. + node_to_kill = cluster.add_node( + num_cpus=1, + resources={"node1": 1}, + object_store_memory=10**8, + _internal_config=config) + cluster.add_node( + num_cpus=1, + resources={"node2": 1}, + object_store_memory=10**8, + _internal_config=config) + cluster.wait_for_nodes() + + @ray.remote + def large_object(): + return np.zeros(10**5, dtype=np.uint8) + + @ray.remote + def dependent_task(x): + return + + for _ in range(3): + obj = large_object.options(resources={"node1": 1}).remote() + ray.get(dependent_task.options(resources={"node2": 1}).remote(obj)) + + outputs = [ + large_object.options(resources={ + "node1": 1 + }).remote() for _ in range(1000) + ] + outputs = [ + dependent_task.options(resources={ + "node2": 1 + }).remote(obj) for obj in outputs + ] + + cluster.remove_node(node_to_kill, allow_graceful=False) + node_to_kill = cluster.add_node( + num_cpus=1, + resources={"node1": 1}, + object_store_memory=10**8, + _internal_config=config) + + i = 0 + while outputs: + ray.get(outputs.pop(0)) + print(i) + i += 1 + + if __name__ == "__main__": import pytest sys.exit(pytest.main(["-v", __file__])) diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index c19b4f9d0..b24a99340 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -373,6 +373,16 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_ options_.ref_counting_enabled ? reference_counter_ : nullptr, local_raylet_client_, options_.check_signals)); + auto check_node_alive_fn = [this](const ClientID &node_id) { + auto node = gcs_client_->Nodes().Get(node_id); + RAY_CHECK(node.has_value()); + return node->state() == rpc::GcsNodeInfo::ALIVE; + }; + auto reconstruct_object_callback = [this](const ObjectID &object_id) { + io_service_.post([this, object_id]() { + RAY_CHECK_OK(object_recovery_manager_->RecoverObject(object_id)); + }); + }; task_manager_.reset(new TaskManager( memory_store_, reference_counter_, actor_manager_, [this](const TaskSpecification &spec, bool delay) { @@ -387,7 +397,8 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_ } else { RAY_CHECK_OK(direct_task_submitter_->SubmitTask(spec)); } - })); + }, + check_node_alive_fn, reconstruct_object_callback)); // Create an entry for the driver task in the task table. This task is // added immediately with status RUNNING. This allows us to push errors diff --git a/src/ray/core_worker/task_manager.cc b/src/ray/core_worker/task_manager.cc index 153464fe3..7180d7329 100644 --- a/src/ray/core_worker/task_manager.cc +++ b/src/ray/core_worker/task_manager.cc @@ -186,11 +186,17 @@ void TaskManager::CompletePendingTask(const TaskID &task_id, reference_counter_->UpdateObjectSize(object_id, return_object.size()); if (return_object.in_plasma()) { - // Mark it as in plasma with a dummy object. - RAY_CHECK( - in_memory_store_->Put(RayObject(rpc::ErrorType::OBJECT_IN_PLASMA), object_id)); const auto pinned_at_raylet_id = ClientID::FromBinary(worker_addr.raylet_id()); - reference_counter_->UpdateObjectPinnedAtRaylet(object_id, pinned_at_raylet_id); + if (check_node_alive_(pinned_at_raylet_id)) { + // Mark it as in plasma with a dummy object. + RAY_CHECK(in_memory_store_->Put(RayObject(rpc::ErrorType::OBJECT_IN_PLASMA), + object_id)); + reference_counter_->UpdateObjectPinnedAtRaylet(object_id, pinned_at_raylet_id); + } else { + RAY_LOG(INFO) << "Task " << task_id << " returned object " << object_id + << " in plasma on a dead node, attempting to recover"; + reconstruct_object_callback_(object_id); + } } else { // NOTE(swang): If a direct object was promoted to plasma, then we do not // record the node ID that it was pinned at, which means that we will not diff --git a/src/ray/core_worker/task_manager.h b/src/ray/core_worker/task_manager.h index 9bc603bd9..ef646ee42 100644 --- a/src/ray/core_worker/task_manager.h +++ b/src/ray/core_worker/task_manager.h @@ -53,17 +53,22 @@ class TaskResubmissionInterface { }; using RetryTaskCallback = std::function; +using ReconstructObjectCallback = std::function; class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterface { public: TaskManager(std::shared_ptr in_memory_store, std::shared_ptr reference_counter, std::shared_ptr actor_manager, - RetryTaskCallback retry_task_callback) + RetryTaskCallback retry_task_callback, + const std::function &check_node_alive, + ReconstructObjectCallback reconstruct_object_callback) : in_memory_store_(in_memory_store), reference_counter_(reference_counter), actor_manager_(actor_manager), - retry_task_callback_(retry_task_callback) { + retry_task_callback_(retry_task_callback), + check_node_alive_(check_node_alive), + reconstruct_object_callback_(reconstruct_object_callback) { reference_counter_->SetReleaseLineageCallback( [this](const ObjectID &object_id, std::vector *ids_to_release) { RemoveLineageReference(object_id, ids_to_release); @@ -238,6 +243,17 @@ class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterfa /// Called when a task should be retried. const RetryTaskCallback retry_task_callback_; + /// Called to check whether a raylet is still alive. This is used when + /// processing a worker's reply to check whether the node that the worker + /// was on is still alive. If the node is down, the plasma objects returned by the task + /// are marked as failed. + const std::function check_node_alive_; + /// Called when processing a worker's reply if the node that the worker was + /// on died. This should be called to attempt to recover a plasma object + /// returned by the task (or store an error if the object is not + /// recoverable). + const ReconstructObjectCallback reconstruct_object_callback_; + // The number of task failures we have logged total. int64_t num_failure_logs_ GUARDED_BY(mu_) = 0; diff --git a/src/ray/core_worker/test/task_manager_test.cc b/src/ray/core_worker/test/task_manager_test.cc index d37b90dce..f02e3eb4f 100644 --- a/src/ray/core_worker/test/task_manager_test.cc +++ b/src/ray/core_worker/test/task_manager_test.cc @@ -54,11 +54,17 @@ class TaskManagerTest : public ::testing::Test { [this](const TaskSpecification &spec, bool delay) { num_retries_++; return Status::OK(); + }, + [this](const ClientID &node_id) { return all_nodes_alive_; }, + [this](const ObjectID &object_id) { + objects_to_recover_.push_back(object_id); }) {} std::shared_ptr store_; std::shared_ptr reference_counter_; std::shared_ptr actor_manager_; + bool all_nodes_alive_ = true; + std::vector objects_to_recover_; TaskManager manager_; int num_retries_ = 0; }; @@ -142,6 +148,33 @@ TEST_F(TaskManagerTest, TestTaskFailure) { ASSERT_EQ(reference_counter_->NumObjectIDsInScope(), 0); } +TEST_F(TaskManagerTest, TestPlasmaConcurrentFailure) { + TaskID caller_id = TaskID::Nil(); + rpc::Address caller_address; + auto spec = CreateTaskHelper(1, {}); + ASSERT_FALSE(manager_.IsTaskPending(spec.TaskId())); + manager_.AddPendingTask(caller_id, caller_address, spec, ""); + ASSERT_TRUE(manager_.IsTaskPending(spec.TaskId())); + auto return_id = spec.ReturnId(0); + WorkerContext ctx(WorkerType::WORKER, WorkerID::FromRandom(), JobID::FromInt(0)); + + ASSERT_TRUE(objects_to_recover_.empty()); + all_nodes_alive_ = false; + + rpc::PushTaskReply reply; + auto return_object = reply.add_return_objects(); + return_object->set_object_id(return_id.Binary()); + return_object->set_in_plasma(true); + manager_.CompletePendingTask(spec.TaskId(), reply, rpc::Address()); + + ASSERT_FALSE(manager_.IsTaskPending(spec.TaskId())); + + std::vector> results; + ASSERT_FALSE(store_->Get({return_id}, 1, 0, ctx, false, &results).ok()); + ASSERT_EQ(objects_to_recover_.size(), 1); + ASSERT_EQ(objects_to_recover_[0], return_id); +} + TEST_F(TaskManagerTest, TestTaskReconstruction) { TaskID caller_id = TaskID::Nil(); rpc::Address caller_address;