[GCS] global state query node info table from GCS. (#8498)

python/ray/tests/test_actor_failures.py

@@ -11,9 +11,15 @@ import ray
 import ray.ray_constants as ray_constants
 import ray.test_utils
 import ray.cluster_utils
-from ray.test_utils import (relevant_errors, wait_for_condition,
-                            wait_for_errors, wait_for_pid_to_exit,
-                            generate_internal_config_map)
+from ray.test_utils import (
+    relevant_errors,
+    wait_for_condition,
+    wait_for_errors,
+    wait_for_pid_to_exit,
+    generate_internal_config_map,
+    get_non_head_nodes,
+    get_other_nodes,
+)
 
 SIGKILL = signal.SIGKILL if sys.platform != "win32" else signal.SIGTERM
 
@@ -305,7 +311,7 @@ def test_actor_restart_on_node_failure(ray_start_cluster):
     ray.get(actor.ready.remote())
     results = [actor.increase.remote() for _ in range(100)]
     # Kill actor node, while the above task is still being executed.
-    cluster.remove_node(cluster.list_all_nodes()[-1])
+    cluster.remove_node(get_non_head_nodes(cluster)[-1])
     cluster.add_node(num_cpus=1, _internal_config=config)
     cluster.wait_for_nodes()
     # Check that none of the tasks failed and the actor is restarted.
@@ -821,7 +827,7 @@ def test_decorated_method(ray_start_regular):
 @pytest.mark.parametrize(
     "ray_start_cluster", [{
         "num_cpus": 1,
-        "num_nodes": 2,
+        "num_nodes": 3,
     }], indirect=True)
 def test_ray_wait_dead_actor(ray_start_cluster):
     """Tests that methods completed by dead actors are returned as ready"""
@@ -857,8 +863,8 @@ def test_ray_wait_dead_actor(ray_start_cluster):
         except ray.exceptions.RayActorError:
             return True
 
-    # Kill a node.
-    cluster.remove_node(cluster.list_all_nodes()[-1])
+    # Kill a node that must not be driver node or head node.
+    cluster.remove_node(get_other_nodes(cluster, exclude_head=True)[-1])
     # Repeatedly submit tasks and call ray.wait until the exception for the
     # dead actor is received.
     assert wait_for_condition(actor_dead)