[Dashboard] Ray memory dashboard backend (#8461)

python/ray/tests/test_metrics.py

@@ -10,9 +10,11 @@ from ray.core.generated import node_manager_pb2
 from ray.core.generated import node_manager_pb2_grpc
 from ray.core.generated import reporter_pb2
 from ray.core.generated import reporter_pb2_grpc
+from ray.dashboard.memory import (ReferenceType, decode_object_id_if_needed,
+                                  MemoryTableEntry, MemoryTable, SortingType)
 from ray.test_utils import (RayTestTimeoutException,
                             wait_until_succeeded_without_exception,
-                            wait_until_server_available)
+                            wait_until_server_available, wait_for_condition)
 
 import psutil  # We must import psutil after ray because we bundle it with ray.
 
@@ -377,6 +379,446 @@ def test_profiling_info_endpoint(shutdown_only):
     assert profiling_stats is not None
 
 
+# This variable is used inside test_memory_dashboard.
+# It is defined as a global variable to be used across all nested test
+# functions. We use it because memory table is updated every one second,
+# and we need to have a way to verify if the test is running with a fresh
+# new memory table.
+prev_memory_table = MemoryTable([]).__dict__()["group"]
+
+
+def test_memory_dashboard(shutdown_only):
+    """Test Memory table.
+
+    These tests verify examples in this document.
+    https://docs.ray.io/en/latest/memory-management.html#debugging-using-ray-memory
+    """
+    addresses = ray.init(num_cpus=2)
+    webui_url = addresses["webui_url"].replace("localhost", "http://127.0.0.1")
+    assert (wait_until_server_available(addresses["webui_url"]) is True)
+
+    def get_memory_table():
+        memory_table = requests.get(webui_url + "/api/memory_table").json()
+        return memory_table["result"]
+
+    def memory_table_ready():
+        """Wait until the new fresh memory table is ready."""
+        global prev_memory_table
+        memory_table = get_memory_table()
+        from pprint import pprint
+        print("Current")
+        pprint(memory_table)
+        print("Prev")
+        pprint(prev_memory_table)
+        is_ready = memory_table["group"] != prev_memory_table
+        prev_memory_table = memory_table["group"]
+        return is_ready
+
+    def stop_memory_table():
+        requests.get(webui_url + "/api/stop_memory_table").json()
+
+    def test_local_reference():
+        @ray.remote
+        def f(arg):
+            return arg
+
+        # a and b are local references.
+        a = ray.put(None)  # Noqa F841
+        b = f.remote(None)  # Noqa F841
+
+        wait_for_condition(memory_table_ready)
+        memory_table = get_memory_table()
+        summary = memory_table["summary"]
+        group = memory_table["group"]
+        assert summary["total_captured_in_objects"] == 0
+        assert summary["total_pinned_in_memory"] == 0
+        assert summary["total_used_by_pending_task"] == 0
+        assert summary["total_local_ref_count"] == 2
+        for table in group.values():
+            for entry in table["entries"]:
+                assert (
+                    entry["reference_type"] == ReferenceType.LOCAL_REFERENCE)
+        stop_memory_table()
+        return True
+
+    def test_object_pineed_in_memory():
+        import numpy as np
+
+        a = ray.put(np.zeros(1))
+        b = ray.get(a)  # Noqa F841
+        del a
+
+        wait_for_condition(memory_table_ready)
+        memory_table = get_memory_table()
+        summary = memory_table["summary"]
+        group = memory_table["group"]
+        assert summary["total_captured_in_objects"] == 0
+        assert summary["total_pinned_in_memory"] == 1
+        assert summary["total_used_by_pending_task"] == 0
+        assert summary["total_local_ref_count"] == 0
+        for table in group.values():
+            for entry in table["entries"]:
+                assert (
+                    entry["reference_type"] == ReferenceType.PINNED_IN_MEMORY)
+        stop_memory_table()
+        return True
+
+    def test_pending_task_references():
+        @ray.remote
+        def f(arg):
+            time.sleep(1)
+
+        a = ray.put(None)  # Noqa F841
+        b = f.remote(a)  # Noqa F841
+
+        wait_for_condition(memory_table_ready)
+        memory_table = get_memory_table()
+        summary = memory_table["summary"]
+        assert summary["total_captured_in_objects"] == 0
+        assert summary["total_pinned_in_memory"] == 1
+        assert summary["total_used_by_pending_task"] == 1
+        assert summary["total_local_ref_count"] == 1
+        # Make sure the function f is done before going to the next test.
+        # Otherwise, the memory table will be corrupted because the
+        # task f won't be done when the next test is running.
+        ray.get(b)
+        stop_memory_table()
+        return True
+
+    def test_serialized_object_id_reference():
+        @ray.remote
+        def f(arg):
+            time.sleep(1)
+
+        a = ray.put(None)  # Noqa F841
+        b = f.remote([a])  # Noqa F841
+
+        wait_for_condition(memory_table_ready)
+        memory_table = get_memory_table()
+        summary = memory_table["summary"]
+        assert summary["total_captured_in_objects"] == 0
+        assert summary["total_pinned_in_memory"] == 0
+        assert summary["total_used_by_pending_task"] == 1
+        assert summary["total_local_ref_count"] == 2
+        # Make sure the function f is done before going to the next test.
+        # Otherwise, the memory table will be corrupted because the
+        # task f won't be done when the next test is running.
+        ray.get(b)
+        stop_memory_table()
+        return True
+
+    def test_captured_object_id_reference():
+        a = ray.put(None)
+        b = ray.put([a])  # Noqa F841
+        del a
+
+        wait_for_condition(memory_table_ready)
+        memory_table = get_memory_table()
+        summary = memory_table["summary"]
+        assert summary["total_captured_in_objects"] == 1
+        assert summary["total_pinned_in_memory"] == 0
+        assert summary["total_used_by_pending_task"] == 0
+        assert summary["total_local_ref_count"] == 1
+        stop_memory_table()
+        return True
+
+    def test_actor_handle_reference():
+        @ray.remote
+        class Actor:
+            pass
+
+        a = Actor.remote()  # Noqa F841
+        b = Actor.remote()  # Noqa F841
+        c = Actor.remote()  # Noqa F841
+
+        wait_for_condition(memory_table_ready)
+        memory_table = get_memory_table()
+        summary = memory_table["summary"]
+        group = memory_table["group"]
+        assert summary["total_captured_in_objects"] == 0
+        assert summary["total_pinned_in_memory"] == 0
+        assert summary["total_used_by_pending_task"] == 0
+        assert summary["total_local_ref_count"] == 0
+        assert summary["total_actor_handles"] == 3
+        for table in group.values():
+            for entry in table["entries"]:
+                assert (entry["reference_type"] == ReferenceType.ACTOR_HANDLE)
+        stop_memory_table()
+        return True
+
+    # These tests should be retried because it takes at least one second
+    # to get the fresh new memory table. It is because memory table is updated
+    # Whenever raylet and node info is renewed which takes 1 second.
+    assert (wait_for_condition(
+        test_local_reference, timeout=30000, retry_interval_ms=1000) is True)
+
+    assert (wait_for_condition(
+        test_object_pineed_in_memory, timeout=30000, retry_interval_ms=1000) is
+            True)
+
+    assert (wait_for_condition(
+        test_pending_task_references, timeout=30000, retry_interval_ms=1000) is
+            True)
+
+    assert (wait_for_condition(
+        test_serialized_object_id_reference,
+        timeout=30000,
+        retry_interval_ms=1000) is True)
+
+    assert (wait_for_condition(
+        test_captured_object_id_reference,
+        timeout=30000,
+        retry_interval_ms=1000) is True)
+
+    assert (wait_for_condition(
+        test_actor_handle_reference, timeout=30000, retry_interval_ms=1000) is
+            True)
+
+
+"""Memory Table Unit Test"""
+
+NODE_ADDRESS = "127.0.0.1"
+IS_DRIVER = True
+PID = 1
+OBJECT_ID = "7wpsIhgZiBz/////AQAAyAEAAAA="
+ACTOR_ID = "fffffffffffffffff66d17ba010000c801000000"
+DECODED_ID = decode_object_id_if_needed(OBJECT_ID)
+OBJECT_SIZE = 100
+
+
+def build_memory_entry(*,
+                       local_ref_count,
+                       pinned_in_memory,
+                       submitted_task_reference_count,
+                       contained_in_owned,
+                       object_size,
+                       pid,
+                       object_id=OBJECT_ID,
+                       node_address=NODE_ADDRESS):
+    object_ref = {
+        "objectId": object_id,
+        "callSite": "(task call) /Users:458",
+        "objectSize": object_size,
+        "localRefCount": local_ref_count,
+        "pinnedInMemory": pinned_in_memory,
+        "submittedTaskRefCount": submitted_task_reference_count,
+        "containedInOwned": contained_in_owned
+    }
+    return MemoryTableEntry(
+        object_ref=object_ref,
+        node_address=node_address,
+        is_driver=IS_DRIVER,
+        pid=pid)
+
+
+def build_local_reference_entry(object_size=OBJECT_SIZE,
+                                pid=PID,
+                                node_address=NODE_ADDRESS):
+    return build_memory_entry(
+        local_ref_count=1,
+        pinned_in_memory=False,
+        submitted_task_reference_count=0,
+        contained_in_owned=[],
+        object_size=object_size,
+        pid=pid,
+        node_address=node_address)
+
+
+def build_used_by_pending_task_entry(object_size=OBJECT_SIZE,
+                                     pid=PID,
+                                     node_address=NODE_ADDRESS):
+    return build_memory_entry(
+        local_ref_count=0,
+        pinned_in_memory=False,
+        submitted_task_reference_count=2,
+        contained_in_owned=[],
+        object_size=object_size,
+        pid=pid,
+        node_address=node_address)
+
+
+def build_captured_in_object_entry(object_size=OBJECT_SIZE,
+                                   pid=PID,
+                                   node_address=NODE_ADDRESS):
+    return build_memory_entry(
+        local_ref_count=0,
+        pinned_in_memory=False,
+        submitted_task_reference_count=0,
+        contained_in_owned=[OBJECT_ID],
+        object_size=object_size,
+        pid=pid,
+        node_address=node_address)
+
+
+def build_actor_handle_entry(object_size=OBJECT_SIZE,
+                             pid=PID,
+                             node_address=NODE_ADDRESS):
+    return build_memory_entry(
+        local_ref_count=1,
+        pinned_in_memory=False,
+        submitted_task_reference_count=0,
+        contained_in_owned=[],
+        object_size=object_size,
+        pid=pid,
+        node_address=node_address,
+        object_id=ACTOR_ID)
+
+
+def build_pinned_in_memory_entry(object_size=OBJECT_SIZE,
+                                 pid=PID,
+                                 node_address=NODE_ADDRESS):
+    return build_memory_entry(
+        local_ref_count=0,
+        pinned_in_memory=True,
+        submitted_task_reference_count=0,
+        contained_in_owned=[],
+        object_size=object_size,
+        pid=pid,
+        node_address=node_address)
+
+
+def build_entry(object_size=OBJECT_SIZE,
+                pid=PID,
+                node_address=NODE_ADDRESS,
+                reference_type=ReferenceType.PINNED_IN_MEMORY):
+    if reference_type == ReferenceType.USED_BY_PENDING_TASK:
+        return build_used_by_pending_task_entry(
+            pid=pid, object_size=object_size, node_address=node_address)
+    elif reference_type == ReferenceType.LOCAL_REFERENCE:
+        return build_local_reference_entry(
+            pid=pid, object_size=object_size, node_address=node_address)
+    elif reference_type == ReferenceType.PINNED_IN_MEMORY:
+        return build_pinned_in_memory_entry(
+            pid=pid, object_size=object_size, node_address=node_address)
+    elif reference_type == ReferenceType.ACTOR_HANDLE:
+        return build_actor_handle_entry(
+            pid=pid, object_size=object_size, node_address=node_address)
+    elif reference_type == ReferenceType.CAPTURED_IN_OBJECT:
+        return build_captured_in_object_entry(
+            pid=pid, object_size=object_size, node_address=node_address)
+
+
+def test_invalid_memory_entry():
+    memory_entry = build_memory_entry(
+        local_ref_count=0,
+        pinned_in_memory=False,
+        submitted_task_reference_count=0,
+        contained_in_owned=[],
+        object_size=OBJECT_SIZE,
+        pid=PID)
+    assert memory_entry.is_valid() is False
+    memory_entry = build_memory_entry(
+        local_ref_count=0,
+        pinned_in_memory=False,
+        submitted_task_reference_count=0,
+        contained_in_owned=[],
+        object_size=-1,
+        pid=PID)
+    assert memory_entry.is_valid() is False
+
+
+def test_valid_reference_memory_entry():
+    memory_entry = build_local_reference_entry()
+    assert memory_entry.reference_type == ReferenceType.LOCAL_REFERENCE
+    assert memory_entry.object_id == ray.ObjectID(
+        decode_object_id_if_needed(OBJECT_ID))
+    assert memory_entry.is_valid() is True
+
+
+def test_reference_type():
+    # pinned in memory
+    memory_entry = build_pinned_in_memory_entry()
+    assert memory_entry.reference_type == ReferenceType.PINNED_IN_MEMORY
+
+    # used by pending task
+    memory_entry = build_used_by_pending_task_entry()
+    assert memory_entry.reference_type == ReferenceType.USED_BY_PENDING_TASK
+
+    # captued in object
+    memory_entry = build_captured_in_object_entry()
+    assert memory_entry.reference_type == ReferenceType.CAPTURED_IN_OBJECT
+
+    # actor handle
+    memory_entry = build_actor_handle_entry()
+    assert memory_entry.reference_type == ReferenceType.ACTOR_HANDLE
+
+
+def test_memory_table_summary():
+    entries = [
+        build_pinned_in_memory_entry(),
+        build_used_by_pending_task_entry(),
+        build_captured_in_object_entry(),
+        build_actor_handle_entry(),
+        build_local_reference_entry(),
+        build_local_reference_entry()
+    ]
+    memory_table = MemoryTable(entries)
+    assert len(memory_table.group) == 1
+    assert memory_table.summary["total_actor_handles"] == 1
+    assert memory_table.summary["total_captured_in_objects"] == 1
+    assert memory_table.summary["total_local_ref_count"] == 2
+    assert memory_table.summary[
+        "total_object_size"] == len(entries) * OBJECT_SIZE
+    assert memory_table.summary["total_pinned_in_memory"] == 1
+    assert memory_table.summary["total_used_by_pending_task"] == 1
+
+
+def test_memory_table_sort_by_pid():
+    unsort = [1, 3, 2]
+    entries = [build_entry(pid=pid) for pid in unsort]
+    memory_table = MemoryTable(entries, sort_by_type=SortingType.PID)
+    sort = sorted(unsort)
+    for pid, entry in zip(sort, memory_table.table):
+        assert pid == entry.pid
+
+
+def test_memory_table_sort_by_reference_type():
+    unsort = [
+        ReferenceType.USED_BY_PENDING_TASK, ReferenceType.LOCAL_REFERENCE,
+        ReferenceType.LOCAL_REFERENCE, ReferenceType.PINNED_IN_MEMORY
+    ]
+    entries = [
+        build_entry(reference_type=reference_type) for reference_type in unsort
+    ]
+    memory_table = MemoryTable(
+        entries, sort_by_type=SortingType.REFERENCE_TYPE)
+    sort = sorted(unsort)
+    for reference_type, entry in zip(sort, memory_table.table):
+        assert reference_type == entry.reference_type
+
+
+def test_memory_table_sort_by_object_size():
+    unsort = [312, 214, -1, 1244, 642]
+    entries = [build_entry(object_size=object_size) for object_size in unsort]
+    memory_table = MemoryTable(entries, sort_by_type=SortingType.OBJECT_SIZE)
+    sort = sorted(unsort)
+    for object_size, entry in zip(sort, memory_table.table):
+        assert object_size == entry.object_size
+
+
+def test_group_by():
+    node_second = "127.0.0.2"
+    node_first = "127.0.0.1"
+    entries = [
+        build_entry(node_address=node_second, pid=2),
+        build_entry(node_address=node_second, pid=1),
+        build_entry(node_address=node_first, pid=2),
+        build_entry(node_address=node_first, pid=1)
+    ]
+    memory_table = MemoryTable(entries)
+
+    # Make sure it is correctly grouped
+    assert node_first in memory_table.group
+    assert node_second in memory_table.group
+
+    # make sure pid is sorted in the right order.
+    for group_key, group_memory_table in memory_table.group.items():
+        pid = 1
+        for entry in group_memory_table.table:
+            assert pid == entry.pid
+            pid += 1
+
+
 if __name__ == "__main__":
     import pytest
     import sys