[cherry-pick] [core] Fix deadlock when cancelling stale requests on in-order actors (#57746) (#57768)

## Description
Cherry picking #57746

Signed-off-by: dayshah <dhyey2019@gmail.com>

Author: dayshah

python/ray/tests/test_core_worker_fault_tolerance.py

@@ -1,6 +1,85 @@
+import sys
+
+import numpy as np
 import pytest
 
 import ray
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+
+@pytest.mark.parametrize(
+    "allow_out_of_order_execution",
+    [True, False],
+)
+@pytest.mark.parametrize("deterministic_failure", ["request", "response"])
+def test_push_actor_task_failure(
+    monkeypatch,
+    ray_start_cluster,
+    allow_out_of_order_execution: bool,
+    deterministic_failure: str,
+):
+    with monkeypatch.context() as m:
+        m.setenv(
+            "RAY_testing_rpc_failure",
+            "CoreWorkerService.grpc_client.PushTask=2:"
+            + ("100:0" if deterministic_failure == "request" else "0:100"),
+        )
+        m.setenv("RAY_actor_scheduling_queue_max_reorder_wait_seconds", "0")
+        cluster = ray_start_cluster
+        cluster.add_node(num_cpus=1)
+        ray.init(address=cluster.address)
+
+        @ray.remote(
+            max_task_retries=-1,
+            allow_out_of_order_execution=allow_out_of_order_execution,
+        )
+        class RetryActor:
+            def echo(self, value):
+                return value
+
+        refs = []
+        actor = RetryActor.remote()
+        for i in range(10):
+            refs.append(actor.echo.remote(i))
+        assert ray.get(refs) == list(range(10))
+
+
+@pytest.mark.parametrize("deterministic_failure", ["request", "response"])
+def test_update_object_location_batch_failure(
+    monkeypatch, ray_start_cluster, deterministic_failure
+):
+    with monkeypatch.context() as m:
+        m.setenv(
+            "RAY_testing_rpc_failure",
+            "CoreWorkerService.grpc_client.UpdateObjectLocationBatch=1:"
+            + ("100:0" if deterministic_failure == "request" else "0:100"),
+        )
+        cluster = ray_start_cluster
+        head_node_id = cluster.add_node(
+            num_cpus=0,
+        ).node_id
+        ray.init(address=cluster.address)
+        worker_node_id = cluster.add_node(num_cpus=1).node_id
+
+        @ray.remote(num_cpus=1)
+        def create_large_object():
+            return np.zeros(100 * 1024 * 1024, dtype=np.uint8)
+
+        @ray.remote(num_cpus=0)
+        def consume_large_object(obj):
+            return sys.getsizeof(obj)
+
+        obj_ref = create_large_object.options(
+            scheduling_strategy=NodeAffinitySchedulingStrategy(
+                node_id=worker_node_id, soft=False
+            )
+        ).remote()
+        consume_ref = consume_large_object.options(
+            scheduling_strategy=NodeAffinitySchedulingStrategy(
+                node_id=head_node_id, soft=False
+            )
+        ).remote(obj_ref)
+        assert ray.get(consume_ref, timeout=10) > 0
 
 
 @pytest.mark.parametrize("deterministic_failure", ["request", "response"])

python/ray/tests/test_failure.py

@@ -17,7 +17,6 @@
     init_error_pubsub,
 )
 from ray.exceptions import ActorDiedError, GetTimeoutError, RayActorError, RayTaskError
-from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
 
 
 def test_unhandled_errors(ray_start_regular):
@@ -688,70 +687,6 @@ def func():
     caplog.clear()
 
 
-def test_transient_error_retry(monkeypatch, ray_start_cluster):
-    with monkeypatch.context() as m:
-        # This test submits 200 tasks with infinite retries and verifies that all tasks eventually succeed in the unstable network environment.
-        m.setenv(
-            "RAY_testing_rpc_failure",
-            "CoreWorkerService.grpc_client.PushTask=100:25:25",
-        )
-        cluster = ray_start_cluster
-        cluster.add_node(
-            num_cpus=1,
-            resources={"head": 1},
-        )
-        ray.init(address=cluster.address)
-
-        @ray.remote(max_task_retries=-1, resources={"head": 1})
-        class RetryActor:
-            def echo(self, value):
-                return value
-
-        refs = []
-        actor = RetryActor.remote()
-        for i in range(200):
-            refs.append(actor.echo.remote(i))
-        assert ray.get(refs) == list(range(200))
-
-
-@pytest.mark.parametrize("deterministic_failure", ["request", "response"])
-def test_update_object_location_batch_failure(
-    monkeypatch, ray_start_cluster, deterministic_failure
-):
-    with monkeypatch.context() as m:
-        m.setenv(
-            "RAY_testing_rpc_failure",
-            "CoreWorkerService.grpc_client.UpdateObjectLocationBatch=1:"
-            + ("100:0" if deterministic_failure == "request" else "0:100"),
-        )
-        cluster = ray_start_cluster
-        head_node_id = cluster.add_node(
-            num_cpus=0,
-        ).node_id
-        ray.init(address=cluster.address)
-        worker_node_id = cluster.add_node(num_cpus=1).node_id
-
-        @ray.remote(num_cpus=1)
-        def create_large_object():
-            return np.zeros(100 * 1024 * 1024, dtype=np.uint8)
-
-        @ray.remote(num_cpus=0)
-        def consume_large_object(obj):
-            return sys.getsizeof(obj)
-
-        obj_ref = create_large_object.options(
-            scheduling_strategy=NodeAffinitySchedulingStrategy(
-                node_id=worker_node_id, soft=False
-            )
-        ).remote()
-        consume_ref = consume_large_object.options(
-            scheduling_strategy=NodeAffinitySchedulingStrategy(
-                node_id=head_node_id, soft=False
-            )
-        ).remote(obj_ref)
-        assert ray.get(consume_ref, timeout=10) > 0
-
-
 def test_raytaskerror_serialization(ray_start_regular):
     """Test that RayTaskError with dual exception instances can be properly serialized."""
     import ray.cloudpickle as pickle

src/ray/core_worker/task_execution/task_receiver.cc

@@ -171,16 +171,9 @@ void TaskReceiver::HandleTask(rpc::PushTaskRequest request,
     if (canceled_task_spec.IsActorTask()) {
       // If task cancelation is due to worker shutdown, propagate that information
       // to the submitter.
-      bool is_worker_exiting = false;
-      {
-        absl::MutexLock lock(&stop_mu_);
-        is_worker_exiting = stopping_;
-        if (stopping_) {
-          reply->set_worker_exiting(true);
-          reply->set_was_cancelled_before_running(true);
-        }
-      }
-      if (is_worker_exiting) {
+      if (stopping_) {
+        reply->set_worker_exiting(true);
+        reply->set_was_cancelled_before_running(true);
         canceled_send_reply_callback(Status::OK(), nullptr, nullptr);
       } else {
         canceled_send_reply_callback(status, nullptr, nullptr);
@@ -191,82 +184,79 @@ void TaskReceiver::HandleTask(rpc::PushTaskRequest request,
     }
   };
 
-  {
-    absl::MutexLock lock(&stop_mu_);
-    task_spec = TaskSpecification(std::move(*request.mutable_task_spec()));
-    if (stopping_) {
-      reply->set_was_cancelled_before_running(true);
-      if (task_spec.IsActorTask()) {
-        reply->set_worker_exiting(true);
-      }
-      send_reply_callback(Status::OK(), nullptr, nullptr);
-      return;
+  task_spec = TaskSpecification(std::move(*request.mutable_task_spec()));
+  if (stopping_) {
+    reply->set_was_cancelled_before_running(true);
+    if (task_spec.IsActorTask()) {
+      reply->set_worker_exiting(true);
     }
+    send_reply_callback(Status::OK(), nullptr, nullptr);
+    return;
+  }
 
-    if (task_spec.IsActorCreationTask()) {
-      SetupActor(task_spec.IsAsyncioActor(),
-                 task_spec.MaxActorConcurrency(),
-                 task_spec.AllowOutOfOrderExecution());
-    }
+  if (task_spec.IsActorCreationTask()) {
+    SetupActor(task_spec.IsAsyncioActor(),
+               task_spec.MaxActorConcurrency(),
+               task_spec.AllowOutOfOrderExecution());
+  }
 
-    if (!task_spec.IsActorTask()) {
-      resource_ids = ResourceMappingType{};
-      for (const auto &mapping : request.resource_mapping()) {
-        std::vector<std::pair<int64_t, double>> rids;
-        rids.reserve(mapping.resource_ids().size());
-        for (const auto &ids : mapping.resource_ids()) {
-          rids.emplace_back(ids.index(), ids.quantity());
-        }
-        (*resource_ids)[mapping.name()] = std::move(rids);
+  if (!task_spec.IsActorTask()) {
+    resource_ids = ResourceMappingType{};
+    for (const auto &mapping : request.resource_mapping()) {
+      std::vector<std::pair<int64_t, double>> rids;
+      rids.reserve(mapping.resource_ids().size());
+      for (const auto &ids : mapping.resource_ids()) {
+        rids.emplace_back(ids.index(), ids.quantity());
       }
+      (*resource_ids)[mapping.name()] = std::move(rids);
     }
+  }
 
-    if (task_spec.IsActorTask()) {
-      auto it = actor_scheduling_queues_.find(task_spec.CallerWorkerId());
-      if (it == actor_scheduling_queues_.end()) {
-        it = actor_scheduling_queues_
-                 .emplace(
-                     task_spec.CallerWorkerId(),
-                     allow_out_of_order_execution_
-                         ? std::unique_ptr<SchedulingQueue>(
-                               std::make_unique<OutOfOrderActorSchedulingQueue>(
-                                   task_execution_service_,
-                                   waiter_,
-                                   task_event_buffer_,
-                                   pool_manager_,
-                                   fiber_state_manager_,
-                                   is_asyncio_,
-                                   fiber_max_concurrency_,
-                                   concurrency_groups_))
-                         : std::unique_ptr<
-                               SchedulingQueue>(std::make_unique<ActorSchedulingQueue>(
-                               task_execution_service_,
-                               waiter_,
-                               task_event_buffer_,
-                               pool_manager_,
-                               RayConfig::instance()
-                                   .actor_scheduling_queue_max_reorder_wait_seconds())))
-                 .first;
-      }
-
-      auto accept_callback = make_accept_callback();
-      it->second->Add(request.sequence_number(),
-                      request.client_processed_up_to(),
-                      std::move(accept_callback),
-                      std::move(cancel_callback),
-                      std::move(send_reply_callback),
-                      std::move(task_spec));
-    } else {
-      RAY_LOG(DEBUG) << "Adding task " << task_spec.TaskId()
-                     << " to normal scheduling task queue.";
-      auto accept_callback = make_accept_callback();
-      normal_scheduling_queue_->Add(request.sequence_number(),
-                                    request.client_processed_up_to(),
-                                    std::move(accept_callback),
-                                    std::move(cancel_callback),
-                                    std::move(send_reply_callback),
-                                    std::move(task_spec));
+  if (task_spec.IsActorTask()) {
+    auto it = actor_scheduling_queues_.find(task_spec.CallerWorkerId());
+    if (it == actor_scheduling_queues_.end()) {
+      it = actor_scheduling_queues_
+               .emplace(
+                   task_spec.CallerWorkerId(),
+                   allow_out_of_order_execution_
+                       ? std::unique_ptr<SchedulingQueue>(
+                             std::make_unique<OutOfOrderActorSchedulingQueue>(
+                                 task_execution_service_,
+                                 waiter_,
+                                 task_event_buffer_,
+                                 pool_manager_,
+                                 fiber_state_manager_,
+                                 is_asyncio_,
+                                 fiber_max_concurrency_,
+                                 concurrency_groups_))
+                       : std::unique_ptr<SchedulingQueue>(
+                             std::make_unique<ActorSchedulingQueue>(
+                                 task_execution_service_,
+                                 waiter_,
+                                 task_event_buffer_,
+                                 pool_manager_,
+                                 RayConfig::instance()
+                                     .actor_scheduling_queue_max_reorder_wait_seconds())))
+               .first;
     }
+
+    auto accept_callback = make_accept_callback();
+    it->second->Add(request.sequence_number(),
+                    request.client_processed_up_to(),
+                    std::move(accept_callback),
+                    std::move(cancel_callback),
+                    std::move(send_reply_callback),
+                    std::move(task_spec));
+  } else {
+    RAY_LOG(DEBUG) << "Adding task " << task_spec.TaskId()
+                   << " to normal scheduling task queue.";
+    auto accept_callback = make_accept_callback();
+    normal_scheduling_queue_->Add(request.sequence_number(),
+                                  request.client_processed_up_to(),
+                                  std::move(accept_callback),
+                                  std::move(cancel_callback),
+                                  std::move(send_reply_callback),
+                                  std::move(task_spec));
   }
 }
 
@@ -315,12 +305,8 @@ void TaskReceiver::SetupActor(bool is_asyncio,
 }
 
 void TaskReceiver::Stop() {
-  {
-    absl::MutexLock lock(&stop_mu_);
-    if (stopping_) {
-      return;
-    }
-    stopping_ = true;
+  if (stopping_.exchange(true)) {
+    return;
   }
   for (const auto &[_, scheduling_queue] : actor_scheduling_queues_) {
     scheduling_queue->Stop();

src/ray/core_worker/task_execution/task_receiver.h

@@ -104,10 +104,8 @@ class TaskReceiver {
   void SetActorReprName(const std::string &repr_name);
 
  private:
-  /// Guard for shutdown state.
-  absl::Mutex stop_mu_;
   // True once shutdown begins. Requests to execute new tasks will be rejected.
-  bool stopping_ ABSL_GUARDED_BY(stop_mu_) = false;
+  std::atomic<bool> stopping_ = false;
   /// Set up the configs for an actor.
   /// This should be called once for the actor creation task.
   void SetupActor(bool is_asyncio,