chore: import upstream snapshot with attribution

python/ray/tests/test_actor_state_metrics.py

@@ -0,0 +1,326 @@
+import asyncio
+import sys
+import time
+from collections import defaultdict
+from typing import Dict
+
+import pytest
+
+import ray
+from ray._common.test_utils import (
+    PrometheusTimeseries,
+    run_string_as_driver,
+    wait_for_condition,
+)
+from ray._private.test_utils import raw_metric_timeseries
+from ray._private.worker import RayContext
+from ray.util.state import list_actors
+
+_SYSTEM_CONFIG = {
+    "metrics_report_interval_ms": 200,
+}
+
+
+def actors_by_state(info: RayContext, timeseries: PrometheusTimeseries) -> Dict:
+    res = raw_metric_timeseries(info, timeseries)
+    actors_info = defaultdict(int)
+    if "ray_actors" in res:
+        for sample in res["ray_actors"]:
+            actors_info[sample.labels["State"]] += sample.value
+    for k, v in actors_info.copy().items():
+        if v == 0:
+            del actors_info[k]
+    print(f"Actors by state: {actors_info}")
+    return actors_info
+
+
+def actors_by_name(info: RayContext, timeseries: PrometheusTimeseries) -> Dict:
+    res = raw_metric_timeseries(info, timeseries)
+    actors_info = defaultdict(int)
+    if "ray_actors" in res:
+        for sample in res["ray_actors"]:
+            actors_info[sample.labels["Name"]] += sample.value
+    for k, v in actors_info.copy().items():
+        if v == 0:
+            del actors_info[k]
+    print(f"Actors by name: {actors_info}")
+    return actors_info
+
+
+def test_basic_states(shutdown_only):
+    info = ray.init(num_cpus=3, _system_config=_SYSTEM_CONFIG)
+
+    @ray.remote(num_cpus=1)
+    class Actor:
+        def ping(self):
+            pass
+
+        def sleep(self):
+            time.sleep(999)
+
+        def get(self):
+            @ray.remote
+            def sleep():
+                time.sleep(999)
+
+            ray.get(sleep.remote())
+
+        def wait(self):
+            @ray.remote
+            def sleep():
+                time.sleep(999)
+
+            ray.wait([sleep.remote()])
+
+    a = Actor.remote()
+    b = Actor.remote()
+    c = Actor.remote()
+    ray.get(a.ping.remote())
+    ray.get(b.ping.remote())
+    ray.get(c.ping.remote())
+    d = Actor.remote()
+    timeseries = PrometheusTimeseries()
+
+    # Test creation states.
+    expected = {
+        "ALIVE": 3,
+        "ALIVE_IDLE": 3,
+        "PENDING_CREATION": 1,
+    }
+    wait_for_condition(
+        lambda: actors_by_state(info, timeseries) == expected,
+        timeout=20,
+        retry_interval_ms=500,
+    )
+
+    # Test running states.
+    a.sleep.remote()
+    b.get.remote()
+    c.wait.remote()
+    expected = {
+        "ALIVE": 3,
+        "ALIVE_RUNNING_TASKS": 3,
+        "PENDING_CREATION": 1,
+    }
+    wait_for_condition(
+        lambda: actors_by_state(info, timeseries) == expected,
+        timeout=20,
+        retry_interval_ms=500,
+    )
+    del d
+
+
+def test_destroy_actors(shutdown_only):
+    info = ray.init(num_cpus=3, _system_config=_SYSTEM_CONFIG)
+
+    @ray.remote(num_cpus=1)
+    class Actor:
+        def ping(self):
+            pass
+
+    a = Actor.remote()
+    b = Actor.remote()
+    c = Actor.remote()
+    del a
+    del b
+    timeseries = PrometheusTimeseries()
+    expected = {
+        "ALIVE": 1,
+        "ALIVE_IDLE": 1,
+        "DEAD": 2,
+    }
+    wait_for_condition(
+        lambda: actors_by_state(info, timeseries) == expected,
+        timeout=20,
+        retry_interval_ms=500,
+    )
+    del c
+
+
+def test_destroy_actors_from_driver(monkeypatch, shutdown_only):
+    with monkeypatch.context() as m:
+        # Dead actors are not cached.
+        m.setenv("RAY_maximum_gcs_destroyed_actor_cached_count", 5)
+        m.setenv("RAY_WORKER_TIMEOUT_S", 5)
+        driver = """
+import ray
+ray.init("auto")
+@ray.remote(num_cpus=0)
+class Actor:
+    def ready(self):
+        pass
+actors = [Actor.remote() for _ in range(10)]
+ray.get([actor.ready.remote() for actor in actors])
+"""
+        info = ray.init(num_cpus=3, _system_config=_SYSTEM_CONFIG)
+
+        output = run_string_as_driver(driver)
+        print(output)
+        timeseries = PrometheusTimeseries()
+
+        expected = {
+            "DEAD": 10,
+        }
+        wait_for_condition(
+            lambda: expected.items() <= actors_by_state(info, timeseries).items(),
+            timeout=20,
+            retry_interval_ms=500,
+        )
+
+        """
+        Make sure even after the actor entries are deleted from GCS by GC
+        the metrics are correct.
+        """
+        # Wait until the state API returns the # of actors are 0
+        # becasue entries are GC'ed by GCS.
+        wait_for_condition(lambda: len(list_actors()) == 5)
+        # DEAD count shouldn't be changed.
+        wait_for_condition(
+            lambda: expected.items() <= actors_by_state(info, timeseries).items(),
+            timeout=20,
+            retry_interval_ms=500,
+        )
+
+
+def test_dep_wait(shutdown_only):
+    info = ray.init(num_cpus=3, _system_config=_SYSTEM_CONFIG)
+
+    @ray.remote
+    def sleep():
+        time.sleep(999)
+
+    @ray.remote(num_cpus=1)
+    class Actor:
+        def __init__(self, x):
+            pass
+
+    a = Actor.remote(sleep.remote())
+    timeseries = PrometheusTimeseries()
+    expected = {
+        "DEPENDENCIES_UNREADY": 1,
+    }
+    wait_for_condition(
+        lambda: actors_by_state(info, timeseries) == expected,
+        timeout=20,
+        retry_interval_ms=500,
+    )
+    del a
+
+
+def test_async_actor(shutdown_only):
+    info = ray.init(num_cpus=3, _system_config=_SYSTEM_CONFIG)
+
+    @ray.remote
+    def sleep():
+        time.sleep(999)
+
+    @ray.remote(max_concurrency=30)
+    class AsyncActor:
+        async def sleep(self):
+            await asyncio.sleep(300)
+
+        async def do_get(self):
+            await ray.get(sleep.remote())
+
+    a = AsyncActor.remote()
+    a.sleep.remote()
+    timeseries = PrometheusTimeseries()
+    expected = {
+        "ALIVE": 1,
+        "ALIVE_RUNNING_TASKS": 1,
+    }
+    wait_for_condition(
+        lambda: actors_by_state(info, timeseries) == expected,
+        timeout=20,
+        retry_interval_ms=500,
+    )
+
+    # Test that this transitions the entire actor to reporting IN_RAY_GET state.
+    a.do_get.remote()
+    a.do_get.remote()
+    expected = {
+        "ALIVE": 1,
+        "ALIVE_RUNNING_TASKS": 1,
+    }
+    wait_for_condition(
+        lambda: actors_by_state(info, timeseries) == expected,
+        timeout=20,
+        retry_interval_ms=500,
+    )
+    del a
+
+
+@ray.remote(num_cpus=1)
+class Actor1:
+    def sleep(self):
+        time.sleep(999)
+
+
+@ray.remote(num_cpus=1)
+class Actor2:
+    def sleep(self):
+        time.sleep(999)
+
+
+def test_tracking_by_name(shutdown_only):
+    info = ray.init(num_cpus=3, _system_config=_SYSTEM_CONFIG)
+
+    a = Actor1.remote()
+    b = Actor2.remote()
+    timeseries = PrometheusTimeseries()
+    expected = {
+        # one reported by gcs as ALIVE
+        # another reported by core worker as IDLE
+        "Actor1": 2,
+        "Actor2": 2,
+    }
+    wait_for_condition(
+        lambda: actors_by_name(info, timeseries) == expected,
+        timeout=20,
+        retry_interval_ms=500,
+    )
+
+    del a
+    del b
+
+
+def test_get_all_actors_info(shutdown_only):
+    ray.init(num_cpus=2, include_dashboard=True)
+
+    @ray.remote(num_cpus=1)
+    class Actor:
+        def ping(self):
+            pass
+
+    actor_1 = Actor.remote()
+    actor_2 = Actor.remote()
+    ray.get([actor_1.ping.remote(), actor_2.ping.remote()], timeout=5)
+    actors_info = list_actors(detail=True)
+    assert len(actors_info) == 2
+
+    job_id_hex = ray.get_runtime_context().get_job_id()
+    actors_info = list_actors(filters=[("job_id", "=", job_id_hex)], detail=True)
+    assert len(actors_info) == 2
+    actors_info = list_actors(
+        filters=[("job_id", "=", ray.JobID.from_int(100).hex())], detail=True
+    )
+    assert len(actors_info) == 0
+
+    # To filter actors by state
+    actor_3 = Actor.remote()
+    wait_for_condition(
+        lambda: len(list_actors(filters=[("state", "=", "PENDING_CREATION")])) == 1
+    )
+    assert actor_3._actor_id.hex() in list(
+        map(
+            lambda s: s.actor_id,
+            list_actors(filters=[("state", "=", "PENDING_CREATION")]),
+        )
+    )
+
+    with pytest.raises(ray.util.state.exception.RayStateApiException):
+        actors_info = list_actors(filters=[("state", "=", "UNKONWN_STATE")])
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))