def transform_ray_dag_to_serve_dag(
dag_node: DAGNode, node_name_generator: DAGNodeNameGenerator
):
"""
Transform a Ray DAG to a Serve DAG. Map ClassNode to DeploymentNode with
ray decorated body passed in, and ClassMethodNode to DeploymentMethodNode.
"""
if isinstance(dag_node, ClassNode):
deployment_name = node_name_generator.get_node_name(dag_node)
return DeploymentNode(
dag_node._body,
deployment_name,
dag_node.get_args(),
dag_node.get_kwargs(),
dag_node.get_options(),
# TODO: (jiaodong) Support .options(metadata=xxx) for deployment
other_args_to_resolve=dag_node.get_other_args_to_resolve(),
)
elif isinstance(dag_node, ClassMethodNode):
other_args_to_resolve = dag_node.get_other_args_to_resolve()
# TODO: (jiaodong) Need to capture DAGNodes in the parent node
parent_deployment_node = other_args_to_resolve[PARENT_CLASS_NODE_KEY]
return DeploymentMethodNode(
parent_deployment_node._deployment,
dag_node._method_name,
dag_node.get_args(),
dag_node.get_kwargs(),
dag_node.get_options(),
other_args_to_resolve=dag_node.get_other_args_to_resolve(),
)
elif isinstance(
dag_node,
FunctionNode
# TODO (jiaodong): We do not convert ray function to deployment function
# yet, revisit this later
) and dag_node.get_other_args_to_resolve().get("is_from_serve_deployment"):
deployment_name = node_name_generator.get_node_name(dag_node)
return DeploymentFunctionNode(
dag_node._body,
deployment_name,
dag_node.get_args(),
dag_node.get_kwargs(),
dag_node.get_options(),
other_args_to_resolve=dag_node.get_other_args_to_resolve(),
)
else:
# TODO: (jiaodong) Support FunctionNode or leave it as ray task
return dag_node
def dag_to_dot(dag: DAGNode):
"""Create a Dot graph from dag.
TODO(lchu):
1. add more Dot configs in kwargs,
e.g. rankdir, alignment, etc.
2. add more contents to graph,
e.g. args, kwargs and options of each node
"""
# Step 0: check dependencies and init graph
check_pydot_and_graphviz()
import pydot
graph = pydot.Dot(rankdir="LR")
# Step 1: generate unique name for each node in dag
nodes, edges = get_nodes_and_edges(dag)
name_generator = DAGNodeNameGenerator()
node_names = {}
for node in nodes:
node_names[node] = name_generator.get_node_name(node)
# Step 2: create graph with all the edges
for edge in edges:
graph.add_edge(pydot.Edge(node_names[edge[0]], node_names[edge[1]]))
# if there is only one node
if len(nodes) == 1 and len(edges) == 0:
graph.add_node(pydot.Node(node_names[nodes[0]]))
return graph