[core] Enable object reconstruction for retryable actor tasks (#9557)

* Test actor plasma reconstruction

* Allow resubmission of actor tasks

* doc

* Test for actor constructor

* Kill PID before removing node

* Kill pid before node

src/ray/core_worker/actor_handle.cc

@@ -93,6 +93,15 @@ void ActorHandle::SetActorTaskSpec(TaskSpecBuilder &builder, const ObjectID new_
   actor_cursor_ = new_cursor;
 }
 
+void ActorHandle::SetResubmittedActorTaskSpec(TaskSpecification &spec,
+                                              const ObjectID new_cursor) {
+  absl::MutexLock guard(&mutex_);
+  auto mutable_spec = spec.GetMutableMessage().mutable_actor_task_spec();
+  mutable_spec->set_previous_actor_task_dummy_object_id(actor_cursor_.Binary());
+  mutable_spec->set_actor_counter(task_counter_++);
+  actor_cursor_ = new_cursor;
+}
+
 void ActorHandle::Serialize(std::string *output) { inner_.SerializeToString(output); }
 
 }  // namespace ray

src/ray/core_worker/actor_handle.h

@@ -63,8 +63,22 @@ class ActorHandle {
 
   std::string ExtensionData() const { return inner_.extension_data(); }
 
+  /// Set the actor task spec fields.
+  ///
+  /// \param[in] builder Task spec builder.
+  /// \param[in] new_cursor Actor dummy object. This is legacy code needed for
+  /// raylet-based actor restart.
   void SetActorTaskSpec(TaskSpecBuilder &builder, const ObjectID new_cursor);
 
+  /// Reset the actor task spec fields of an existing task so that the task can
+  /// be re-executed.
+  ///
+  /// \param[in] spec An existing task spec that has executed on the actor
+  /// before.
+  /// \param[in] new_cursor Actor dummy object. This is legacy code needed for
+  /// raylet-based actor restart.
+  void SetResubmittedActorTaskSpec(TaskSpecification &spec, const ObjectID new_cursor);
+
   void Serialize(std::string *output);
 
   int64_t MaxTaskRetries() const { return inner_.max_task_retries(); }

src/ray/core_worker/core_worker.cc

@@ -380,7 +380,7 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_
   };
   task_manager_.reset(new TaskManager(
       memory_store_, reference_counter_, actor_reporter_,
-      [this](const TaskSpecification &spec, bool delay) {
+      [this](TaskSpecification &spec, bool delay) {
         if (delay) {
           // Retry after a delay to emulate the existing Raylet reconstruction
           // behaviour. TODO(ekl) backoff exponentially.
@@ -392,7 +392,13 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_
         } else {
           RAY_LOG(ERROR) << "Resubmitting task that produced lost plasma object: "
                          << spec.DebugString();
-          RAY_CHECK_OK(direct_task_submitter_->SubmitTask(spec));
+          if (spec.IsActorTask()) {
+            const auto &actor_handle = actor_manager_->GetActorHandle(spec.ActorId());
+            actor_handle->SetResubmittedActorTaskSpec(spec, spec.ActorDummyObject());
+            RAY_CHECK_OK(direct_actor_submitter_->SubmitTask(spec));
+          } else {
+            RAY_CHECK_OK(direct_task_submitter_->SubmitTask(spec));
+          }
         }
       },
       check_node_alive_fn, reconstruct_object_callback));

src/ray/core_worker/task_manager.cc

@@ -87,9 +87,6 @@ Status TaskManager::ResubmitTask(const TaskID &task_id,
     if (it == submissible_tasks_.end()) {
       return Status::Invalid("Task spec missing");
     }
-    if (it->second.spec.IsActorTask()) {
-      return Status::Invalid("Cannot reconstruct objects returned by actors");
-    }
 
     if (!it->second.pending) {
       resubmit = true;
@@ -118,6 +115,11 @@ Status TaskManager::ResubmitTask(const TaskID &task_id,
     reference_counter_->UpdateResubmittedTaskReferences(*task_deps);
   }
 
+  if (spec.IsActorTask()) {
+    const auto actor_creation_return_id = spec.ActorCreationDummyObjectId();
+    reference_counter_->UpdateResubmittedTaskReferences({actor_creation_return_id});
+  }
+
   if (resubmit) {
     retry_task_callback_(spec, /*delay=*/false);
   }

src/ray/core_worker/task_manager.h

@@ -51,7 +51,7 @@ class TaskResubmissionInterface {
   virtual ~TaskResubmissionInterface() {}
 };
 
-using RetryTaskCallback = std::function<void(const TaskSpecification &spec, bool delay)>;
+using RetryTaskCallback = std::function<void(TaskSpecification &spec, bool delay)>;
 using ReconstructObjectCallback = std::function<void(const ObjectID &object_id)>;
 
 class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterface {

src/ray/core_worker/test/task_manager_test.cc

@@ -52,7 +52,7 @@ class TaskManagerTest : public ::testing::Test {
             /*distributed_ref_counting_enabled=*/true, lineage_pinning_enabled))),
         actor_reporter_(std::shared_ptr<ActorReporterInterface>(new MockActorManager())),
         manager_(store_, reference_counter_, actor_reporter_,
-                 [this](const TaskSpecification &spec, bool delay) {
+                 [this](TaskSpecification &spec, bool delay) {
                    num_retries_++;
                    return Status::OK();
                  },