[core] Adding option for in flight rpc failure injection (#58512)

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

Author: dayshah

python/ray/_private/test_utils.py

@@ -2096,3 +2096,12 @@ def _execute_command_on_node(command: str, node_ip: str):
     except subprocess.CalledProcessError as e:
         print("Exit code:", e.returncode)
         print("Stderr:", e.stderr)
+
+
+RPC_FAILURE_MAP = {
+    "request": "100:0:0",
+    "response": "0:100:0",
+    "in_flight": "0:0:100",
+}
+
+RPC_FAILURE_TYPES = list(RPC_FAILURE_MAP.keys())

python/ray/dashboard/modules/job/tests/test_job_manager.py

@@ -361,7 +361,7 @@ async def test_runtime_env_setup_logged_to_job_driver_logs(
         {
             "cmd": "ray start --head",
             "env": {
-                "RAY_testing_rpc_failure": "ray::rpc::InternalKVGcsService.grpc_client.InternalKVGet=2:50:50,CoreWorkerService.grpc_client.PushTask=3:50:50"
+                "RAY_testing_rpc_failure": "ray::rpc::InternalKVGcsService.grpc_client.InternalKVGet=3:33:33:33,CoreWorkerService.grpc_client.PushTask=3:33:33:33"
             },
         },
     ],

python/ray/tests/test_actor_lineage_reconstruction.py

@@ -6,11 +6,15 @@
 import pytest
 
 import ray
-from ray._common.test_utils import wait_for_condition
+from ray._private.test_utils import (
+    RPC_FAILURE_MAP,
+    RPC_FAILURE_TYPES,
+    wait_for_condition,
+)
 from ray.core.generated import common_pb2, gcs_pb2
 
 
-@pytest.mark.parametrize("deterministic_failure", ["request", "response"])
+@pytest.mark.parametrize("deterministic_failure", RPC_FAILURE_TYPES)
 def test_actor_reconstruction_triggered_by_lineage_reconstruction(
     monkeypatch, ray_start_cluster, deterministic_failure
 ):
@@ -21,11 +25,11 @@ def test_actor_reconstruction_triggered_by_lineage_reconstruction(
     # -> actor goes out of scope again after lineage reconstruction is done
     # -> actor is permanently dead when there is no reference.
     # This test also injects network failure to make sure relevant rpcs are retried.
-    chaos_failure = "100:0" if deterministic_failure == "request" else "0:100"
+    failure = RPC_FAILURE_MAP[deterministic_failure]
     monkeypatch.setenv(
         "RAY_testing_rpc_failure",
-        f"ray::rpc::ActorInfoGcsService.grpc_client.RestartActorForLineageReconstruction=1:{chaos_failure},"
-        f"ray::rpc::ActorInfoGcsService.grpc_client.ReportActorOutOfScope=1:{chaos_failure}",
+        f"ray::rpc::ActorInfoGcsService.grpc_client.RestartActorForLineageReconstruction=1:{failure},"
+        f"ray::rpc::ActorInfoGcsService.grpc_client.ReportActorOutOfScope=1:{failure}",
     )
     cluster = ray_start_cluster
     cluster.add_node(resources={"head": 1})

python/ray/tests/test_core_worker_fault_tolerance.py

@@ -4,7 +4,12 @@
 import pytest
 
 import ray
-from ray._common.test_utils import SignalActor, wait_for_condition
+from ray._common.test_utils import SignalActor
+from ray._private.test_utils import (
+    RPC_FAILURE_MAP,
+    RPC_FAILURE_TYPES,
+    wait_for_condition,
+)
 from ray.core.generated import common_pb2, gcs_pb2
 from ray.exceptions import GetTimeoutError, TaskCancelledError
 from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
@@ -14,18 +19,18 @@
     "allow_out_of_order_execution",
     [True, False],
 )
-@pytest.mark.parametrize("deterministic_failure", ["request", "response"])
+@pytest.mark.parametrize("deterministic_failure", RPC_FAILURE_TYPES)
 def test_push_actor_task_failure(
     monkeypatch,
     ray_start_cluster,
     allow_out_of_order_execution: bool,
     deterministic_failure: str,
 ):
     with monkeypatch.context() as m:
+        failure = RPC_FAILURE_MAP[deterministic_failure]
         m.setenv(
             "RAY_testing_rpc_failure",
-            "CoreWorkerService.grpc_client.PushTask=2:"
-            + ("100:0" if deterministic_failure == "request" else "0:100"),
+            f"CoreWorkerService.grpc_client.PushTask=2:{failure}",
         )
         m.setenv("RAY_actor_scheduling_queue_max_reorder_wait_seconds", "0")
         cluster = ray_start_cluster
@@ -47,15 +52,15 @@ def echo(self, value):
         assert ray.get(refs) == list(range(10))
 
 
-@pytest.mark.parametrize("deterministic_failure", ["request", "response"])
+@pytest.mark.parametrize("deterministic_failure", RPC_FAILURE_TYPES)
 def test_update_object_location_batch_failure(
     monkeypatch, ray_start_cluster, deterministic_failure
 ):
     with monkeypatch.context() as m:
+        failure = RPC_FAILURE_MAP[deterministic_failure]
         m.setenv(
             "RAY_testing_rpc_failure",
-            "CoreWorkerService.grpc_client.UpdateObjectLocationBatch=1:"
-            + ("100:0" if deterministic_failure == "request" else "0:100"),
+            f"CoreWorkerService.grpc_client.UpdateObjectLocationBatch=1:{failure}",
         )
         cluster = ray_start_cluster
         head_node_id = cluster.add_node(
@@ -85,7 +90,7 @@ def consume_large_object(obj):
         assert ray.get(consume_ref, timeout=10) > 0
 
 
-@pytest.mark.parametrize("deterministic_failure", ["request", "response"])
+@pytest.mark.parametrize("deterministic_failure", RPC_FAILURE_TYPES)
 def test_get_object_status_rpc_retry_and_idempotency(
     monkeypatch, shutdown_only, deterministic_failure
 ):
@@ -94,11 +99,10 @@ def test_get_object_status_rpc_retry_and_idempotency(
     Cross_worker_access_task triggers GetObjectStatus because it does
     not own objects and needs to request it from the driver.
     """
-
+    failure = RPC_FAILURE_MAP[deterministic_failure]
     monkeypatch.setenv(
         "RAY_testing_rpc_failure",
-        "CoreWorkerService.grpc_client.GetObjectStatus=1:"
-        + ("100:0" if deterministic_failure == "request" else "0:100"),
+        f"CoreWorkerService.grpc_client.GetObjectStatus=1:{failure}",
     )
 
     ray.init()
@@ -118,7 +122,7 @@ def cross_worker_access_task(objects):
     assert final_result == [0, 2, 4, 6, 8]
 
 
-@pytest.mark.parametrize("deterministic_failure", ["request", "response"])
+@pytest.mark.parametrize("deterministic_failure", RPC_FAILURE_TYPES)
 def test_wait_for_actor_ref_deleted_rpc_retry_and_idempotency(
     monkeypatch, shutdown_only, deterministic_failure
 ):
@@ -127,11 +131,10 @@ def test_wait_for_actor_ref_deleted_rpc_retry_and_idempotency(
     The GCS actor manager will trigger this RPC during actor initialization
     to monitor when the actor handles have gone out of scope and the actor should be destroyed.
     """
-
+    failure = RPC_FAILURE_MAP[deterministic_failure]
     monkeypatch.setenv(
         "RAY_testing_rpc_failure",
-        "CoreWorkerService.grpc_client.WaitForActorRefDeleted=1:"
-        + ("100:0" if deterministic_failure == "request" else "0:100"),
+        f"CoreWorkerService.grpc_client.WaitForActorRefDeleted=1:{failure}",
     )
 
     ray.init()
@@ -166,15 +169,17 @@ def verify_actor_ref_deleted():
 @pytest.fixture
 def inject_cancel_remote_task_rpc_failure(monkeypatch, request):
     deterministic_failure = request.param
+    failure = RPC_FAILURE_MAP[deterministic_failure]
     monkeypatch.setenv(
         "RAY_testing_rpc_failure",
-        "CoreWorkerService.grpc_client.CancelRemoteTask=1:"
-        + ("100:0" if deterministic_failure == "request" else "0:100"),
+        f"CoreWorkerService.grpc_client.CancelRemoteTask=1:{failure}",
     )
 
 
 @pytest.mark.parametrize(
-    "inject_cancel_remote_task_rpc_failure", ["request", "response"], indirect=True
+    "inject_cancel_remote_task_rpc_failure",
+    RPC_FAILURE_TYPES,
+    indirect=True,
 )
 def test_cancel_remote_task_rpc_retry_and_idempotency(
     inject_cancel_remote_task_rpc_failure, ray_start_cluster
@@ -208,7 +213,7 @@ def remote_wait(sg):
 def test_double_borrowing_with_rpc_failure(monkeypatch, shutdown_only):
     """Regression test for https://github.com/ray-project/ray/issues/57997"""
     monkeypatch.setenv(
-        "RAY_testing_rpc_failure", "CoreWorkerService.grpc_client.PushTask=3:0:100"
+        "RAY_testing_rpc_failure", "CoreWorkerService.grpc_client.PushTask=3:0:100:0"
     )
 
     ray.init()

python/ray/tests/test_gcs_fault_tolerance.py

@@ -1258,11 +1258,10 @@ def test_mark_job_finished_rpc_retry_and_idempotency(shutdown_only, monkeypatch)
     Uses RPC failure injection to simulate network retry scenarios.
     """
     # Inject RPC failures for MarkJobFinished - simulate network failures
-    # Format: method_name=max_failures:request_failure_prob:response_failure_prob
     # We inject request failures to force retries and test idempotency
     monkeypatch.setenv(
         "RAY_testing_rpc_failure",
-        "ray::rpc::JobInfoGcsService.grpc_client.MarkJobFinished=3:50:0",
+        "ray::rpc::JobInfoGcsService.grpc_client.MarkJobFinished=3:50:0:0",
     )
 
     ray.init(num_cpus=1)

python/ray/tests/test_gcs_utils.py

@@ -109,8 +109,8 @@ def test_kv_timeout(ray_start_regular):
 def test_kv_transient_network_error(shutdown_only, monkeypatch):
     monkeypatch.setenv(
         "RAY_testing_rpc_failure",
-        "ray::rpc::InternalKVGcsService.grpc_client.InternalKVGet=5:25:25,"
-        "ray::rpc::InternalKVGcsService.grpc_client.InternalKVPut=5:25:25",
+        "ray::rpc::InternalKVGcsService.grpc_client.InternalKVGet=5:25:25:25,"
+        "ray::rpc::InternalKVGcsService.grpc_client.InternalKVPut=5:25:25:25",
     )
     ray.init()
     gcs_address = ray._private.worker.global_worker.gcs_client.address

python/ray/tests/test_object_manager_fault_tolerance.py

@@ -5,18 +5,22 @@
 
 import ray
 from ray._private.internal_api import get_memory_info_reply, get_state_from_address
-from ray._private.test_utils import wait_for_condition
+from ray._private.test_utils import (
+    RPC_FAILURE_MAP,
+    RPC_FAILURE_TYPES,
+    wait_for_condition,
+)
 from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
 
 
-@pytest.mark.parametrize("deterministic_failure", ["request", "response"])
+@pytest.mark.parametrize("deterministic_failure", RPC_FAILURE_TYPES)
 def test_free_objects_idempotent(
     monkeypatch, shutdown_only, deterministic_failure, ray_start_cluster
 ):
+    failure = RPC_FAILURE_MAP[deterministic_failure]
     monkeypatch.setenv(
         "RAY_testing_rpc_failure",
-        "ObjectManagerService.grpc_client.FreeObjects=1:"
-        + ("100:0" if deterministic_failure == "request" else "0:100"),
+        f"ObjectManagerService.grpc_client.FreeObjects=1:{failure}",
     )
 
     @ray.remote

python/ray/tests/test_raylet_fault_tolerance.py

@@ -4,7 +4,11 @@
 import pytest
 
 import ray
-from ray._private.test_utils import 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.util.placement_group import placement_group, remove_placement_group
 from ray.util.scheduling_strategies import (
@@ -15,14 +19,14 @@
 import psutil
 
 
-@pytest.mark.parametrize("deterministic_failure", ["request", "response"])
+@pytest.mark.parametrize("deterministic_failure", RPC_FAILURE_TYPES)
 def test_request_worker_lease_idempotent(
     monkeypatch, shutdown_only, deterministic_failure, ray_start_cluster
 ):
+    failure = RPC_FAILURE_MAP[deterministic_failure]
     monkeypatch.setenv(
         "RAY_testing_rpc_failure",
-        "NodeManagerService.grpc_client.RequestWorkerLease=1:"
-        + ("100:0" if deterministic_failure == "request" else "0:100"),
+        f"NodeManagerService.grpc_client.RequestWorkerLease=1:{failure}",
     )
 
     @ray.remote
@@ -57,7 +61,7 @@ def test_drain_node_idempotent(monkeypatch, shutdown_only, ray_start_cluster):
     # NOTE: not testing response failure since the node is already marked as draining and shuts down gracefully.
     monkeypatch.setenv(
         "RAY_testing_rpc_failure",
-        "NodeManagerService.grpc_client.DrainRaylet=1:100:0",
+        "NodeManagerService.grpc_client.DrainRaylet=1:100:0:0",
     )
 
     cluster = ray_start_cluster
@@ -94,16 +98,18 @@ def node_is_dead():
 @pytest.fixture
 def inject_release_unused_bundles_rpc_failure(monkeypatch, request):
     deterministic_failure = request.param
+    failure = RPC_FAILURE_MAP[deterministic_failure]
     monkeypatch.setenv(
         "RAY_testing_rpc_failure",
-        "NodeManagerService.grpc_client.ReleaseUnusedBundles=1:"
-        + ("100:0" if deterministic_failure == "request" else "0:100")
-        + ",NodeManagerService.grpc_client.CancelResourceReserve=-1:100:0",
+        f"NodeManagerService.grpc_client.ReleaseUnusedBundles=1:{failure}"
+        + ",NodeManagerService.grpc_client.CancelResourceReserve=-1:100:0:0",
     )
 
 
 @pytest.mark.parametrize(
-    "inject_release_unused_bundles_rpc_failure", ["request", "response"], indirect=True
+    "inject_release_unused_bundles_rpc_failure",
+    RPC_FAILURE_TYPES,
+    indirect=True,
 )
 @pytest.mark.parametrize(
     "ray_start_cluster_head_with_external_redis",
@@ -146,15 +152,17 @@ def task():
 @pytest.fixture
 def inject_notify_gcs_restart_rpc_failure(monkeypatch, request):
     deterministic_failure = request.param
+    failure = RPC_FAILURE_MAP[deterministic_failure]
     monkeypatch.setenv(
         "RAY_testing_rpc_failure",
-        "NodeManagerService.grpc_client.NotifyGCSRestart=1:"
-        + ("100:0" if deterministic_failure == "request" else "0:100"),
+        f"NodeManagerService.grpc_client.NotifyGCSRestart=1:{failure}",
     )
 
 
 @pytest.mark.parametrize(
-    "inject_notify_gcs_restart_rpc_failure", ["request", "response"], indirect=True
+    "inject_notify_gcs_restart_rpc_failure",
+    RPC_FAILURE_TYPES,
+    indirect=True,
 )
 @pytest.mark.parametrize(
     "ray_start_cluster_head_with_external_redis",
@@ -207,7 +215,7 @@ def test_kill_local_actor_rpc_retry_and_idempotency(monkeypatch, shutdown_only):
 
     monkeypatch.setenv(
         "RAY_testing_rpc_failure",
-        "NodeManagerService.grpc_client.KillLocalActor=1:100:0",
+        "NodeManagerService.grpc_client.KillLocalActor=1:100:0:0",
     )
 
     ray.init()

python/ray/tests/test_streaming_generator_4.py

@@ -190,7 +190,7 @@ def test_many_tasks_lineage_reconstruction_mini_stress_test(
         )
         m.setenv(
             "RAY_testing_rpc_failure",
-            "CoreWorkerService.grpc_client.ReportGeneratorItemReturns=5:25:25",
+            "CoreWorkerService.grpc_client.ReportGeneratorItemReturns=5:25:25:25",
         )
         cluster = ray_start_cluster
         cluster.add_node(

src/ray/common/ray_config_def.h

@@ -851,32 +851,37 @@ RAY_CONFIG(std::string, REDIS_CLIENT_KEY, "")
 RAY_CONFIG(std::string, REDIS_SERVER_NAME, "")
 
 /// grpc delay testing flags
-///  To use this, simply do
+///  To use this,
 ///      export RAY_testing_asio_delay_us="method1=min_val:max_val,method2=20:100"
 //  The delay is a random number between the interval. If method equals '*',
 //  it will apply to all methods.
 RAY_CONFIG(std::string, testing_asio_delay_us, "")
 
-/// To use this, simply do
-///     export
-///     RAY_testing_rpc_failure="method1=max_num_failures:req_failure_prob:resp_failure_prob,method2=max_num_failures:req_failure_prob:resp_failure_prob"
+/// To use this,
+///  export
+///  RAY_testing_rpc_failure="method1=max_num_failures:req_failure_prob:resp_failure_prob:in_flight_failure_prob,method2=max_num_failures:req_failure_prob:resp_failure_prob:in_flight_failure_prob"
 /// If you want to test all rpc failures you can use * as the method name and you can set
 /// -1 max_num_failures to have unlimited failures.
-/// Ex. unlimited failures for all rpc's with 25% request failures and 50% response
-/// failures.
-///     export RAY_testing_rpc_failure="*=-1:25:50"
+/// Ex. unlimited failures for all rpc's with 25% request failures, 50% response
+/// failures, and 10% in-flight failures.
+///     export RAY_testing_rpc_failure="*=-1:25:50:10"
 /// NOTE: Setting the wildcard will override any configuration for other methods.
 ///
-/// You can also provide an optional fourth and/or fifth parameter to specify that there
-/// should be at least a certain amount of failures on the request and response. The
-/// fourth parameter is for the request and the fifth parameter is for the response. By
-/// default these are set to 0, but by setting them to positive values it guarantees
-/// that the first X request RPCs will fail, followed by Y response RPCs that will fail.
+/// You can also provide an optional fifth, sixth, and/or seventh parameter to specify
+/// that there should be at least a certain amount of failures.
+//  The 5th parameter is for request failures.
+//  The 6th parameter is for response failures.
+//  The 7th parameter is for in-flight failures.
+/// By default these are set to 0, but by setting them to positive values it guarantees
+/// that the first X request RPCs will fail, followed by Y response RPCs that will fail,
+/// followed by Z in-flight RPCs that will fail.
 /// Afterwards, it will revert to the probabilistic failures. You can combine this with
 /// the wildcard so that each RPC method will have the same lower bounds applied.
-/// Ex. unlimited failures for all rpc's with 25% request failures and 50% response
-/// failures with at least 2 request failures and 3 response failures.
-///     export RAY_testing_rpc_failure="*=-1:25:50:2:3"
+///
+/// Ex. unlimited failures for all rpc's with 25% request failures, 50% response failures,
+/// and 10% in-flight failures with at least 2 request failures, 3 response failures, and
+/// 1 in-flight failure.
+///     export RAY_testing_rpc_failure="*=-1:25:50:10:2:3:1"
 RAY_CONFIG(std::string, testing_rpc_failure, "")
 /// If this is set, when injecting RPC failures, we'll check if the server and client have
 /// the same address. If they do, we won't inject the failure.