def get_node_name(self, node: DAGNode):
if isinstance(node, InputNode):
node_name = "INPUT_NODE"
elif isinstance(node, InputAttributeNode):
node_name = "INPUT_ATTRIBUTE_NODE"
elif isinstance(node, ClassMethodNode):
node_name = node.get_options().get("name",
None) or node._method_name
elif isinstance(node, (ClassNode, FunctionNode)):
node_name = node.get_options().get("name",
None) or node._body.__name__
# we use instance class name check here to avoid importing ServeNodes as
# serve components are not included in Ray Core.
elif type(node).__name__ in ("DeploymentNode",
"DeploymentFunctionNode"):
node_name = node.get_deployment_name()
elif type(node).__name__ == "DeploymentMethodNode":
node_name = node.get_deployment_method_name()
else:
raise ValueError(
"get_node_name() should only be called on DAGNode instances.")
if node_name not in self.name_to_suffix:
self.name_to_suffix[node_name] = 0
return node_name
else:
self.name_to_suffix[node_name] += 1
suffix_num = self.name_to_suffix[node_name]
return f"{node_name}_{suffix_num}"
def transform_ray_dag_to_serve_dag(
dag_node: DAGNode, deployment_name_generator: DeploymentNameGenerator):
"""
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 = deployment_name_generator.get_deployment_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 = deployment_name_generator.get_deployment_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