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 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