Files
ray-project--ray/python/ray/dag/tests/experimental/test_dag_visualization.py
T
2026-07-13 13:17:40 +08:00

713 lines
21 KiB
Python

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__]))