Add a flag to disable reconstruction for a killed actor (#7346)

src/ray/core_worker/core_worker.cc

@@ -872,11 +872,12 @@ Status CoreWorker::SubmitActorTask(const ActorID &actor_id, const RayFunction &f
   return status;
 }
 
-Status CoreWorker::KillActor(const ActorID &actor_id, bool force_kill) {
+Status CoreWorker::KillActor(const ActorID &actor_id, bool force_kill,
+                             bool no_reconstruction) {
   ActorHandle *actor_handle = nullptr;
   RAY_RETURN_NOT_OK(GetActorHandle(actor_id, &actor_handle));
   RAY_CHECK(actor_handle->IsDirectCallActor());
-  direct_actor_submitter_->KillActor(actor_id, force_kill);
+  direct_actor_submitter_->KillActor(actor_id, force_kill, no_reconstruction);
   return Status::OK();
 }
 
@@ -977,7 +978,8 @@ bool CoreWorker::AddActorHandle(std::unique_ptr<ActorHandle> actor_handle,
             RAY_LOG(INFO) << "Owner's handle and creation ID " << object_id
                           << " has gone out of scope, sending message to actor "
                           << actor_id << " to do a clean exit.";
-            RAY_CHECK_OK(KillActor(actor_id, /*intentional=*/true));
+            RAY_CHECK_OK(
+                KillActor(actor_id, /*force_kill=*/true, /*no_reconstruction=*/false));
           }
         }));
   }
@@ -1397,6 +1399,9 @@ void CoreWorker::HandleKillActor(const rpc::KillActorRequest &request,
 
   if (request.force_kill()) {
     RAY_LOG(INFO) << "Got KillActor, exiting immediately...";
+    if (request.no_reconstruction()) {
+      RAY_IGNORE_EXPR(local_raylet_client_->Disconnect());
+    }
     if (log_dir_ != "") {
       RayLog::ShutDownRayLog();
     }

src/ray/core_worker/core_worker.h

@@ -394,8 +394,10 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
   /// Tell an actor to exit immediately, without completing outstanding work.
   ///
   /// \param[in] actor_id ID of the actor to kill.
+  /// \param[in] no_reconstruction If set to true, the killed actor will not be
+  /// reconstructed anymore.
   /// \param[out] Status
-  Status KillActor(const ActorID &actor_id, bool force_kill);
+  Status KillActor(const ActorID &actor_id, bool force_kill, bool no_reconstruction);
 
   /// Decrease the reference count for this actor. Should be called by the
   /// language frontend when a reference to the ActorHandle destroyed.

src/ray/core_worker/lib/java/org_ray_runtime_RayNativeRuntime.cc

@@ -158,10 +158,11 @@ JNIEXPORT void JNICALL Java_org_ray_runtime_RayNativeRuntime_nativeSetResource(
 }
 
 JNIEXPORT void JNICALL Java_org_ray_runtime_RayNativeRuntime_nativeKillActor(
-    JNIEnv *env, jclass, jlong nativeCoreWorkerPointer, jbyteArray actorId) {
+    JNIEnv *env, jclass, jlong nativeCoreWorkerPointer, jbyteArray actorId,
+    jboolean noReconstruction) {
   auto core_worker = reinterpret_cast<ray::CoreWorker *>(nativeCoreWorkerPointer);
   auto status = core_worker->KillActor(JavaByteArrayToId<ActorID>(env, actorId),
-                                       /*force_kill=*/true);
+                                       /*force_kill=*/true, noReconstruction);
   THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, (void)0);
 }
 

src/ray/core_worker/lib/java/org_ray_runtime_RayNativeRuntime.h

@@ -73,12 +73,10 @@ JNIEXPORT void JNICALL Java_org_ray_runtime_RayNativeRuntime_nativeSetResource(
 /*
  * Class:     org_ray_runtime_RayNativeRuntime
  * Method:    nativeKillActor
- * Signature: (J[B)V
+ * Signature: (J[BZ)V
  */
-JNIEXPORT void JNICALL Java_org_ray_runtime_RayNativeRuntime_nativeKillActor(JNIEnv *,
-                                                                             jclass,
-                                                                             jlong,
-                                                                             jbyteArray);
+JNIEXPORT void JNICALL Java_org_ray_runtime_RayNativeRuntime_nativeKillActor(
+    JNIEnv *, jclass, jlong, jbyteArray, jboolean);
 
 #ifdef __cplusplus
 }

src/ray/core_worker/transport/direct_actor_transport.cc

@@ -23,13 +23,23 @@ using ray::rpc::ActorTableData;
 namespace ray {
 
 void CoreWorkerDirectActorTaskSubmitter::KillActor(const ActorID &actor_id,
-                                                   bool force_kill) {
+                                                   bool force_kill,
+                                                   bool no_reconstruction) {
   absl::MutexLock lock(&mu_);
-  auto inserted = pending_force_kills_.emplace(actor_id, force_kill);
+  rpc::KillActorRequest request;
+  request.set_intended_actor_id(actor_id.Binary());
+  request.set_force_kill(force_kill);
+  request.set_no_reconstruction(no_reconstruction);
+  auto inserted = pending_force_kills_.emplace(actor_id, request);
   if (!inserted.second && force_kill) {
     // Overwrite the previous request to kill the actor if the new request is a
     // force kill.
-    inserted.first->second = true;
+    inserted.first->second.set_force_kill(true);
+    if (no_reconstruction) {
+      // Overwrite the previous request to disable reconstruction if the new request's
+      // no_reconstruction flag is set to true.
+      inserted.first->second.set_no_reconstruction(true);
+    }
   }
   auto it = rpc_clients_.find(actor_id);
   if (it == rpc_clients_.end()) {
@@ -102,9 +112,7 @@ void CoreWorkerDirectActorTaskSubmitter::ConnectActor(const ActorID &actor_id,
     rpc_clients_[actor_id] =
         std::shared_ptr<rpc::CoreWorkerClientInterface>(client_factory_(address));
   }
-  if (pending_requests_.count(actor_id) > 0) {
-    SendPendingTasks(actor_id);
-  }
+  SendPendingTasks(actor_id);
 }
 
 void CoreWorkerDirectActorTaskSubmitter::DisconnectActor(const ActorID &actor_id,
@@ -135,6 +143,8 @@ void CoreWorkerDirectActorTaskSubmitter::DisconnectActor(const ActorID &actor_id
     // replies. They will be treated as failed once the connection dies.
     // We retain the sequencing information so that we can properly fail
     // any tasks submitted after the actor death.
+
+    pending_force_kills_.erase(actor_id);
   }
 }
 
@@ -145,11 +155,9 @@ void CoreWorkerDirectActorTaskSubmitter::SendPendingTasks(const ActorID &actor_i
   // client.
   auto it = pending_force_kills_.find(actor_id);
   if (it != pending_force_kills_.end()) {
-    rpc::KillActorRequest request;
-    request.set_intended_actor_id(actor_id.Binary());
-    request.set_force_kill(it->second);
+    RAY_LOG(INFO) << "Sending KillActor request to actor " << actor_id;
     // It's okay if this fails because this means the worker is already dead.
-    RAY_UNUSED(client->KillActor(request, nullptr));
+    RAY_UNUSED(client->KillActor(it->second, nullptr));
     pending_force_kills_.erase(it);
   }
 

src/ray/core_worker/transport/direct_actor_transport.h

@@ -70,7 +70,9 @@ class CoreWorkerDirectActorTaskSubmitter {
   /// \param[in] actor_id The actor_id of the actor to kill.
   /// \param[in] force_kill Whether to force kill the actor, or let the actor
   /// try a clean exit.
-  void KillActor(const ActorID &actor_id, bool force_kill);
+  /// \param[in] no_reconstruction If set to true, the killed actor will not be
+  /// reconstructed anymore.
+  void KillActor(const ActorID &actor_id, bool force_kill, bool no_reconstruction);
 
   /// Create connection to actor and send all pending tasks.
   ///
@@ -134,8 +136,10 @@ class CoreWorkerDirectActorTaskSubmitter {
   /// rpc_clients_ map.
   absl::flat_hash_map<ActorID, std::string> worker_ids_ GUARDED_BY(mu_);
 
-  /// Set of actor ids that should be force killed once a client is available.
-  absl::flat_hash_map<ActorID, bool> pending_force_kills_ GUARDED_BY(mu_);
+  /// Map from actor ids that should be force killed once a client is available to the
+  /// pending kill actor requests.
+  absl::flat_hash_map<ActorID, rpc::KillActorRequest> pending_force_kills_
+      GUARDED_BY(mu_);
 
   /// Map from actor id to the actor's pending requests. Each actor's requests
   /// are ordered by the task number in the request.

src/ray/protobuf/core_worker.proto

@@ -172,6 +172,8 @@ message KillActorRequest {
   bytes intended_actor_id = 1;
   // Whether to force kill the actor.
   bool force_kill = 2;
+  // If set to true, the killed actor will not be reconstructed anymore.
+  bool no_reconstruction = 3;
 }
 
 message KillActorReply {