[core] Fix race condition for object reconstruction (#8791)

* Fix

* doc

* Unit test

* Update src/ray/core_worker/task_manager.h

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

* Update src/ray/core_worker/task_manager.h

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

* Update src/ray/core_worker/task_manager.h

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

* lint

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

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__]))

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

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

src/ray/core_worker/task_manager.h

@@ -53,17 +53,22 @@ class TaskResubmissionInterface {
 };
 
 using RetryTaskCallback = std::function<void(const TaskSpecification &spec, bool delay)>;
+using ReconstructObjectCallback = std::function<void(const ObjectID &object_id)>;
 
 class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterface {
  public:
   TaskManager(std::shared_ptr<CoreWorkerMemoryStore> in_memory_store,
               std::shared_ptr<ReferenceCounter> reference_counter,
               std::shared_ptr<ActorManagerInterface> actor_manager,
-              RetryTaskCallback retry_task_callback)
+              RetryTaskCallback retry_task_callback,
+              const std::function<bool(const ClientID &node_id)> &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<ObjectID> *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<bool(const ClientID &node_id)> 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;
 

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<CoreWorkerMemoryStore> store_;
   std::shared_ptr<ReferenceCounter> reference_counter_;
   std::shared_ptr<ActorManagerInterface> actor_manager_;
+  bool all_nodes_alive_ = true;
+  std::vector<ObjectID> 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<std::shared_ptr<RayObject>> 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;