chore: import upstream snapshot with attribution

python/ray/dag/tests/experimental/test_dag_visualization.py

@@ -0,0 +1,712 @@
+import os
+import sys
+
+import pydot
+import pytest
+
+import ray
+from ray.dag import InputNode, MultiOutputNode
+from ray.tests.conftest import *  # noqa
+
+
+@pytest.fixture
+def cleanup_files():
+    """Clean up files generated during the test."""
+
+    def _cleanup_files(filename: str):
+        for ext in ["", ".png", ".pdf", ".jpeg", ".dot"]:
+            file_path = filename + ext
+            if os.path.exists(file_path):
+                os.remove(file_path)
+
+    return _cleanup_files
+
+
+def test_visualize_basic(ray_start_regular, cleanup_files):
+    """
+    Expect output or dot_source:
+        MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
+            0 -> 1 [label=SharedMemoryType]
+            1 -> 2 [label=SharedMemoryType]
+    """
+
+    @ray.remote
+    class Actor:
+        def echo(self, x):
+            return x
+
+    actor = Actor.remote()
+
+    with InputNode() as i:
+        dag = actor.echo.bind(i)
+
+    compiled_dag = dag.experimental_compile()
+
+    # Call the visualize method
+    dot_source = compiled_dag.visualize()
+
+    graphs = pydot.graph_from_dot_data(dot_source)
+    graph = graphs[0]
+
+    node_names = {node.get_name() for node in graph.get_nodes()}
+    edge_pairs = {
+        (edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
+    }
+
+    expected_nodes = {"0", "1", "2"}
+    assert expected_nodes.issubset(
+        node_names
+    ), f"Expected nodes {expected_nodes} not found."
+
+    expected_edges = {("0", "1"), ("1", "2")}
+    assert expected_edges.issubset(
+        edge_pairs
+    ), f"Expected edges {expected_edges} not found."
+
+    cleanup_files("compiled_graph")
+
+
+def test_visualize_multi_return(ray_start_regular, cleanup_files):
+    """
+    Expect output or dot_source:
+        MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
+            0 -> 1 [label=SharedMemoryType]
+            1 -> 2 [label=SharedMemoryType]
+            1 -> 3 [label=SharedMemoryType]
+            2 -> 4 [label=SharedMemoryType]
+            3 -> 4 [label=SharedMemoryType]
+    """
+
+    @ray.remote
+    class Actor:
+        @ray.method(num_returns=2)
+        def return_two(self, x):
+            return x, x + 1
+
+    actor = Actor.remote()
+
+    with InputNode() as i:
+        o1, o2 = actor.return_two.bind(i)
+        dag = MultiOutputNode([o1, o2])
+
+    compiled_dag = dag.experimental_compile()
+
+    # Get the DOT source
+    dot_source = compiled_dag.visualize()
+
+    graphs = pydot.graph_from_dot_data(dot_source)
+    graph = graphs[0]
+
+    node_names = {node.get_name() for node in graph.get_nodes()}
+    edge_pairs = {
+        (edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
+    }
+
+    expected_nodes = {"0", "1", "2", "3", "4"}
+    assert expected_nodes.issubset(
+        node_names
+    ), f"Expected nodes {expected_nodes} not found."
+
+    expected_edges = {("0", "1"), ("1", "2"), ("1", "3"), ("2", "4"), ("3", "4")}
+    assert expected_edges.issubset(
+        edge_pairs
+    ), f"Expected edges {expected_edges} not found."
+    cleanup_files("compiled_graph")
+
+
+def test_visualize_multi_return2(ray_start_regular, cleanup_files):
+    """
+    Expect output or dot_source:
+        MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
+            0 -> 1 [label=SharedMemoryType]
+            1 -> 2 [label=SharedMemoryType]
+            1 -> 3 [label=SharedMemoryType]
+            2 -> 4 [label=SharedMemoryType]
+            3 -> 5 [label=SharedMemoryType]
+            4 -> 6 [label=SharedMemoryType]
+            5 -> 6 [label=SharedMemoryType]
+    """
+
+    @ray.remote
+    class Actor:
+        @ray.method(num_returns=2)
+        def return_two(self, x):
+            return x, x + 1
+
+        def echo(self, x):
+            return x
+
+    a = Actor.remote()
+    b = Actor.remote()
+    with InputNode() as i:
+        o1, o2 = a.return_two.bind(i)
+        o3 = b.echo.bind(o1)
+        o4 = b.echo.bind(o2)
+        dag = MultiOutputNode([o3, o4])
+
+    compiled_dag = dag.experimental_compile()
+
+    # Get the DOT source
+    dot_source = compiled_dag.visualize()
+
+    graphs = pydot.graph_from_dot_data(dot_source)
+    graph = graphs[0]
+
+    node_names = {node.get_name() for node in graph.get_nodes()}
+    edge_pairs = {
+        (edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
+    }
+
+    expected_nodes = {"0", "1", "2", "3", "4", "5", "6"}
+    assert expected_nodes.issubset(
+        node_names
+    ), f"Expected nodes {expected_nodes} not found."
+
+    expected_edges = {
+        ("0", "1"),
+        ("1", "2"),
+        ("1", "3"),
+        ("2", "4"),
+        ("3", "5"),
+        ("4", "6"),
+        ("5", "6"),
+    }
+    assert expected_edges.issubset(
+        edge_pairs
+    ), f"Expected edges {expected_edges} not found."
+    cleanup_files("compiled_graph")
+
+
+def test_visualize_multi_input_nodes(ray_start_regular, cleanup_files):
+    """
+    Expect output or dot_source:
+        MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
+            0 -> 1
+            0 -> 2
+            0 -> 3
+            1 -> 4
+            2 -> 5
+            3 -> 6
+            4 -> 7
+            5 -> 7
+            6 -> 7
+    """
+
+    @ray.remote
+    class Actor:
+        def echo(self, x):
+            return x
+
+    actor = Actor.remote()
+
+    with InputNode() as inp:
+        o1 = actor.echo.bind(inp.x)
+        o2 = actor.echo.bind(inp.y)
+        o3 = actor.echo.bind(inp.z)
+        dag = MultiOutputNode([o1, o2, o3])
+
+    compiled_dag = dag.experimental_compile()
+
+    # Get the DOT source
+    dot_source = compiled_dag.visualize()
+
+    graphs = pydot.graph_from_dot_data(dot_source)
+    graph = graphs[0]
+
+    node_names = {node.get_name() for node in graph.get_nodes()}
+    edge_pairs = {
+        (edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
+    }
+
+    expected_nodes = {"0", "1", "2", "3", "4", "5", "6", "7"}
+    assert expected_nodes.issubset(
+        node_names
+    ), f"Expected nodes {expected_nodes} not found."
+
+    expected_edges = {
+        ("0", "1"),
+        ("0", "2"),
+        ("0", "3"),
+        ("1", "4"),
+        ("2", "5"),
+        ("3", "6"),
+        ("4", "7"),
+        ("5", "7"),
+        ("6", "7"),
+    }
+    assert expected_edges.issubset(
+        edge_pairs
+    ), f"Expected edges {expected_edges} not found."
+    cleanup_files("compiled_graph")
+
+
+class TestVisualizationAscii:
+
+    """Tests for the visualize_ascii method of compiled DAGs."""
+
+    @staticmethod
+    def parse_ascii_visualization(ascii_visualization: str):
+        """
+        Parses the ASCII visualization output to extract node names and edge pairs.
+
+        Args:
+            ascii_visualization: The ASCII visualization
+                output generated by the `visualize` function.
+
+        Returns:
+            tuple: A tuple containing:
+                - node_names: A set of strings representing node names.
+                - edge_pairs: A set of tuples representing edge
+                pairs with type hints.
+        """
+        import re
+
+        # Sets to store unique nodes and edges
+        node_names = set()
+        edge_pairs = set()
+
+        # Extract nodes from "Nodes Information" section
+        node_pattern = re.compile(r'^(\d+) \[label="Task \d+')
+        edge_pattern = re.compile(r"^(\d+) (--->|\+\+\+>) (\d+)")
+
+        lines = ascii_visualization.splitlines()
+        in_nodes_section = False
+        in_edges_section = False
+
+        for line in lines:
+            line = line.strip()
+
+            # Check for nodes section
+            if line.startswith("Nodes Information:"):
+                in_nodes_section = True
+                in_edges_section = False
+                continue
+
+            # Check for edges section
+            if line.startswith("Edges Information:"):
+                in_edges_section = True
+                in_nodes_section = False
+                continue
+
+            # Collect nodes
+            if in_nodes_section:
+                node_match = node_pattern.match(line)
+                if node_match:
+                    node_id = node_match.group(1)
+                    node_names.add(node_id)
+
+            # Collect edges
+            if in_edges_section:
+                edge_match = edge_pattern.match(line)
+                if edge_match:
+                    from_node, _, to_node = edge_match.groups()
+                    edge_pairs.add((from_node, to_node))
+
+        return node_names, edge_pairs
+
+    def test_visualize_ascii_basic(self, ray_start_regular):
+        """
+        Expect output:
+            Nodes Information:
+            0 [label="Task 0  InputNode"]
+            1 [label="Task 1  Actor: d6c5c4... Method: echo"]
+            2 [label="Task 2  MultiOutputNode"]
+
+            Edges Information:
+            0 ---> 1
+            1 ---> 2
+
+            Legend:
+            +++> : Represents Nccl-type data channels
+            ---> : Represents Shared Memory data channels
+
+            Experimental Graph:
+            0:InputNode
+            |
+            1:Actor_d6c5c4:echo
+            |
+            2:MultiOutputNode
+        """
+
+        @ray.remote
+        class Actor:
+            def echo(self, x):
+                return x
+
+        actor = Actor.remote()
+
+        with InputNode() as i:
+            dag = actor.echo.bind(i)
+
+        compiled_dag = dag.experimental_compile()
+
+        # Call the visualize method
+        ascii_visualization = compiled_dag.visualize(format="ascii")
+
+        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
+            ascii_visualization
+        )
+        print(node_names, edge_pairs)
+        expected_nodes = {"0", "1", "2"}
+        assert expected_nodes.issubset(
+            node_names
+        ), f"Expected nodes {expected_nodes} not found."
+
+        expected_edges = {("0", "1"), ("1", "2")}
+        assert expected_edges.issubset(
+            edge_pairs
+        ), f"Expected edges {expected_edges} not found."
+
+    def test_visualize_ascii_multi_return(self, ray_start_regular):
+        """
+        Expect output:
+            Nodes Information:
+            0 [label="Task 0  InputNode"]
+            1 [label="Task 1  Actor: 885f1d... Method: return_two"]
+            2 [label="Task 2  ClassMethodOutputNode[0]"]
+            3 [label="Task 3  ClassMethodOutputNode[1]"]
+            4 [label="Task 4  MultiOutputNode"]
+
+            Edges Information:
+            0 ---> 1
+            1 ---> 2
+            1 ---> 3
+            2 ---> 4
+            3 ---> 4
+
+            Legend:
+            +++> : Represents Nccl-type data channels
+            ---> : Represents Shared Memory data channels
+
+            Graph Built:
+            0:InputNode
+            |
+            1:Actor_885f1d:return_two
+            |---------------------------->|
+            2:Output[0]                   3:Output[1]
+            |<----------------------------|
+            4:MultiOutputNode
+        """
+
+        @ray.remote
+        class Actor:
+            @ray.method(num_returns=2)
+            def return_two(self, x):
+                return x, x + 1
+
+        actor = Actor.remote()
+
+        with InputNode() as i:
+            o1, o2 = actor.return_two.bind(i)
+            dag = MultiOutputNode([o1, o2])
+
+        compiled_dag = dag.experimental_compile()
+
+        ascii_visualization = compiled_dag.visualize(format="ascii")
+
+        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
+            ascii_visualization
+        )
+
+        expected_nodes = {"0", "1", "2", "3", "4"}
+        assert expected_nodes.issubset(
+            node_names
+        ), f"Expected nodes {expected_nodes} not found."
+
+        expected_edges = {("0", "1"), ("1", "2"), ("1", "3"), ("2", "4"), ("3", "4")}
+        assert expected_edges.issubset(
+            edge_pairs
+        ), f"Expected edges {expected_edges} not found."
+
+    def test_visualize_ascii_multi_return2(self, ray_start_regular):
+        """
+        Expect output:
+            Nodes Information:
+            0 [label="Task 0  InputNode"]
+            1 [label="Task 1  Actor: f3e919... Method: return_two"]
+            2 [label="Task 2  ClassMethodOutputNode[0]"]
+            3 [label="Task 3  ClassMethodOutputNode[1]"]
+            4 [label="Task 4  Actor: 15ec69... Method: echo"]
+            5 [label="Task 5  Actor: 15ec69... Method: echo"]
+            6 [label="Task 6  MultiOutputNode"]
+
+            Edges Information:
+            0 ---> 1
+            1 ---> 2
+            1 ---> 3
+            2 ---> 4
+            3 ---> 5
+            4 ---> 6
+            5 ---> 6
+
+            Legend:
+            +++> : Represents Nccl-type data channels
+            ---> : Represents Shared Memory data channels
+
+            Graph Built:
+            0:InputNode
+            |
+            1:Actor_f3e919:return_two
+            |---------------------------->|
+            2:Output[0]                   3:Output[1]
+            |                             |
+            4:Actor_15ec69:echo           5:Actor_15ec69:echo
+            |<----------------------------|
+            6:MultiOutputNode
+        """
+
+        @ray.remote
+        class Actor:
+            @ray.method(num_returns=2)
+            def return_two(self, x):
+                return x, x + 1
+
+            def echo(self, x):
+                return x
+
+        a = Actor.remote()
+        b = Actor.remote()
+        with InputNode() as i:
+            o1, o2 = a.return_two.bind(i)
+            o3 = b.echo.bind(o1)
+            o4 = b.echo.bind(o2)
+            dag = MultiOutputNode([o3, o4])
+
+        compiled_dag = dag.experimental_compile()
+
+        ascii_visualization = compiled_dag.visualize(format="ascii")
+
+        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
+            ascii_visualization
+        )
+
+        expected_nodes = {"0", "1", "2", "3", "4", "5", "6"}
+        assert expected_nodes.issubset(
+            node_names
+        ), f"Expected nodes {expected_nodes} not found."
+
+        expected_edges = {
+            ("0", "1"),
+            ("1", "2"),
+            ("1", "3"),
+            ("2", "4"),
+            ("3", "5"),
+            ("4", "6"),
+            ("5", "6"),
+        }
+        assert expected_edges.issubset(
+            edge_pairs
+        ), f"Expected edges {expected_edges} not found."
+
+    def test_visualize_ascii_complicate(self, ray_start_regular):
+        """
+        Expect output:
+            Nodes Information:
+            0 [label="Task 0  InputNode"]
+            1 [label="Task 1  Actor: 54777d... Method: return_three"]
+            2 [label="Task 2  ClassMethodOutputNode[0]"]
+            3 [label="Task 3  ClassMethodOutputNode[1]"]
+            4 [label="Task 4  ClassMethodOutputNode[2]"]
+            5 [label="Task 5  Actor: c927c9... Method: echo"]
+            6 [label="Task 6  Actor: c927c9... Method: echo"]
+            7 [label="Task 7  Actor: c927c9... Method: return_two"]
+            8 [label="Task 8  MultiOutputNode"]
+            9 [label="Task 9  ClassMethodOutputNode[0]"]
+            10 [label="Task 10  ClassMethodOutputNode[1]"]
+
+            Edges Information:
+            0 ---> 1
+            1 ---> 2
+            1 ---> 3
+            1 ---> 4
+            2 ---> 5
+            3 ---> 6
+            4 ---> 7
+            5 ---> 8
+            6 ---> 8
+            9 ---> 8
+            10 ---> 8
+            7 ---> 9
+            7 ---> 10
+
+            Legend:
+            +++> : Represents Nccl-type data channels
+            ---> : Represents Shared Memory data channels
+
+            Graph Built:
+            0:InputNode
+            |
+            1:Actor_54777d:return_three
+            |---------------------------->|---------------------------->|                                                  # noqa
+            2:Output[0]                   3:Output[1]                   4:Output[2]                                        # noqa
+            |                             |                             |                                                  # noqa
+            5:Actor_c927c9:echo           6:Actor_c927c9:echo           7:Actor_c927c9:return_two                          # noqa
+            |                             |                             |---------------------------->|                    # noqa
+            |                             |                             9:Output[0]                   10:Output[1]         # noqa
+            |<----------------------------|-----------------------------|-----------------------------|                    # noqa
+            8:MultiOutputNode
+        """
+
+        @ray.remote
+        class Actor:
+            @ray.method(num_returns=3)
+            def return_three(self, x):
+                return x, x + 1, x + 2
+
+            def echo(self, x):
+                return x
+
+            @ray.method(num_returns=2)
+            def return_two(self, x):
+                return x, x + 1
+
+        a = Actor.remote()
+        b = Actor.remote()
+        with InputNode() as i:
+            o1, o2, o3 = a.return_three.bind(i)
+            o4 = b.echo.bind(o1)
+            o5 = b.echo.bind(o2)
+            o6, o7 = b.return_two.bind(o3)
+            dag = MultiOutputNode([o4, o5, o6, o7])
+
+        compiled_dag = dag.experimental_compile()
+
+        ascii_visualization = compiled_dag.visualize(format="ascii")
+
+        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
+            ascii_visualization
+        )
+
+        expected_nodes = {"0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10"}
+        assert expected_nodes.issubset(
+            node_names
+        ), f"Expected nodes {expected_nodes} not found."
+
+        expected_edges = {
+            ("0", "1"),
+            ("1", "2"),
+            ("1", "3"),
+            ("1", "4"),
+            ("2", "5"),
+            ("3", "6"),
+            ("4", "7"),
+            ("5", "8"),
+            ("6", "8"),
+            ("9", "8"),
+            ("10", "8"),
+            ("7", "9"),
+            ("7", "10"),
+        }
+        assert expected_edges.issubset(
+            edge_pairs
+        ), f"Expected edges {expected_edges} not found."
+
+    def test_visualize_ascii_cross_line(self, ray_start_regular):
+        """
+        Expect output:
+            Nodes Information:
+            0 [label="Task 0  InputNode"]
+            1 [label="Task 1  Actor: 84835a... Method: return_three"]
+            2 [label="Task 2  ClassMethodOutputNode[0]"]
+            3 [label="Task 3  ClassMethodOutputNode[1]"]
+            4 [label="Task 4  ClassMethodOutputNode[2]"]
+            5 [label="Task 5  Actor: 02a6a1... Method: echo"]
+            6 [label="Task 6  Actor: 02a6a1... Method: return_two"]
+            7 [label="Task 7  Actor: 02a6a1... Method: echo"]
+            8 [label="Task 8  MultiOutputNode"]
+            9 [label="Task 9  ClassMethodOutputNode[0]"]
+            10 [label="Task 10  ClassMethodOutputNode[1]"]
+
+            Edges Information:
+            0 ---> 1
+            1 ---> 2
+            1 ---> 3
+            1 ---> 4
+            2 ---> 5
+            3 ---> 6
+            4 ---> 7
+            5 ---> 8
+            7 ---> 8
+            9 ---> 8
+            10 ---> 8
+            6 ---> 9
+            6 ---> 10
+
+            Legend:
+            +++> : Represents Nccl-type data channels
+            ---> : Represents Shared Memory data channels
+
+            Graph Built:
+            0:InputNode
+            |
+            1:Actor_84835a:return_three
+            |---------------------------->|---------------------------->|                            # noqa
+            2:Output[0]                   3:Output[1]                   4:Output[2]                  # noqa
+            |                             |                             |                            # noqa
+            5:Actor_02a6a1:echo           6:Actor_02a6a1:return_two     7:Actor_02a6a1:echo          # noqa
+            |                             |---------------------------->|                            # noqa
+            |                             9:Output[0]                   10:Output[1]                 # noqa
+            |<----------------------------------------------------------|                            # noqa
+            8:MultiOutputNod
+        """
+
+        @ray.remote
+        class Actor:
+            @ray.method(num_returns=3)
+            def return_three(self, x):
+                return x, x + 1, x + 2
+
+            def echo(self, x):
+                return x
+
+            @ray.method(num_returns=2)
+            def return_two(self, x):
+                return x, x + 1
+
+        a = Actor.remote()
+        b = Actor.remote()
+        with InputNode() as i:
+            o1, o2, o3 = a.return_three.bind(i)
+            o4 = b.echo.bind(o1)
+            o5 = b.echo.bind(o3)
+            o6, o7 = b.return_two.bind(o2)
+            dag = MultiOutputNode([o4, o5, o6, o7])
+
+        compiled_dag = dag.experimental_compile()
+
+        ascii_visualization = compiled_dag.visualize(format="ascii")
+
+        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
+            ascii_visualization
+        )
+
+        expected_nodes = {"0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10"}
+        assert expected_nodes.issubset(
+            node_names
+        ), f"Expected nodes {expected_nodes} not found."
+
+        expected_edges = {
+            ("0", "1"),
+            ("1", "2"),
+            ("1", "3"),
+            ("1", "4"),
+            ("2", "5"),
+            ("3", "6"),
+            ("4", "7"),
+            ("5", "8"),
+            ("7", "8"),
+            ("9", "8"),
+            ("10", "8"),
+            ("6", "9"),
+            ("6", "10"),
+        }
+        assert expected_edges.issubset(
+            edge_pairs
+        ), f"Expected edges {expected_edges} not found."
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))