[core] Make CancelTask RPC Fault Tolerant

python/ray/serve/tests/test_streaming_response.py

@@ -11,7 +11,11 @@
 import ray
 from ray import serve
 from ray._common.test_utils import SignalActor
-from ray.serve._private.test_utils import get_application_url, get_application_urls
+from ray.serve._private.test_utils import (
+    get_application_url,
+    get_application_urls,
+    send_signal_on_cancellation,
+)
 from ray.serve.handle import DeploymentHandle
 
 
@@ -326,12 +330,9 @@ def test_http_disconnect(serve_instance):
     class SimpleGenerator:
         def __call__(self, request: Request) -> StreamingResponse:
             async def wait_for_disconnect():
-                try:
-                    yield "hi"
-                    await asyncio.sleep(100)
-                except asyncio.CancelledError:
-                    print("Cancelled!")
-                    signal_actor.send.remote()
+                yield "hi"
+                async with send_signal_on_cancellation(signal_actor):
+                    pass
 
             return StreamingResponse(wait_for_disconnect())
 

python/ray/tests/test_core_worker_fault_tolerance.py

@@ -172,7 +172,7 @@ def inject_cancel_remote_task_rpc_failure(monkeypatch, request):
     failure = RPC_FAILURE_MAP[deterministic_failure]
     monkeypatch.setenv(
         "RAY_testing_rpc_failure",
-        f"CoreWorkerService.grpc_client.CancelRemoteTask=1:{failure}",
+        f"CoreWorkerService.grpc_client.RequestOwnerToCancelTask=1:{failure}",
     )
 
 

python/ray/tests/test_raylet_fault_tolerance.py

@@ -4,12 +4,13 @@
 import pytest
 
 import ray
+from ray._common.test_utils import SignalActor, wait_for_condition
 from ray._private.test_utils import (
     RPC_FAILURE_MAP,
     RPC_FAILURE_TYPES,
-    wait_for_condition,
 )
 from ray.core.generated import autoscaler_pb2
+from ray.exceptions import GetTimeoutError, TaskCancelledError
 from ray.util.placement_group import placement_group, remove_placement_group
 from ray.util.scheduling_strategies import (
     NodeAffinitySchedulingStrategy,
@@ -247,5 +248,57 @@ def verify_process_killed():
     wait_for_condition(verify_process_killed, timeout=30)
 
 
+@pytest.fixture
+def inject_cancel_local_task_rpc_failure(monkeypatch, request):
+    failure = RPC_FAILURE_MAP[request.param]
+    monkeypatch.setenv(
+        "RAY_testing_rpc_failure",
+        f"NodeManagerService.grpc_client.CancelLocalTask=1:{failure}",
+    )
+
+
+@pytest.mark.parametrize(
+    "inject_cancel_local_task_rpc_failure", RPC_FAILURE_TYPES, indirect=True
+)
+@pytest.mark.parametrize("force_kill", [True, False])
+def test_cancel_local_task_rpc_retry_and_idempotency(
+    inject_cancel_local_task_rpc_failure, force_kill, shutdown_only
+):
+    """Test that CancelLocalTask RPC retries work correctly.
+
+    Verify that the RPC is idempotent when network failures occur.
+    When force_kill=True, verify the worker process is actually killed using psutil.
+    """
+    ray.init(num_cpus=1)
+    signaler = SignalActor.remote()
+
+    @ray.remote(num_cpus=1)
+    def get_pid():
+        return os.getpid()
+
+    @ray.remote(num_cpus=1)
+    def blocking_task():
+        return ray.get(signaler.wait.remote())
+
+    worker_pid = ray.get(get_pid.remote())
+
+    blocking_ref = blocking_task.remote()
+
+    with pytest.raises(GetTimeoutError):
+        ray.get(blocking_ref, timeout=1)
+
+    ray.cancel(blocking_ref, force=force_kill)
+
+    with pytest.raises(TaskCancelledError):
+        ray.get(blocking_ref, timeout=10)
+
+    if force_kill:
+
+        def verify_process_killed():
+            return not psutil.pid_exists(worker_pid)
+
+        wait_for_condition(verify_process_killed, timeout=30)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-sv", __file__]))

src/mock/ray/core_worker/core_worker.h

@@ -89,9 +89,9 @@ class MockCoreWorker : public CoreWorker {
                rpc::SendReplyCallback send_reply_callback),
               (override));
   MOCK_METHOD(void,
-              HandleCancelRemoteTask,
-              (rpc::CancelRemoteTaskRequest request,
-               rpc::CancelRemoteTaskReply *reply,
+              HandleRequestOwnerToCancelTask,
+              (rpc::RequestOwnerToCancelTaskRequest request,
+               rpc::RequestOwnerToCancelTaskReply *reply,
                rpc::SendReplyCallback send_reply_callback),
               (override));
   MOCK_METHOD(void,

src/mock/ray/raylet_client/raylet_client.h

@@ -156,6 +156,11 @@ class MockRayletClientInterface : public RayletClientInterface {
               (const rpc::ClientCallback<rpc::GlobalGCReply> &callback),
               (override));
   MOCK_METHOD(int64_t, GetPinsInFlight, (), (const, override));
+  MOCK_METHOD(void,
+              CancelLocalTask,
+              (const rpc::CancelLocalTaskRequest &request,
+               const rpc::ClientCallback<rpc::CancelLocalTaskReply> &callback),
+              (override));
 };
 
 }  // namespace ray

src/mock/ray/rpc/worker/core_worker_client.h

@@ -87,9 +87,9 @@ class MockCoreWorkerClientInterface : public CoreWorkerClientInterface {
                const ClientCallback<CancelTaskReply> &callback),
               (override));
   MOCK_METHOD(void,
-              CancelRemoteTask,
-              (CancelRemoteTaskRequest && request,
-               const ClientCallback<CancelRemoteTaskReply> &callback),
+              RequestOwnerToCancelTask,
+              (RequestOwnerToCancelTaskRequest && request,
+               const ClientCallback<RequestOwnerToCancelTaskReply> &callback),
               (override));
   MOCK_METHOD(void,
               GetCoreWorkerStats,

src/ray/core_worker/core_worker.cc

@@ -2460,12 +2460,13 @@ Status CoreWorker::CancelTask(const ObjectID &object_id,
   }
 
   if (obj_addr.SerializeAsString() != rpc_address_.SerializeAsString()) {
-    // We don't have CancelRemoteTask for actor_task_submitter_
+    // We don't have RequestOwnerToCancelTask for actor_task_submitter_
     // because it requires the same implementation.
     RAY_LOG(DEBUG).WithField(object_id)
         << "Request to cancel a task of object to an owner "
         << obj_addr.SerializeAsString();
-    normal_task_submitter_->CancelRemoteTask(object_id, obj_addr, force_kill, recursive);
+    normal_task_submitter_->RequestOwnerToCancelTask(
+        object_id, obj_addr, force_kill, recursive);
     return Status::OK();
   }
 
@@ -3910,9 +3911,10 @@ void CoreWorker::ProcessSubscribeForRefRemoved(
   reference_counter_->SubscribeRefRemoved(object_id, contained_in_id, owner_address);
 }
 
-void CoreWorker::HandleCancelRemoteTask(rpc::CancelRemoteTaskRequest request,
-                                        rpc::CancelRemoteTaskReply *reply,
-                                        rpc::SendReplyCallback send_reply_callback) {
+void CoreWorker::HandleRequestOwnerToCancelTask(
+    rpc::RequestOwnerToCancelTaskRequest request,
+    rpc::RequestOwnerToCancelTaskReply *reply,
+    rpc::SendReplyCallback send_reply_callback) {
   auto status = CancelTask(ObjectID::FromBinary(request.remote_object_id()),
                            request.force_kill(),
                            request.recursive());

src/ray/core_worker/core_worker.h

@@ -1208,9 +1208,9 @@ class CoreWorker {
                         rpc::SendReplyCallback send_reply_callback);
 
   /// Implements gRPC server handler.
-  void HandleCancelRemoteTask(rpc::CancelRemoteTaskRequest request,
-                              rpc::CancelRemoteTaskReply *reply,
-                              rpc::SendReplyCallback send_reply_callback);
+  void HandleRequestOwnerToCancelTask(rpc::RequestOwnerToCancelTaskRequest request,
+                                      rpc::RequestOwnerToCancelTaskReply *reply,
+                                      rpc::SendReplyCallback send_reply_callback);
 
   /// Implements gRPC server handler.
   void HandlePlasmaObjectReady(rpc::PlasmaObjectReadyRequest request,

src/ray/core_worker/core_worker_process.cc

@@ -500,6 +500,8 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
 
   auto actor_task_submitter = std::make_unique<ActorTaskSubmitter>(
       *core_worker_client_pool,
+      *raylet_client_pool,
+      gcs_client,
       *memory_store,
       *task_manager,
       *actor_creator,
@@ -538,6 +540,7 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
       local_raylet_rpc_client,
       core_worker_client_pool,
       raylet_client_pool,
+      gcs_client,
       std::move(lease_policy),
       memory_store,
       *task_manager,
@@ -554,7 +557,7 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
         // OBJECT_STORE.
         return rpc::TensorTransport::OBJECT_STORE;
       },
-      boost::asio::steady_timer(io_service_),
+      io_service_,
       *scheduler_placement_time_ms_histogram_);
 
   auto report_locality_data_callback = [this](

src/ray/core_worker/core_worker_rpc_proxy.h

@@ -56,7 +56,7 @@ class CoreWorkerServiceHandlerProxy : public rpc::CoreWorkerServiceHandler {
   RAY_CORE_WORKER_RPC_PROXY(ReportGeneratorItemReturns)
   RAY_CORE_WORKER_RPC_PROXY(KillActor)
   RAY_CORE_WORKER_RPC_PROXY(CancelTask)
-  RAY_CORE_WORKER_RPC_PROXY(CancelRemoteTask)
+  RAY_CORE_WORKER_RPC_PROXY(RequestOwnerToCancelTask)
   RAY_CORE_WORKER_RPC_PROXY(RegisterMutableObjectReader)
   RAY_CORE_WORKER_RPC_PROXY(GetCoreWorkerStats)
   RAY_CORE_WORKER_RPC_PROXY(LocalGC)

src/ray/core_worker/grpc_service.cc

@@ -77,7 +77,7 @@ void CoreWorkerGrpcService::InitServerCallFactories(
                                                           max_active_rpcs_per_handler_,
                                                           ClusterIdAuthType::NO_AUTH);
   RPC_SERVICE_HANDLER_CUSTOM_AUTH_SERVER_METRICS_DISABLED(CoreWorkerService,
-                                                          CancelRemoteTask,
+                                                          RequestOwnerToCancelTask,
                                                           max_active_rpcs_per_handler_,
                                                           ClusterIdAuthType::NO_AUTH);
   RPC_SERVICE_HANDLER_CUSTOM_AUTH_SERVER_METRICS_DISABLED(CoreWorkerService,

src/ray/core_worker/grpc_service.h

@@ -96,9 +96,9 @@ class CoreWorkerServiceHandler : public DelayedServiceHandler {
                                 CancelTaskReply *reply,
                                 SendReplyCallback send_reply_callback) = 0;
 
-  virtual void HandleCancelRemoteTask(CancelRemoteTaskRequest request,
-                                      CancelRemoteTaskReply *reply,
-                                      SendReplyCallback send_reply_callback) = 0;
+  virtual void HandleRequestOwnerToCancelTask(RequestOwnerToCancelTaskRequest request,
+                                              RequestOwnerToCancelTaskReply *reply,
+                                              SendReplyCallback send_reply_callback) = 0;
 
   virtual void HandleRegisterMutableObjectReader(
       RegisterMutableObjectReaderRequest request,

src/ray/core_worker/task_submission/BUILD.bazel

@@ -53,6 +53,17 @@ ray_cc_library(
     ],
 )
 
+ray_cc_library(
+    name = "task_submission_util",
+    hdrs = ["task_submission_util.h"],
+    visibility = [":__subpackages__"],
+    deps = [
+        "//src/ray/common:asio",
+        "//src/ray/common:id",
+        "//src/ray/gcs_rpc_client:gcs_client",
+    ],
+)
+
 ray_cc_library(
     name = "actor_task_submitter",
     srcs = ["actor_task_submitter.cc"],
@@ -66,12 +77,16 @@ ray_cc_library(
         ":dependency_resolver",
         ":out_of_order_actor_submit_queue",
         ":sequential_actor_submit_queue",
+        ":task_submission_util",
         "//src/ray/common:asio",
         "//src/ray/common:id",
         "//src/ray/common:protobuf_utils",
         "//src/ray/core_worker:actor_creator",
         "//src/ray/core_worker:reference_counter_interface",
         "//src/ray/core_worker_rpc_client:core_worker_client_pool",
+        "//src/ray/gcs_rpc_client:gcs_client",
+        "//src/ray/raylet_rpc_client:raylet_client_interface",
+        "//src/ray/raylet_rpc_client:raylet_client_pool",
         "//src/ray/rpc:rpc_callback_types",
         "//src/ray/util:time",
         "@com_google_absl//absl/base:core_headers",
@@ -90,14 +105,17 @@ ray_cc_library(
     ],
     deps = [
         ":dependency_resolver",
+        ":task_submission_util",
         "//src/ray/common:id",
         "//src/ray/common:lease",
         "//src/ray/common:protobuf_utils",
         "//src/ray/core_worker:lease_policy",
         "//src/ray/core_worker:memory_store",
         "//src/ray/core_worker:task_manager_interface",
         "//src/ray/core_worker_rpc_client:core_worker_client_pool",
+        "//src/ray/gcs_rpc_client:gcs_client",
         "//src/ray/raylet_rpc_client:raylet_client_interface",
+        "//src/ray/raylet_rpc_client:raylet_client_pool",
         "//src/ray/util:time",
         "@com_google_absl//absl/base:core_headers",
     ],