Пример #1
0
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
Пример #2
0
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
Пример #3
0
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,
    )