[serve] Remove usage of ray._private.state (ray-project#54140)

Replace with `ray.util.state`.

There is one remaining:
```
(ray) eoakes@eoakes-CQTYX0Y0RV serve % grep -r "_private.state" *
_private/cluster_node_info_cache.py:            ray._private.state.available_resources_per_node()
```

---------

Signed-off-by: Edward Oakes <[email protected]>
Signed-off-by: doyoung <[email protected]>

Author: edoakes

python/ray/serve/tests/conftest.py

@@ -52,7 +52,7 @@ def ray_shutdown():
 @pytest.fixture
 def ray_cluster():
     cluster = Cluster()
-    yield Cluster()
+    yield cluster
     serve.shutdown()
     ray.shutdown()
     cluster.shutdown()

python/ray/serve/tests/test_cluster.py

@@ -22,6 +22,7 @@
 from ray.serve.context import _get_global_client
 from ray.serve.handle import DeploymentHandle
 from ray.serve.schema import ServeDeploySchema
+from ray.util.state import list_actors
 
 
 def get_pids(expected, deployment_name="D", app_name="default", timeout=30):
@@ -220,14 +221,13 @@ def test_intelligent_scale_down(ray_cluster):
     client = _get_global_client()
 
     def get_actor_distributions():
-        actors = ray._private.state.actors()
         node_to_actors = defaultdict(list)
-        for actor in actors.values():
-            if "ServeReplica" not in actor["ActorClassName"]:
+        for actor in list_actors(
+            address=cluster.address, filters=[("STATE", "=", "ALIVE")]
+        ):
+            if "ServeReplica" not in actor.class_name:
                 continue
-            if actor["State"] != "ALIVE":
-                continue
-            node_to_actors[actor["Address"]["NodeID"]].append(actor)
+            node_to_actors[actor.node_id].append(actor)
 
         return set(map(len, node_to_actors.values()))
 
@@ -446,19 +446,22 @@ def __call__(self):
 
         # Ensure worker node has both replicas.
         def check_replicas_on_worker_nodes():
-            _actors = ray._private.state.actors().values()
-            replica_nodes = [
-                a["Address"]["NodeID"]
-                for a in _actors
-                if a["ActorClassName"].startswith("ServeReplica")
-            ]
-            return len(set(replica_nodes)) == 1
+            return (
+                len(
+                    {
+                        a.node_id
+                        for a in list_actors(address=cluster.address)
+                        if a.class_name.startswith("ServeReplica")
+                    }
+                )
+                == 1
+            )
 
         wait_for_condition(check_replicas_on_worker_nodes)
 
         # Ensure total actors of 2 proxies, 1 controller, and 2 replicas,
         # and 2 nodes exist.
-        wait_for_condition(lambda: len(ray._private.state.actors()) == 5)
+        wait_for_condition(lambda: len(list_actors(address=cluster.address)) == 5)
         assert len(ray.nodes()) == 2
 
         # Ensure `/-/healthz` and `/-/routes` return 200 and expected responses
@@ -490,21 +493,12 @@ def check_request(url: str, expected_code: int, expected_text: str):
         # replicas on all nodes.
         serve.delete(name=SERVE_DEFAULT_APP_NAME)
 
-        def _check():
-            _actors = ray._private.state.actors().values()
-            return (
-                len(
-                    list(
-                        filter(
-                            lambda a: a["State"] == "ALIVE",
-                            _actors,
-                        )
-                    )
-                )
-                == 3
+        wait_for_condition(
+            lambda: len(
+                list_actors(address=cluster.address, filters=[("STATE", "=", "ALIVE")])
             )
-
-        wait_for_condition(_check)
+            == 3,
+        )
 
         # Ensure head node `/-/healthz` and `/-/routes` continue to
         # return 200 and expected responses. Also, the worker node

python/ray/serve/tests/test_grpc.py

@@ -10,7 +10,6 @@
 import ray
 from ray import serve
 from ray._common.test_utils import SignalActor, wait_for_condition
-from ray.cluster_utils import Cluster
 from ray.serve._private.constants import SERVE_NAMESPACE
 from ray.serve._private.test_utils import (
     ping_fruit_stand,
@@ -26,6 +25,7 @@
 from ray.serve.generated import serve_pb2, serve_pb2_grpc
 from ray.serve.grpc_util import RayServegRPCContext
 from ray.serve.tests.test_config_files.grpc_deployment import g, g2
+from ray.util.state import list_actors
 
 
 def test_serving_request_through_grpc_proxy(ray_cluster):
@@ -246,7 +246,7 @@ def test_grpc_proxy_on_draining_nodes(ray_cluster):
     os.environ["TEST_WORKER_NODE_GRPC_PORT"] = str(worker_node_grpc_port)
 
     # Set up a cluster with 2 nodes.
-    cluster = Cluster()
+    cluster = ray_cluster
     cluster.add_node(num_cpus=0)
     cluster.add_node(num_cpus=2)
     cluster.wait_for_nodes()
@@ -277,18 +277,21 @@ def __call__(self):
 
     # Ensure worker node has both replicas.
     def check_replicas_on_worker_nodes():
-        _actors = ray._private.state.actors().values()
-        replica_nodes = [
-            a["Address"]["NodeID"]
-            for a in _actors
-            if a["ActorClassName"].startswith("ServeReplica")
-        ]
-        return len(set(replica_nodes)) == 1
+        return (
+            len(
+                {
+                    a.node_id
+                    for a in list_actors(address=cluster.address)
+                    if a.class_name.startswith("ServeReplica")
+                }
+            )
+            == 1
+        )
 
     wait_for_condition(check_replicas_on_worker_nodes)
 
     # Ensure total actors of 2 proxies, 1 controller, and 2 replicas, and 2 nodes exist.
-    wait_for_condition(lambda: len(ray._private.state.actors()) == 5)
+    wait_for_condition(lambda: len(list_actors(address=cluster.address)) == 5)
     assert len(ray.nodes()) == 2
 
     # Set up gRPC channels.
@@ -318,21 +321,12 @@ def check_replicas_on_worker_nodes():
     # replicas on all nodes.
     serve.delete(name=app_name)
 
-    def _check():
-        _actors = ray._private.state.actors().values()
-        return (
-            len(
-                list(
-                    filter(
-                        lambda a: a["State"] == "ALIVE",
-                        _actors,
-                    )
-                )
-            )
-            == 3
+    wait_for_condition(
+        lambda: len(
+            list_actors(address=cluster.address, filters=[("STATE", "=", "ALIVE")])
         )
-
-    wait_for_condition(_check)
+        == 3,
+    )
 
     # Ensures ListApplications method on the head node is succeeding.
     wait_for_condition(

python/ray/serve/tests/test_proxy.py

@@ -9,17 +9,15 @@
     DEFAULT_UVICORN_KEEP_ALIVE_TIMEOUT_S,
     SERVE_NAMESPACE,
 )
+from ray.util.state import list_actors
 
 
 class TestTimeoutKeepAliveConfig:
     """Test setting keep_alive_timeout_s in config and env."""
 
     def get_proxy_actor(self) -> ActorHandle:
-        proxy_actor_name = None
-        for actor in ray._private.state.actors().values():
-            if actor["ActorClassName"] == "ProxyActor":
-                proxy_actor_name = actor["Name"]
-        return ray.get_actor(proxy_actor_name, namespace=SERVE_NAMESPACE)
+        [proxy_actor] = list_actors(filters=[("class_name", "=", "ProxyActor")])
+        return ray.get_actor(proxy_actor.name, namespace=SERVE_NAMESPACE)
 
     def test_default_keep_alive_timeout_s(self, ray_shutdown):
         """Test when no keep_alive_timeout_s is set.

python/ray/serve/tests/test_standalone_2.py

@@ -7,7 +7,6 @@
 import pytest
 
 import ray
-import ray._private.state
 import ray.actor
 from ray import serve
 from ray._common.test_utils import wait_for_condition

python/ray/serve/tests/test_standalone_3.py

@@ -8,7 +8,6 @@
 import pytest
 
 import ray
-import ray._private.state
 import ray.actor
 from ray import serve
 from ray._common.test_utils import SignalActor, wait_for_condition
@@ -20,6 +19,7 @@
 from ray.serve.context import _get_global_client
 from ray.serve.schema import ProxyStatus, ServeInstanceDetails
 from ray.tests.conftest import call_ray_stop_only  # noqa: F401
+from ray.util.state import list_actors
 
 
 @pytest.fixture
@@ -287,23 +287,15 @@ def __call__(self, *args):
     serve.run(A.bind(), name="app_f")
 
     # 2 proxies, 1 controller, 2 replicas.
-    wait_for_condition(lambda: len(ray._private.state.actors()) == 5)
+    wait_for_condition(lambda: len(list_actors()) == 5)
     assert len(ray.nodes()) == 2
 
     # Stop all deployment replicas.
     serve.delete("app_f")
 
     # The http proxy on worker node should exit as well.
     wait_for_condition(
-        lambda: len(
-            list(
-                filter(
-                    lambda a: a["State"] == "ALIVE",
-                    ray._private.state.actors().values(),
-                )
-            )
-        )
-        == 2
+        lambda: len(list_actors(filters=[("STATE", "=", "ALIVE")])) == 2,
     )
 
     client = _get_global_client()
@@ -356,7 +348,7 @@ def __call__(self):
     serve.run(HelloModel.options(num_replicas=2).bind())
 
     # 3 proxies, 1 controller, 2 replicas.
-    wait_for_condition(lambda: len(ray._private.state.actors()) == 6)
+    wait_for_condition(lambda: len(list_actors()) == 6)
     assert len(ray.nodes()) == 3
 
     client = _get_global_client()
@@ -435,7 +427,7 @@ def __call__(self):
     serve.run(target=model)
 
     # Ensure total actors of 2 proxies, 1 controller, and 2 replicas
-    wait_for_condition(lambda: len(ray._private.state.actors()) == 5)
+    wait_for_condition(lambda: len(list_actors()) == 5)
     assert len(ray.nodes()) == 2
 
     # Call `graceful_shutdown()` on the controller, so it will start shutdown.
@@ -450,9 +442,7 @@ def __call__(self):
 
     # Ensure the all resources are shutdown.
     wait_for_condition(
-        lambda: all(
-            [actor["State"] == "DEAD" for actor in ray._private.state.actors().values()]
-        )
+        lambda: len(list_actors(filters=[("STATE", "=", "ALIVE")])) == 0,
     )
 
     # Clean up serve.
@@ -496,7 +486,7 @@ def __call__(self):
     serve.run(target=model)
 
     # Ensure total actors of 2 proxies, 1 controller, and 2 replicas
-    wait_for_condition(lambda: len(ray._private.state.actors()) == 5)
+    wait_for_condition(lambda: len(list_actors()) == 5)
     assert len(ray.nodes()) == 2
 
     # Ensure client times out if the controller does not shutdown within timeout.
@@ -510,9 +500,7 @@ def __call__(self):
 
     # Ensure the all resources are shutdown gracefully.
     wait_for_condition(
-        lambda: all(
-            [actor["State"] == "DEAD" for actor in ray._private.state.actors().values()]
-        ),
+        lambda: len(list_actors(filters=[("STATE", "=", "ALIVE")])) == 0,
     )
 
     # Clean up serve.
@@ -543,9 +531,7 @@ def __call__(self):
 
     # Ensure the all resources are shutdown gracefully.
     wait_for_condition(
-        lambda: all(
-            [actor["State"] == "DEAD" for actor in ray._private.state.actors().values()]
-        ),
+        lambda: len(list_actors(filters=[("STATE", "=", "ALIVE")])) == 0,
     )
 
     all_serve_logs = ""