chore: import upstream snapshot with attribution

python/ray/tests/test_memstat.py

@@ -0,0 +1,379 @@
+import os
+import sys
+import time
+
+import numpy as np
+import pytest
+
+import ray
+from ray._common.test_utils import Semaphore, wait_for_condition
+from ray._private.internal_api import memory_summary
+from ray.cluster_utils import Cluster, cluster_not_supported
+
+# RayConfig to enable recording call sites during ObjectRej creations.
+ray_config = {"record_ref_creation_sites": True}
+
+# Unique strings.
+DRIVER_PID = "Driver"
+WORKER_PID = "Worker"
+UNKNOWN_SIZE = " ? "
+
+# Reference status.
+PINNED_IN_MEMORY = "PINNED_IN_MEMORY"
+LOCAL_REF = "LOCAL_REFERENCE"
+USED_BY_PENDING_TASK = "USED_BY_PENDING_TASK"
+CAPTURED_IN_OBJECT = "CAPTURED_IN_OBJECT"
+ACTOR_HANDLE = "ACTOR_HANDLE"
+
+# Call sites.
+PUT_OBJ = "(put object)"
+TASK_CALL_OBJ = "(task call)"
+ACTOR_TASK_CALL_OBJ = "(actor call)"
+DESER_TASK_ARG = "(deserialize task arg)"
+# Only 22 characters can be matched because longer strings are wrapped around.
+DESER_ACTOR_TASK_ARG = "(deserialize actor tas"
+
+# Group by and sort by parameters.
+NODE_ADDRESS = "node address"
+STACK_TRACE = "stack trace"
+PID = "pid"
+OBJECT_SIZE = "object size"
+REFERENCE_TYPE = "reference type"
+
+# Task status.
+WAITING_FOR_DEPENDENCIES = "PENDING_ARGS_AVAIL"
+SCHEDULED = "PENDING_NODE_ASSIGNMENT"
+FINISHED = "FINISHED"
+WAITING_FOR_EXECUTION = "SUBMITTED_TO_WORKER"
+
+
+def data_lines(memory_str):
+    for line in memory_str.split("\n"):
+        if (
+            PINNED_IN_MEMORY in line
+            or LOCAL_REF in line
+            or USED_BY_PENDING_TASK in line
+            or CAPTURED_IN_OBJECT in line
+            or ACTOR_HANDLE in line
+        ):
+            yield line
+        else:
+            continue
+
+
+def num_objects(memory_str):
+    n = 0
+    for line in data_lines(memory_str):
+        n += 1
+    return n
+
+
+def count(memory_str, substr):
+    substr = substr[:42]
+    n = 0
+    for line in memory_str.split("\n"):
+        if substr in line:
+            n += 1
+    return n
+
+
+@pytest.mark.parametrize(
+    "ray_start_regular", [{"_system_config": ray_config}], indirect=True
+)
+def test_driver_put_ref(ray_start_regular):
+    address = ray_start_regular["address"]
+    info = memory_summary(address)
+    assert num_objects(info) == 0, info
+    x_id = ray.put("HI")
+    info = memory_summary(address)
+    print(info)
+    assert num_objects(info) == 1, info
+    assert count(info, DRIVER_PID) == 1, info
+    assert count(info, WORKER_PID) == 0, info
+    del x_id
+    info = memory_summary(address)
+    assert num_objects(info) == 0, info
+
+
+@pytest.mark.parametrize(
+    "ray_start_regular", [{"_system_config": ray_config}], indirect=True
+)
+def test_worker_task_refs(ray_start_regular):
+    address = ray_start_regular["address"]
+
+    @ray.remote
+    def f(y):
+        from ray._private.internal_api import memory_summary
+
+        x_id = ray.put("HI")
+        info = memory_summary(address)
+        del x_id
+        return info
+
+    x_id = f.remote(np.zeros(100000))
+    info = ray.get(x_id)
+    print(info)
+    assert num_objects(info) == 4, info
+    # Task argument plus task return ids.
+    assert count(info, TASK_CALL_OBJ) == 2, info
+    assert count(info, DRIVER_PID) == 2, info
+    assert count(info, WORKER_PID) == 2, info
+    assert count(info, LOCAL_REF) == 2, info
+    assert count(info, PINNED_IN_MEMORY) == 1, info
+    assert count(info, PUT_OBJ) == 1, info
+    assert count(info, DESER_TASK_ARG) == 1, info
+    assert count(info, UNKNOWN_SIZE) == 1, info
+
+    print(ray_start_regular)
+    info = memory_summary(address)
+    print(info)
+    assert num_objects(info) == 1, info
+    assert count(info, DRIVER_PID) == 1, info
+    assert count(info, TASK_CALL_OBJ) == 1, info
+    assert count(info, UNKNOWN_SIZE) == 0, info
+    assert count(info, x_id.hex()) == 1, info
+
+    del x_id
+    info = memory_summary(address)
+    assert num_objects(info) == 0, info
+
+
+@pytest.mark.parametrize(
+    "ray_start_regular", [{"_system_config": ray_config}], indirect=True
+)
+def test_actor_task_refs(ray_start_regular):
+    address = ray_start_regular["address"]
+
+    @ray.remote
+    class Actor:
+        def __init__(self):
+            self.refs = []
+
+        def f(self, x):
+            from ray._private.internal_api import memory_summary
+
+            self.refs.append(x)
+            return memory_summary(address)
+
+    def make_actor():
+        return Actor.remote()
+
+    actor = make_actor()
+    x_id = actor.f.remote(np.zeros(100000))
+    info = ray.get(x_id)
+    print(info)
+    assert num_objects(info) == 4, info
+    # Actor handle, task argument id, task return id.
+    assert count(info, ACTOR_TASK_CALL_OBJ) == 3, info
+    assert count(info, DRIVER_PID) == 3, info
+    assert count(info, WORKER_PID) == 1, info
+    assert count(info, LOCAL_REF) == 1, info
+    assert count(info, PINNED_IN_MEMORY) == 1, info
+    assert count(info, USED_BY_PENDING_TASK) == 1, info
+    assert count(info, ACTOR_HANDLE) == 1, info
+    assert count(info, DESER_ACTOR_TASK_ARG) == 1, info
+    del x_id
+
+    # These should accumulate in the actor.
+    for _ in range(5):
+        ray.get(actor.f.remote([ray.put(np.zeros(100000))]))
+    info = memory_summary(address)
+    print(info)
+    assert count(info, DESER_ACTOR_TASK_ARG) == 5, info
+    assert count(info, ACTOR_TASK_CALL_OBJ) == 1, info
+
+    # Cleanup.
+    del actor
+    time.sleep(1)
+    info = memory_summary(address)
+    assert num_objects(info) == 0, info
+
+
+@pytest.mark.parametrize(
+    "ray_start_regular", [{"_system_config": ray_config}], indirect=True
+)
+def test_nested_object_refs(ray_start_regular):
+    address = ray_start_regular["address"]
+    x_id = ray.put(np.zeros(100000))
+    y_id = ray.put([x_id])
+    z_id = ray.put([y_id])
+    del x_id, y_id
+    info = memory_summary(address)
+    print(info)
+    assert num_objects(info) == 3, info
+    assert count(info, LOCAL_REF) == 1, info
+    assert count(info, CAPTURED_IN_OBJECT) == 2, info
+    del z_id
+
+
+@pytest.mark.parametrize(
+    "ray_start_regular", [{"_system_config": ray_config}], indirect=True
+)
+def test_pinned_object_call_site(ray_start_regular):
+    address = ray_start_regular["address"]
+    # Local ref only.
+    x_id = ray.put(np.zeros(100000))
+    info = memory_summary(address)
+    print(info)
+    assert num_objects(info) == 1, info
+    assert count(info, LOCAL_REF) == 1, info
+    assert count(info, PINNED_IN_MEMORY) == 0, info
+
+    # Local ref + pinned buffer.
+    buf = ray.get(x_id)
+    info = memory_summary(address)
+    print(info)
+    assert num_objects(info) == 1, info
+    assert count(info, LOCAL_REF) == 0, info
+    assert count(info, PINNED_IN_MEMORY) == 1, info
+
+    # Just pinned buffer.
+    del x_id
+    info = memory_summary(address)
+    print(info)
+    assert num_objects(info) == 1, info
+    assert count(info, LOCAL_REF) == 0, info
+    assert count(info, PINNED_IN_MEMORY) == 1, info
+
+    # Nothing.
+    del buf
+    info = memory_summary(address)
+    print(info)
+    assert num_objects(info) == 0, info
+
+
+@pytest.mark.xfail(cluster_not_supported, reason="cluster not supported")
+def test_multi_node_stats(shutdown_only):
+    # NOTE(mwtian): using env var only enables the feature on workers, while
+    # using head_node_args={"_system_config": ray_config} only enables the
+    # feature on the driver.
+    os.environ["RAY_record_ref_creation_sites"] = "1"
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=1)
+
+    ray.init(address=cluster.address)
+
+    @ray.remote(num_cpus=1)
+    class Actor:
+        def __init__(self):
+            self.ref = ray.put(np.zeros(100000))
+
+        def ping(self):
+            pass
+
+    # Each actor will be on a different node.
+    a = Actor.remote()
+    b = Actor.remote()
+    ray.get(a.ping.remote())
+    ray.get(b.ping.remote())
+
+    # Verify we have collected stats across the nodes.
+    info = memory_summary(cluster.address)
+    print(info)
+    assert count(info, PUT_OBJ) == 2, info
+
+
+@pytest.mark.parametrize(
+    "ray_start_regular", [{"_system_config": ray_config}], indirect=True
+)
+def test_group_by_sort_by(ray_start_regular):
+    address = ray_start_regular["address"]
+
+    @ray.remote
+    def f(y):
+        x_id = ray.put("HI")
+        info_a = memory_summary(
+            address, group_by="STACK_TRACE", sort_by="REFERENCE_TYPE"
+        )
+        info_b = memory_summary(address, group_by="NODE_ADDRESS", sort_by="OBJECT_SIZE")
+        info_c = memory_summary(address, group_by="NODE_ADDRESS", sort_by="PID")
+        del x_id
+        return info_a, info_b, info_c
+
+    x_id = f.remote(np.zeros(100000))
+    info_a, info_b, info_c = ray.get(x_id)
+    print(info_c)
+    assert count(info_a, STACK_TRACE) == 7, info_a
+    assert count(info_a, REFERENCE_TYPE) == 1, info_a
+    assert count(info_b, NODE_ADDRESS) == 3, info_b
+    assert count(info_b, OBJECT_SIZE) == 1, info_b
+    assert count(info_c, NODE_ADDRESS) == 3, info_c
+    assert count(info_c, PID) == 1, info_c
+
+
+@pytest.mark.parametrize(
+    "ray_start_regular", [{"_system_config": ray_config}], indirect=True
+)
+def test_memory_used_output(ray_start_regular):
+    address = ray_start_regular["address"]
+    _ = ray.put(np.ones(8 * 1024 * 1024, dtype=np.int8))
+
+    info = memory_summary(address)
+    print(info)
+    assert count(info, "Plasma memory usage 8 MiB") == 1, info
+    assert count(info, "8388862.0 B") == 2, info
+
+
+def test_task_status(ray_start_regular):
+    address = ray_start_regular["address"]
+
+    @ray.remote
+    def dep(sema, x=None):
+        ray.get(sema.acquire.remote())
+        return
+
+    @ray.remote(num_gpus=1)
+    def impossible():
+        pass
+
+    # Filter out actor handle refs.
+    def filtered_summary():
+        data = "\n".join(
+            [
+                line
+                for line in memory_summary(address, line_wrap=False).split("\n")
+                if "ACTOR_HANDLE" not in line
+            ]
+        )
+        print(data)
+        return data
+
+    sema = Semaphore.remote(value=0)
+    x = dep.remote(sema)
+    y = dep.remote(sema, x=x)
+    im = impossible.remote()  # noqa
+    # x and its semaphore task are scheduled. im cannot
+    # be scheduled, so it is pending forever.
+    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_EXECUTION) == 2)
+    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_DEPENDENCIES) == 1)
+    wait_for_condition(lambda: count(filtered_summary(), SCHEDULED) == 1)
+
+    z = dep.remote(sema, x=x)
+    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_DEPENDENCIES) == 2)
+    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_EXECUTION) == 2)
+    wait_for_condition(lambda: count(filtered_summary(), FINISHED) == 0)
+
+    sema.release.remote()
+    time.sleep(2)
+    wait_for_condition(lambda: count(filtered_summary(), FINISHED) == 1)
+    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_DEPENDENCIES) == 0)
+    # y, z, and two semaphore tasks are scheduled.
+    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_EXECUTION) == 4)
+
+    sema.release.remote()
+    wait_for_condition(lambda: count(filtered_summary(), FINISHED) == 2)
+    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_DEPENDENCIES) == 0)
+    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_EXECUTION) == 2)
+
+    sema.release.remote()
+    ray.get(y)
+    ray.get(z)
+    wait_for_condition(lambda: count(filtered_summary(), FINISHED) == 3)
+    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_DEPENDENCIES) == 0)
+    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_EXECUTION) == 0)
+    wait_for_condition(lambda: count(filtered_summary(), SCHEDULED) == 1)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))