chore: import upstream snapshot with attribution

python/ray/tests/test_actor_lineage_reconstruction.py

@@ -0,0 +1,126 @@
+import gc
+import json
+import os
+import signal
+import sys
+
+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,
+)
+from ray.core.generated import common_pb2, gcs_pb2
+
+
+@pytest.mark.parametrize("deterministic_failure", RPC_FAILURE_TYPES)
+def test_actor_reconstruction_triggered_by_lineage_reconstruction(
+    monkeypatch, ray_start_cluster, deterministic_failure
+):
+    # Test the sequence of events:
+    # actor goes out of scope and killed
+    # -> lineage reconstruction triggered by object lost
+    # -> actor is restarted
+    # -> 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.
+    failure = RPC_FAILURE_MAP[deterministic_failure].copy()
+    failure["num_failures"] = 1
+    monkeypatch.setenv(
+        "RAY_testing_rpc_failure",
+        json.dumps(
+            {
+                "ray::rpc::ActorInfoGcsService.grpc_client.RestartActorForLineageReconstruction": failure,
+                "ray::rpc::ActorInfoGcsService.grpc_client.ReportActorOutOfScope": failure,
+            }
+        ),
+    )
+    cluster = ray_start_cluster
+    cluster.add_node(resources={"head": 1})
+    ray.init(address=cluster.address)
+    worker1 = cluster.add_node(resources={"worker": 1})
+
+    @ray.remote(
+        num_cpus=1, resources={"worker": 1}, max_restarts=-1, max_task_retries=-1
+    )
+    class Actor:
+        def ping(self):
+            return [1] * 1024 * 1024
+
+        def pid(self):
+            return os.getpid()
+
+    actor = Actor.remote()
+    actor_id = actor._actor_id
+
+    obj1 = actor.ping.remote()
+    os.kill(ray.get(actor.pid.remote()), signal.SIGKILL)
+
+    # obj2 should be ready after actor is restarted
+    obj2 = actor.ping.remote()
+
+    # Make the actor out of scope
+    actor = None
+
+    def verify1():
+        gc.collect()
+        actor_info = ray._private.state.state.get_actor_info(actor_id)
+        assert actor_info is not None
+        actor_info = gcs_pb2.ActorTableData.FromString(actor_info)
+        assert actor_info.state == gcs_pb2.ActorTableData.ActorState.DEAD
+        assert (
+            actor_info.death_cause.actor_died_error_context.reason
+            == common_pb2.ActorDiedErrorContext.Reason.OUT_OF_SCOPE
+        )
+        assert actor_info.num_restarts_due_to_lineage_reconstruction == 0
+        return True
+
+    wait_for_condition(lambda: verify1())
+
+    # objs will be lost and recovered
+    # during the process, actor will be reconstructured
+    # and dead again after lineage reconstruction finishes
+    cluster.remove_node(worker1)
+    cluster.add_node(resources={"worker": 1})
+
+    assert ray.get(obj1) == [1] * 1024 * 1024
+    assert ray.get(obj2) == [1] * 1024 * 1024
+
+    def verify2():
+        actor_info = ray._private.state.state.get_actor_info(actor_id)
+        assert actor_info is not None
+        actor_info = gcs_pb2.ActorTableData.FromString(actor_info)
+        assert actor_info.state == gcs_pb2.ActorTableData.ActorState.DEAD
+        assert (
+            actor_info.death_cause.actor_died_error_context.reason
+            == common_pb2.ActorDiedErrorContext.Reason.OUT_OF_SCOPE
+        )
+        # 1 restart recovers two objects
+        assert actor_info.num_restarts_due_to_lineage_reconstruction == 1
+        return True
+
+    wait_for_condition(lambda: verify2())
+
+    # actor can be permanently dead since no lineage reconstruction will happen
+    del obj1
+    del obj2
+
+    def verify3():
+        actor_info = ray._private.state.state.get_actor_info(actor_id)
+        assert actor_info is not None
+        actor_info = gcs_pb2.ActorTableData.FromString(actor_info)
+        assert actor_info.state == gcs_pb2.ActorTableData.ActorState.DEAD
+        assert (
+            actor_info.death_cause.actor_died_error_context.reason
+            == common_pb2.ActorDiedErrorContext.Reason.REF_DELETED
+        )
+        assert actor_info.num_restarts_due_to_lineage_reconstruction == 1
+        return True
+
+    wait_for_condition(lambda: verify3())
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))