Use Public State API for Fetching Actors

Instead of fetching actors from the internal state API / migrating it to
the common module, simply reference the public state API that relies on the dashboard
server.

Signed-off-by: Jason <jcarlson212@gmail.com>

Author: jcarlson212

python/ray/_common/state.py

@@ -548,11 +548,10 @@ def wait_for_num_actors(num_actors, state=None, timeout=10):
     while time.time() - start_time < timeout:
         if (
             len(
-                [
-                    _
-                    for _ in ray._common.state.actors().values()
-                    if state is None or _["State"] == state
-                ]
+                ray.util.state.list_actors(
+                    filters=("state", "=", state) if state else None,
+                    limit=num_actors,
+                )
             )
             >= num_actors
         ):
@@ -563,14 +562,14 @@ def wait_for_num_actors(num_actors, state=None, timeout=10):
 
 def kill_actor_and_wait_for_failure(actor, timeout=10, retry_interval_ms=100):
     actor_id = actor._actor_id.hex()
-    current_num_restarts = ray._common.state.actors(actor_id)["NumRestarts"]
+    current_num_restarts = ray.util.state.get_actor(actor_id)["num_restarts"]
     ray.kill(actor)
     start = time.time()
     while time.time() - start <= timeout:
-        actor_status = ray._common.state.actors(actor_id)
+        actor_state = ray.util.state.get_actor(actor_id)
         if (
-            actor_status["State"] == convert_actor_state(gcs_utils.ActorTableData.DEAD)
-            or actor_status["NumRestarts"] > current_num_restarts
+            actor_state["state"] == convert_actor_state(gcs_utils.ActorTableData.DEAD)
+            or actor_state["num_restarts"] > current_num_restarts
         ):
             return
         time.sleep(retry_interval_ms / 1000.0)

python/ray/_private/test_utils.py

@@ -2086,12 +2086,17 @@ def build_block_refs_by_node_id(
 
         def build_node_id_by_actor(actors: List[Any]) -> Dict[Any, str]:
             """Build a map from a actor to its node_id."""
-            actors_state = ray._common.state.actors()
+            actors_state = {
+                actor.actor_id: actor.node_id
+                for actor in ray.util.state.list_actors(
+                    limit=ray.util.state.summarize_actors()
+                    .get("cluster", {})
+                    .get("total_actors", 0)
+                    + 100  # fetch current actors. Some staleness is assumed fine.
+                )
+            }
             return {
-                actor: actors_state.get(actor._actor_id.hex(), {})
-                .get("Address", {})
-                .get("NodeID")
-                for actor in actors
+                actor: actors_state.get(actor._actor_id.hex(), {}) for actor in actors
             }
 
         # expected number of blocks to be allocated for each actor

python/ray/data/dataset.py

@@ -155,6 +155,11 @@ def test_equal_split_balanced_grid(ray_start_regular_shared_2_cpus):
                 _test_equal_split_balanced(block_sizes, num_splits)
 
 
+@pytest.mark.parametrize(
+    "ray_start_regular_shared_2_cpus",
+    [{"num_cpus": 1, "include_dashboard": True}],
+    indirect=True,
+)
 def test_split_small(ray_start_regular_shared_2_cpus):
     x = [Counter.remote() for _ in range(4)]
     data = ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"]
@@ -366,6 +371,11 @@ def test_split(ray_start_regular_shared_2_cpus):
     assert 190 == sum([dataset.sum("id") or 0 for dataset in datasets])
 
 
+@pytest.mark.parametrize(
+    "ray_start_regular_shared_2_cpus",
+    [{"num_cpus": 1, "include_dashboard": True}],
+    indirect=True,
+)
 def test_split_hints(ray_start_regular_shared_2_cpus):
     @ray.remote
     class Actor(object):
@@ -394,30 +404,45 @@ def assert_split_assignment(block_node_ids, actor_node_ids, expected_split_resul
         blocks = _ref_bundles_iterator_to_block_refs_list(bundles)
         assert len(block_node_ids) == len(blocks)
         actors = [Actor.remote() for i in range(len(actor_node_ids))]
-        with patch("ray.experimental.get_object_locations") as location_mock:
-            with patch("ray._common.state.actors") as state_mock:
-                block_locations = {}
-                for i, node_id in enumerate(block_node_ids):
-                    if node_id:
-                        block_locations[blocks[i]] = {"node_ids": [node_id]}
-                location_mock.return_value = block_locations
-
-                actor_state = {}
-                for i, node_id in enumerate(actor_node_ids):
-                    actor_state[actors[i]._actor_id.hex()] = {
-                        "Address": {"NodeID": node_id}
-                    }
-
-                state_mock.return_value = actor_state
-
-                datasets = ds.split(len(actors), locality_hints=actors)
-                assert len(datasets) == len(actors)
-                for i in range(len(actors)):
-                    assert {blocks[j] for j in expected_split_result[i]} == set(
-                        _ref_bundles_iterator_to_block_refs_list(
-                            datasets[i].iter_internal_ref_bundles()
+        with patch("ray.util.state.summarize_actors") as state_summary_mock:
+            state_summary_mock.return_value = {"cluster": {"total_actors": len(actors)}}
+            with patch("ray.experimental.get_object_locations") as location_mock:
+                with patch("ray.util.state.list_actors") as state_mock:
+                    block_locations = {}
+                    for i, node_id in enumerate(block_node_ids):
+                        if node_id:
+                            block_locations[blocks[i]] = {"node_ids": [node_id]}
+                    location_mock.return_value = block_locations
+
+                    class ActorStateMock:
+                        def __init__(self, actor_id, node_id):
+                            self._actor_id = actor_id
+                            self._node_id = node_id
+
+                        @property
+                        def actor_id(self):
+                            return self._actor_id
+
+                        @property
+                        def node_id(self):
+                            return self._node_id
+
+                    actor_state = []
+                    for i, node_id in enumerate(actor_node_ids):
+                        actor_state.append(
+                            ActorStateMock(actors[i]._actor_id.hex(), node_id)
+                        )
+
+                    state_mock.return_value = actor_state
+
+                    datasets = ds.split(len(actors), locality_hints=actors)
+                    assert len(datasets) == len(actors)
+                    for i in range(len(actors)):
+                        assert {blocks[j] for j in expected_split_result[i]} == set(
+                            _ref_bundles_iterator_to_block_refs_list(
+                                datasets[i].iter_internal_ref_bundles()
+                            )
                         )
-                    )
 
     assert_split_assignment(
         ["node2", "node1", "node1"], ["node1", "node2"], [[1, 2], [0]]

python/ray/data/tests/test_split.py

@@ -405,8 +405,8 @@ def was_current_actor_reconstructed(self):
         assert (
             not self.actor_id.is_nil()
         ), "This method should't be called inside Ray tasks."
-        actor_info = ray._common.state.actors(self.actor_id.hex())
-        return actor_info and actor_info["NumRestarts"] != 0
+        actors = ray.util.state.list_actors(filters=[("id", "=", self.actor_id.hex())])
+        return any(act.num_restarts or 0 for act in actors)
 
     @property
     @Deprecated(message="Use get_placement_group_id() instead", warning=True)

python/ray/runtime_context.py

@@ -526,6 +526,11 @@ def foobar(self):
     assert ray.get(detached_actor.foobar.remote()) == ["bar", "bar"]
 
 
+@pytest.mark.parametrize(
+    "ray_start_regular",
+    [{"num_cpus": 1, "include_dashboard": True}],
+    indirect=True,
+)
 def test_detached_actor_cleanup(ray_start_regular):
     @ray.remote
     class DetachedActor:
@@ -546,15 +551,11 @@ def create_and_kill_actor(actor_name):
         detached_actor = ray.get_actor(dup_actor_name)
         ray.kill(detached_actor)
         # Wait until actor dies.
-        actor_status = ray._common.state.actors(actor_id=detached_actor._actor_id.hex())
+        actor_status = ray.util.state.get_actor(id=detached_actor._actor_id.hex())
         max_wait_time = 10
         wait_time = 0
-        while actor_status["State"] != convert_actor_state(
-            gcs_utils.ActorTableData.DEAD
-        ):
-            actor_status = ray._common.state.actors(
-                actor_id=detached_actor._actor_id.hex()
-            )
+        while actor_status.state != convert_actor_state(gcs_utils.ActorTableData.DEAD):
+            actor_status = ray.util.state.get_actor(id=detached_actor._actor_id.hex())
             time.sleep(1.0)
             wait_time += 1
             if wait_time >= max_wait_time:
@@ -586,11 +587,11 @@ def ping(self):
 assert ray.get(detached_actor.ping.remote()) == "pong"
 ray.kill(detached_actor)
 # Wait until actor dies.
-actor_status = ray._common.state.actors(actor_id=detached_actor._actor_id.hex())
+actor_status = ray.util.state.get_actor(id=detached_actor._actor_id.hex())
 max_wait_time = 10
 wait_time = 0
-while actor_status["State"] != convert_actor_state(gcs_utils.ActorTableData.DEAD): # noqa
-    actor_status = ray._common.state.actors(actor_id=detached_actor._actor_id.hex())
+while actor_status.state != convert_actor_state(gcs_utils.ActorTableData.DEAD): # noqa
+    actor_status = ray.util.state.get_actor(id=detached_actor._actor_id.hex())
     time.sleep(1.0)
     wait_time += 1
     if wait_time >= max_wait_time:
@@ -640,7 +641,14 @@ def hi(self):
 
 @pytest.mark.parametrize(
     "ray_start_cluster",
-    [{"num_cpus": 3, "num_nodes": 1, "resources": {"first_node": 5}}],
+    [
+        {
+            "num_cpus": 3,
+            "num_nodes": 1,
+            "resources": {"first_node": 5},
+            "include_dashboard": True,
+        }
+    ],
     indirect=True,
 )
 def test_detached_actor_cleanup_due_to_failure(ray_start_cluster):
@@ -661,13 +669,13 @@ def kill_itself(self):
     node_failure_actor_name = "node_failure_actor_name"
 
     def wait_until_actor_dead(handle):
-        actor_status = ray._common.state.actors(actor_id=handle._actor_id.hex())
+        actor_status = ray.util.state.get_actor(id=handle._actor_id.hex())
         max_wait_time = 10
         wait_time = 0
-        while actor_status["State"] != convert_actor_state(
+        while actor_status["state"] != convert_actor_state(
             gcs_utils.ActorTableData.DEAD
         ):
-            actor_status = ray._common.state.actors(actor_id=handle._actor_id.hex())
+            actor_status = ray.util.state.get_actor(id=handle._actor_id.hex())
             time.sleep(1.0)
             wait_time += 1
             if wait_time >= max_wait_time:
@@ -991,6 +999,8 @@ def condition2():
 
 
 def test_actor_timestamps(ray_start_regular):
+    """Tests the internal actor state representation's timestamps. The dashboard APIs don't support this yet."""
+
     @ray.remote
     class Foo:
         def get_id(self):
@@ -1003,11 +1013,12 @@ def graceful_exit():
         actor = Foo.remote()
         actor_id = ray.get(actor.get_id.remote())
 
-        state_after_starting = ray._common.state.actors()[actor_id]
+        state_after_starting = ray._private.state.actors()[actor_id]
+
         time.sleep(1)
         del actor
         time.sleep(1)
-        state_after_ending = ray._common.state.actors()[actor_id]
+        state_after_ending = ray._private.state.actors()[actor_id]
 
         assert state_after_starting["StartTime"] == state_after_ending["StartTime"]
         start_time = state_after_ending["StartTime"]
@@ -1018,11 +1029,11 @@ def not_graceful_exit():
         actor = Foo.remote()
         actor_id = ray.get(actor.get_id.remote())
 
-        state_after_starting = ray._common.state.actors()[actor_id]
+        state_after_starting = ray._private.state.actors()[actor_id]
         time.sleep(1)
         actor.kill_self.remote()
         time.sleep(1)
-        state_after_ending = ray._common.state.actors()[actor_id]
+        state_after_ending = ray._private.state.actors()[actor_id]
 
         assert state_after_starting["StartTime"] == state_after_ending["StartTime"]
 
@@ -1034,13 +1045,13 @@ def restarted():
         actor = Foo.options(max_restarts=1, max_task_retries=-1).remote()
         actor_id = ray.get(actor.get_id.remote())
 
-        state_after_starting = ray._common.state.actors()[actor_id]
+        state_after_starting = ray._private.state.actors()[actor_id]
         time.sleep(1)
         actor.kill_self.remote()
         time.sleep(1)
         actor.kill_self.remote()
         time.sleep(1)
-        state_after_ending = ray._common.state.actors()[actor_id]
+        state_after_ending = ray._private.state.actors()[actor_id]
 
         assert state_after_starting["StartTime"] == state_after_ending["StartTime"]
 

python/ray/tests/test_actor_advanced.py

@@ -42,12 +42,18 @@ def test_actor_creation_num_cpus(ray_start_2_cpus):
     ag.shutdown()
 
 
+@pytest.mark.parametrize(
+    "ray_start_2_cpus",
+    [{"num_cpus": 1, "include_dashboard": True}],
+    indirect=True,
+)
 def test_actor_shutdown(ray_start_2_cpus):
     assert ray.available_resources()["CPU"] == 2
     ag = ActorGroup(actor_cls=DummyActor, num_actors=2)
     time.sleep(1)
     assert "CPU" not in ray.available_resources()
-    assert len(ray._common.state.actors()) == 2
+
+    assert len(ray.util.state.list_actors()) == 2
     ag.shutdown()
     time.sleep(1)
     assert ray.available_resources()["CPU"] == 2

python/ray/tests/test_actor_group.py

@@ -20,6 +20,11 @@
 logger = logging.getLogger(__name__)
 
 
+@pytest.mark.parametrize(
+    "shutdown_only",
+    [{"num_cpus": 1, "include_dashboard": True}],
+    indirect=True,
+)
 def test_global_state_api(shutdown_only):
 
     ray.init(num_cpus=5, num_gpus=3, resources={"CustomResource": 1})
@@ -47,15 +52,14 @@ def __init__(self):
     # Wait for actor to be created
     wait_for_num_actors(1)
 
-    actor_table = ray._common.state.actors()
+    actor_table = (
+        ray.util.state.list_actors()
+    )  # should be using this API now for fetching actors
     assert len(actor_table) == 1
 
     (actor_info,) = actor_table.values()
-    assert actor_info["JobID"] == job_id.hex()
-    assert actor_info["Name"] == "test_actor"
-    assert "IPAddress" in actor_info["Address"]
-    assert "IPAddress" in actor_info["OwnerAddress"]
-    assert actor_info["Address"]["Port"] != actor_info["OwnerAddress"]["Port"]
+    assert actor_info.job_id == job_id.hex()
+    assert actor_info.name == "test_actor"
 
     job_table = ray._private.state.jobs()