chore: import upstream snapshot with attribution

python/ray/tests/test_raylet_fault_tolerance.py

@@ -0,0 +1,340 @@
+import json
+import os
+import sys
+
+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,
+)
+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,
+    PlacementGroupSchedulingStrategy,
+)
+
+import psutil
+
+
+@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].copy()
+    failure["num_failures"] = 1
+    monkeypatch.setenv(
+        "RAY_testing_rpc_failure",
+        json.dumps({"NodeManagerService.grpc_client.RequestWorkerLease": failure}),
+    )
+
+    @ray.remote
+    def simple_task_1():
+        return 0
+
+    @ray.remote
+    def simple_task_2():
+        return 1
+
+    # Spin up a two-node cluster where we're targeting scheduling on the
+    # remote node via NodeAffinitySchedulingStrategy to test remote RequestWorkerLease
+    # calls.
+    cluster = ray_start_cluster
+    remote_node = cluster.add_node(num_cpus=1)
+
+    result_ref1 = simple_task_1.options(
+        scheduling_strategy=NodeAffinitySchedulingStrategy(
+            node_id=remote_node.node_id, soft=False
+        )
+    ).remote()
+    result_ref2 = simple_task_2.options(
+        scheduling_strategy=NodeAffinitySchedulingStrategy(
+            node_id=remote_node.node_id, soft=False
+        )
+    ).remote()
+
+    assert ray.get([result_ref1, result_ref2]) == [0, 1]
+
+
+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",
+        json.dumps(
+            {
+                "NodeManagerService.grpc_client.DrainRaylet": {
+                    "num_failures": 1,
+                    "req_failure_prob": 100,
+                    "resp_failure_prob": 0,
+                    "in_flight_failure_prob": 0,
+                }
+            }
+        ),
+    )
+
+    cluster = ray_start_cluster
+    worker_node = cluster.add_node(num_cpus=1)
+    ray.init(address=cluster.address)
+
+    worker_node_id = worker_node.node_id
+
+    gcs_client = ray._raylet.GcsClient(address=cluster.address)
+
+    is_accepted = gcs_client.drain_node(
+        worker_node_id,
+        autoscaler_pb2.DrainNodeReason.DRAIN_NODE_REASON_IDLE_TERMINATION,
+        "Test drain",
+        0,
+    )
+    assert is_accepted
+
+    # After drain is accepted on an idle node since no tasks are running nor primary objects kept
+    # on that raylet, it should be marked idle and gracefully shut down.
+    def node_is_dead():
+        nodes = ray.nodes()
+        for node in nodes:
+            if node["NodeID"] == worker_node_id:
+                return not node["Alive"]
+        return True
+
+    wait_for_condition(node_is_dead, timeout=1)
+
+
+# Bundles can be leaked if the gcs dies before the RemovePlacementGroupBundles RPCs are
+# propagated to all the raylets. Since this is inherently racy, we block RemovePlacementGroupBundles RPCs
+# from ever succeeding to make this test deterministic.
+@pytest.fixture
+def inject_release_unused_bundles_rpc_failure(monkeypatch, request):
+    deterministic_failure = request.param
+    failure = RPC_FAILURE_MAP[deterministic_failure].copy()
+    failure["num_failures"] = 1
+    monkeypatch.setenv(
+        "RAY_testing_rpc_failure",
+        json.dumps(
+            {
+                "NodeManagerService.grpc_client.ReleaseUnusedBundles": failure,
+                "NodeManagerService.grpc_client.RemovePlacementGroupBundles": {
+                    "num_failures": -1,
+                    "req_failure_prob": 100,
+                    "resp_failure_prob": 0,
+                    "in_flight_failure_prob": 0,
+                },
+            }
+        ),
+    )
+
+
+@pytest.mark.parametrize(
+    "inject_release_unused_bundles_rpc_failure",
+    RPC_FAILURE_TYPES,
+    indirect=True,
+)
+@pytest.mark.parametrize(
+    "ray_start_cluster_head_with_external_redis",
+    [{"num_cpus": 1}],
+    indirect=True,
+)
+def test_release_unused_bundles_idempotent(
+    inject_release_unused_bundles_rpc_failure,
+    ray_start_cluster_head_with_external_redis,
+):
+    cluster = ray_start_cluster_head_with_external_redis
+
+    @ray.remote(num_cpus=1)
+    def task():
+        return "success"
+
+    pg = placement_group(name="test_pg", strategy="PACK", bundles=[{"CPU": 1}])
+
+    result_ref = task.options(
+        scheduling_strategy=PlacementGroupSchedulingStrategy(
+            placement_group=pg,
+            placement_group_bundle_index=0,
+        )
+    ).remote()
+    assert ray.get(result_ref) == "success"
+
+    # Remove the placement group. This will trigger RemovePlacementGroupBundles RPCs which need to be blocked
+    # for the placement group bundle to be leaked.
+    remove_placement_group(pg)
+
+    cluster.head_node.kill_gcs_server()
+    # ReleaseUnusedBundles only triggers after GCS restart to clean up potentially leaked bundles.
+    cluster.head_node.start_gcs_server()
+
+    # If the leaked bundle wasn't cleaned up, this task will hang due to resource unavailability
+    result = ray.get(task.remote())
+    assert result == "success"
+
+
+@pytest.fixture
+def inject_notify_gcs_restart_rpc_failure(monkeypatch, request):
+    deterministic_failure = request.param
+    failure = RPC_FAILURE_MAP[deterministic_failure].copy()
+    failure["num_failures"] = 1
+    monkeypatch.setenv(
+        "RAY_testing_rpc_failure",
+        json.dumps({"NodeManagerService.grpc_client.NotifyGCSRestart": failure}),
+    )
+
+
+@pytest.mark.parametrize(
+    "inject_notify_gcs_restart_rpc_failure",
+    RPC_FAILURE_TYPES,
+    indirect=True,
+)
+@pytest.mark.parametrize(
+    "ray_start_cluster_head_with_external_redis",
+    [
+        {
+            "_system_config": {
+                # Extending the fallback timeout to focus on death
+                # notification received from GCS_ACTOR_CHANNEL pubsub
+                "timeout_ms_task_wait_for_death_info": 10000,
+            }
+        }
+    ],
+    indirect=True,
+)
+def test_notify_gcs_restart_idempotent(
+    inject_notify_gcs_restart_rpc_failure,
+    ray_start_cluster_head_with_external_redis,
+):
+    cluster = ray_start_cluster_head_with_external_redis
+
+    @ray.remote(num_cpus=1, max_restarts=0)
+    class DummyActor:
+        def get_pid(self):
+            return psutil.Process().pid
+
+        def ping(self):
+            return "pong"
+
+    actor = DummyActor.remote()
+    ray.get(actor.ping.remote())
+    actor_pid = ray.get(actor.get_pid.remote())
+
+    cluster.head_node.kill_gcs_server()
+    cluster.head_node.start_gcs_server()
+
+    p = psutil.Process(actor_pid)
+    p.kill()
+
+    # If the actor death notification is not received from the GCS pubsub, this will timeout since
+    # the fallback via wait_for_death_info_tasks in the actor task submitter will never trigger
+    # since it's set to 10 seconds.
+    with pytest.raises(ray.exceptions.RayActorError):
+        ray.get(actor.ping.remote(), timeout=5)
+
+
+def test_kill_local_actor_rpc_retry_and_idempotency(monkeypatch, shutdown_only):
+    """Test that KillLocalActor RPC retries work correctly and guarantee actor death.
+    Not testing response since the actor is killed either way.
+    """
+
+    monkeypatch.setenv(
+        "RAY_testing_rpc_failure",
+        json.dumps(
+            {
+                "NodeManagerService.grpc_client.KillLocalActor": {
+                    "num_failures": 1,
+                    "req_failure_prob": 100,
+                    "resp_failure_prob": 0,
+                    "in_flight_failure_prob": 0,
+                }
+            }
+        ),
+    )
+
+    ray.init()
+
+    @ray.remote
+    class SimpleActor:
+        def ping(self):
+            return "pong"
+
+        def get_pid(self):
+            return os.getpid()
+
+    actor = SimpleActor.remote()
+
+    result = ray.get(actor.ping.remote())
+    assert result == "pong"
+
+    worker_pid = ray.get(actor.get_pid.remote())
+
+    # NOTE: checking the process is still alive rather than checking the actor state from the GCS
+    # since as long as KillActor is sent the GCS will mark the actor as dead even though it may not actually be
+    assert psutil.pid_exists(worker_pid)
+
+    ray.kill(actor)
+
+    def verify_process_killed():
+        return not psutil.pid_exists(worker_pid)
+
+    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].copy()
+    failure["num_failures"] = 1
+    monkeypatch.setenv(
+        "RAY_testing_rpc_failure",
+        json.dumps(
+            {
+                "NodeManagerService.grpc_client.CancelLocalTask": 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__]))