def transform_serve_dag_to_serve_executor_dag(serve_dag_root_node: DAGNode):
"""Given a runnable serve dag with deployment init args and options
processed, transform into an equivalent, but minimal dag optimized for
execution.
"""
if isinstance(serve_dag_root_node, DeploymentNode):
return DeploymentExecutorNode(
serve_dag_root_node._deployment_handle,
serve_dag_root_node.get_args(),
serve_dag_root_node.get_kwargs(),
)
elif isinstance(serve_dag_root_node, DeploymentMethodNode):
return DeploymentMethodExecutorNode(
# Deployment method handle
serve_dag_root_node._deployment_method_name,
serve_dag_root_node.get_args(),
serve_dag_root_node.get_kwargs(),
other_args_to_resolve=serve_dag_root_node.get_other_args_to_resolve(),
)
elif isinstance(serve_dag_root_node, DeploymentFunctionNode):
return DeploymentFunctionExecutorNode(
serve_dag_root_node._deployment_handle,
serve_dag_root_node.get_args(),
serve_dag_root_node.get_kwargs(),
)
else:
return serve_dag_root_node
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
def generate_executor_dag_driver_deployment(
serve_executor_dag_root_node: DAGNode, original_driver_deployment: Deployment
):
"""Given a transformed minimal execution serve dag, and original DAGDriver
deployment, generate new DAGDriver deployment that uses new serve executor
dag as init_args.
Args:
serve_executor_dag_root_node (DeploymentExecutorNode): Transformed
executor serve dag with only barebone deployment handles.
original_driver_deployment (Deployment): User's original DAGDriver
deployment that wrapped Ray DAG as init args.
Returns:
executor_dag_driver_deployment (Deployment): New DAGDriver deployment
with executor serve dag as init args.
"""
def replace_with_handle(node):
if isinstance(node, DeploymentExecutorNode):
return node._deployment_handle
elif isinstance(
node,
(
DeploymentMethodExecutorNode,
DeploymentFunctionExecutorNode,
),
):
serve_dag_root_json = json.dumps(node, cls=DAGNodeEncoder)
return RayServeDAGHandle(serve_dag_root_json)
(
replaced_deployment_init_args,
replaced_deployment_init_kwargs,
) = serve_executor_dag_root_node.apply_functional(
[
serve_executor_dag_root_node.get_args(),
serve_executor_dag_root_node.get_kwargs(),
],
predictate_fn=lambda node: isinstance(
node,
(
DeploymentExecutorNode,
DeploymentFunctionExecutorNode,
DeploymentMethodExecutorNode,
),
),
apply_fn=replace_with_handle,
)
return original_driver_deployment.options(
init_args=replaced_deployment_init_args,
init_kwargs=replaced_deployment_init_kwargs,
)