chore: import upstream snapshot with attribution

python/ray/tests/test_metrics.py

@@ -0,0 +1,346 @@
+import os
+import platform
+import sys
+
+import pytest
+import requests
+
+import ray
+from ray._common.network_utils import build_address
+from ray._common.test_utils import wait_for_condition
+from ray._private.test_utils import (
+    get_node_stats,
+    wait_until_succeeded_without_exception,
+)
+from ray.core.generated import common_pb2
+
+import psutil
+
+_WIN32 = os.name == "nt"
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Hangs on Windows.")
+def test_worker_stats(shutdown_only):
+    ray.init(num_cpus=2, include_dashboard=True)
+    raylet = ray.nodes()[0]
+    reply = get_node_stats(raylet)
+    # Check that there is one connected driver.
+    drivers = [
+        worker
+        for worker in reply.core_workers_stats
+        if worker.worker_type == common_pb2.DRIVER
+    ]
+    assert len(drivers) == 1
+    assert os.getpid() == drivers[0].pid
+
+    @ray.remote
+    def f():
+        return os.getpid()
+
+    @ray.remote(num_cpus=1)
+    class Actor:
+        def __init__(self):
+            pass
+
+        def f(self):
+            return os.getpid()
+
+    # Run a remote function and actor to create workers.
+    ray.get(f.remote())
+    a = Actor.remote()
+    ray.get(a.f.remote())
+
+    # 1 actor + 1 worker for task + 1 driver
+    num_workers = 3
+
+    def verify():
+        reply = get_node_stats(raylet)
+        # Check that the rest of the processes are workers, 1 for each CPU.
+
+        assert len(reply.core_workers_stats) == num_workers
+        # Check that all processes are Python.
+        pids = [worker.pid for worker in reply.core_workers_stats]
+        processes = [
+            p.info["name"]
+            for p in psutil.process_iter(attrs=["pid", "name"])
+            if p.info["pid"] in pids
+        ]
+        for process in processes:
+            # TODO(ekl) why does travis/mi end up in the process list
+            assert (
+                "python" in process
+                or "mini" in process
+                or "conda" in process
+                or "travis" in process
+                or "runner" in process
+                or "pytest" in process
+                or "ray" in process
+            ), process
+
+        return True
+
+    wait_for_condition(verify)
+
+
+def get_owner_info(node_ids):
+    node_addrs = {
+        n["NodeID"]: (n["NodeManagerAddress"], n["NodeManagerPort"])
+        for n in ray.nodes()
+    }
+    # Force a global gc to clean up the object store.
+    ray._private.internal_api.global_gc()
+    owner_stats = {n: 0 for n in node_ids}
+    primary_copy_stats = {n: 0 for n in node_ids}
+
+    for node_id in node_ids:
+        node_stats = ray._private.internal_api.node_stats(
+            node_addrs[node_id][0], node_addrs[node_id][1], False
+        )
+        num_owned_objects = sum(
+            [stats.num_owned_objects for stats in node_stats.core_workers_stats]
+        )
+        num_owned_actors = sum(
+            [stats.num_owned_actors for stats in node_stats.core_workers_stats]
+        )
+        owner_stats[node_id] = (num_owned_objects, num_owned_actors)
+        primary_copy_stats[
+            node_id
+        ] = node_stats.store_stats.num_object_store_primary_copies
+
+    print(owner_stats)
+    print(node_ids)
+    owner_stats = [owner_stats.get(node_id, 0) for node_id in node_ids]
+    primary_copy_stats = [primary_copy_stats.get(node_id, 0) for node_id in node_ids]
+    print("owner_stats", owner_stats)
+    print("primary_copy_stats", primary_copy_stats)
+
+    return owner_stats, primary_copy_stats
+
+
+def test_node_object_metrics(ray_start_cluster):
+    NUM_NODES = 3
+    cluster = ray_start_cluster
+    for i in range(NUM_NODES):
+        cluster.add_node(True, resources={f"node_{i}": 1})
+        if i == 0:
+            ray.init(address=cluster.address)
+    node_ids = []
+
+    for i in range(NUM_NODES):
+
+        @ray.remote(resources={f"node_{i}": 1})
+        def get_node_id():
+            return ray.get_runtime_context().get_node_id()
+
+        node_ids.append(ray.get(get_node_id.remote()))
+
+    # Object store stats
+    # x is owned by node_0
+    # x is stored at node_0
+    x = ray.put([1])  # noqa: F841
+    wait_for_condition(
+        lambda: get_owner_info(node_ids) == ([(1, 0), (0, 0), (0, 0)], [1, 0, 0])
+    )
+
+    # Test nested with put
+    @ray.remote(resources={"node_1": 1})
+    def big_obj():
+        # b is owned by node_1
+        # b is stored at node_1
+        b = ray.put([1] * 1024 * 1024 * 10)
+        return b
+
+    # Object store stats
+    # big_obj is owned by node_0
+    # big_obj is stored in memory (no primary copy)
+    big_obj_ref = big_obj.remote()  # noqa: F841
+    wait_for_condition(
+        lambda: get_owner_info(node_ids) == ([(2, 0), (1, 0), (0, 0)], [1, 1, 0])
+    )
+
+    # Test nested with task (small output)
+    @ray.remote(resources={"node_1": 1})
+    def nest_task(s):
+        @ray.remote(resources={"node_2": 1})
+        def task():
+            return [1] * s
+
+        # t is owned by node_1
+        # if s is small,
+        # then it's is stored in memory of node_1 (no primary copy)
+        # else it's stored in object store of node_1
+        t = task.remote()
+        return t
+
+    # nest_ref is owned by node_0
+    # nest_ref is stored in memory (no primary copy)
+    nest_ref = nest_task.remote(1)  # noqa: F841
+    wait_for_condition(
+        lambda: get_owner_info(node_ids) == ([(3, 0), (2, 0), (0, 0)], [1, 1, 0])
+    )
+
+    big_nest = nest_task.remote(1024 * 1024 * 10)  # noqa: F841
+
+    wait_for_condition(
+        lambda: get_owner_info(node_ids) == ([(4, 0), (3, 0), (0, 0)], [1, 1, 1])
+    )
+
+    @ray.remote(resources={"node_2": 0.5}, num_cpus=0)
+    class A:
+        def ready(self):
+            return
+
+        def gen(self):
+            return ray.put(10)
+
+    # actor is owned by node_0
+    # actor is not an object, so no object store copies
+    actor = A.remote()  # noqa: F841
+    ray.get(actor.ready.remote())
+    wait_for_condition(
+        lambda: get_owner_info(node_ids) == ([(4, 1), (3, 0), (0, 0)], [1, 1, 1])
+    )
+
+    # Test with detached owned
+    # detached actor is owned by GCS. So it's not counted in the owner stats
+    detached_actor = A.options(lifetime="detached", name="A").remote()
+    ray.get(detached_actor.ready.remote())
+    for i in range(3):
+        assert get_owner_info(node_ids) == ([(4, 1), (3, 0), (0, 0)], [1, 1, 1])
+        import time
+
+        time.sleep(1)
+    # gen_obj is owned by node_0
+    # the inner object is owned by A (node_2)
+    # the inner object is stored in object store of node_2
+    gen_obj = detached_actor.gen.remote()  # noqa: F841
+    wait_for_condition(
+        lambda: get_owner_info(node_ids) == ([(5, 1), (3, 0), (1, 0)], [1, 1, 2])
+    )
+
+
+def test_running_tasks(ray_start_cluster):
+    NUM_NODES = 3
+    cluster = ray_start_cluster
+    for i in range(NUM_NODES):
+        cluster.add_node(True, resources={f"node_{i}": 1})
+        if i == 0:
+            ray.init(address=cluster.address)
+    node_ids = []
+
+    for i in range(NUM_NODES):
+
+        @ray.remote(resources={f"node_{i}": 1})
+        def get_node_id():
+            return ray.get_runtime_context().get_node_id()
+
+        node_ids.append(ray.get(get_node_id.remote()))
+
+    @ray.remote
+    def f(t):
+        import time
+
+        time.sleep(t)
+
+    tasks = [
+        f.options(resources={"node_0": 1}).remote(0),
+        f.options(resources={"node_1": 1}).remote(100000),
+        f.options(resources={"node_2": 1}).remote(100000),
+    ]
+
+    ready, pending = ray.wait(tasks)
+    assert len(ready) == 1
+    assert len(pending) == 2
+
+    node_addrs = {
+        n["NodeID"]: (n["NodeManagerAddress"], n["NodeManagerPort"])
+        for n in ray.nodes()
+    }
+
+    def check():
+        for i in range(NUM_NODES):
+            node_stats = ray._private.internal_api.node_stats(
+                node_addrs[node_ids[i]][0], node_addrs[node_ids[i]][1], False
+            )
+
+            if i == 0:
+                assert (
+                    sum(
+                        [
+                            stats.num_running_tasks
+                            for stats in node_stats.core_workers_stats
+                        ]
+                    )
+                    == 0
+                )
+            else:
+                assert (
+                    sum(
+                        [
+                            stats.num_running_tasks
+                            for stats in node_stats.core_workers_stats
+                        ]
+                    )
+                    == 1
+                )
+        return True
+
+    wait_for_condition(check)
+
+
+def test_multi_node_metrics_export_port_discovery(ray_start_cluster):
+    NUM_NODES = 3
+    cluster = ray_start_cluster
+    nodes = [cluster.add_node() for _ in range(NUM_NODES)]
+    nodes = {
+        node.address_info["metrics_export_port"]: node.address_info for node in nodes
+    }
+    cluster.wait_for_nodes()
+    ray.init(address=cluster.address)
+    node_info_list = ray.nodes()
+
+    for node_info in node_info_list:
+        metrics_export_port = node_info["MetricsExportPort"]
+        address_info = nodes[metrics_export_port]
+        assert address_info["raylet_socket_name"] == node_info["RayletSocketName"]
+
+        # Make sure we can ping Prometheus endpoints.
+        def test_prometheus_endpoint():
+            response = requests.get(
+                f"http://{build_address('localhost', metrics_export_port)}",
+                # Fail the request early on if connection timeout
+                timeout=1.0,
+            )
+            return response.status_code == 200
+
+        assert wait_until_succeeded_without_exception(
+            test_prometheus_endpoint,
+            (requests.exceptions.ConnectionError,),
+            # The dashboard takes more than 2s to startup.
+            timeout_ms=10 * 1000,
+        )
+
+
+def test_opentelemetry_conflict(shutdown_only):
+    ray.init()
+
+    # After ray.init(), opencensus protobuf should not be registered.
+    # Otherwise, it might conflict with other versions generated opencensus protobuf.
+
+    from google.protobuf.descriptor_pool import Default as DefaultPool
+
+    pool = DefaultPool()
+
+    try:
+        found_file = pool.FindFileByName("opencensus/proto/resource/v1/resource.proto")
+    except KeyError:
+        found_file = None
+
+    assert found_file is None, "opencensus protobuf registered after ray.init()"
+
+    # Make sure the similar resource protobuf also doesn't raise an exception.
+    from opentelemetry.proto.resource.v1 import resource_pb2  # noqa
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))