def __init__(
self,
func_body: Union[Callable, str],
deployment_name,
func_args,
func_kwargs,
func_options,
other_args_to_resolve=None,
):
self._body = func_body
self._deployment_name = deployment_name
super().__init__(
func_args,
func_kwargs,
func_options,
other_args_to_resolve=other_args_to_resolve,
)
if "deployment_schema" in self._bound_other_args_to_resolve:
deployment_schema: DeploymentSchema = self._bound_other_args_to_resolve[
"deployment_schema"]
deployment_shell = schema_to_deployment(deployment_schema)
# Prefer user specified name to override the generated one.
if (inspect.isfunction(func_body)
and deployment_shell.name != func_body.__name__):
self._deployment_name = deployment_shell.name
# Set the route prefix, prefer the one user supplied,
# otherwise set it to /deployment_name
if (deployment_shell.route_prefix is None
or deployment_shell.route_prefix !=
f"/{deployment_shell.name}"):
route_prefix = deployment_shell.route_prefix
else:
route_prefix = f"/{deployment_name}"
self._deployment = deployment_shell.options(
func_or_class=func_body,
name=self._deployment_name,
init_args=(),
init_kwargs=dict(),
route_prefix=route_prefix,
)
else:
self._deployment: Deployment = Deployment(
func_body,
deployment_name,
DeploymentConfig(),
init_args=tuple(),
init_kwargs=dict(),
ray_actor_options=func_options,
_internal=True,
)
# TODO (jiaodong): Polish with async handle support later
self._deployment_handle = RayServeLazySyncHandle(deployment_name)
def _get_serve_deployment_handle(
self,
deployment: Deployment,
bound_other_args_to_resolve: Dict[str, Any],
) -> Union[RayServeHandle, RayServeSyncHandle]:
"""
Return a sync or async handle of the encapsulated Deployment based on
config.
Args:
deployment (Deployment): Deployment instance wrapped in the DAGNode.
bound_other_args_to_resolve (Dict[str, Any]): Contains args used
to configure DeploymentNode.
Returns:
RayServeHandle: Default and catch-all is to return sync handle.
return async handle only if user explicitly set
USE_SYNC_HANDLE_KEY with value of False.
"""
# TODO (jiaodong): Support configurable async handle
if USE_SYNC_HANDLE_KEY not in bound_other_args_to_resolve:
# Return sync RayServeLazySyncHandle
return RayServeLazySyncHandle(deployment.name)
elif bound_other_args_to_resolve.get(USE_SYNC_HANDLE_KEY) is True:
# Return sync RayServeSyncHandle
return deployment.get_handle(sync=True)
elif bound_other_args_to_resolve.get(USE_SYNC_HANDLE_KEY) is False:
# Return async RayServeHandle
return deployment.get_handle(sync=False)
else:
raise ValueError(
f"{USE_SYNC_HANDLE_KEY} should only be set with a boolean value."
)
def dagnode_from_json(input_json: Any) -> Union[DAGNode, RayServeHandle, Any]:
"""
Decode a DAGNode from given input json dictionary. JSON serialization is
only used and enforced in ray serve from ray core API authored DAGNode(s).
Covers both RayServeHandle and DAGNode types.
Assumptions:
- User object's JSON dict does not have keys that collide with our
reserved DAGNODE_TYPE_KEY
- RayServeHandle and Deployment can be re-constructed without losing
states needed for their functionality or correctness.
- DAGNode type can be re-constructed with new stable_uuid upon each
deserialization without effective correctness of execution.
- Only exception is ClassNode used as parent of ClassMethodNode
that we perserve the same parent node.
- .options() does not contain any DAGNode type
"""
node_type_to_cls = {
# Ray DAG Inputs
InputNode.__name__: InputNode,
InputAttributeNode.__name__: InputAttributeNode,
# Deployment graph execution nodes
DeploymentExecutorNode.__name__: DeploymentExecutorNode,
DeploymentMethodExecutorNode.__name__: DeploymentMethodExecutorNode,
DeploymentFunctionExecutorNode.__name__:
DeploymentFunctionExecutorNode,
}
# Deserialize RayServeHandle type
if SERVE_HANDLE_JSON_KEY in input_json:
return serve_handle_from_json_dict(input_json)
# Base case for plain objects
elif DAGNODE_TYPE_KEY not in input_json:
return input_json
elif input_json[DAGNODE_TYPE_KEY] == RayServeDAGHandle.__name__:
return RayServeDAGHandle(input_json["dag_node_json"])
elif input_json[DAGNODE_TYPE_KEY] == "DeploymentSchema":
return DeploymentSchema.parse_obj(input_json["schema"])
elif input_json[DAGNODE_TYPE_KEY] == RayServeLazySyncHandle.__name__:
return RayServeLazySyncHandle(
input_json["deployment_name"],
HandleOptions(input_json["handle_options_method_name"]),
)
# Deserialize DAGNode type
elif input_json[DAGNODE_TYPE_KEY] in node_type_to_cls:
return node_type_to_cls[input_json[DAGNODE_TYPE_KEY]].from_json(
input_json)
else:
# Class and Function nodes require original module as body.
module_name, attr_name = parse_import_path(input_json["import_path"])
module = getattr(import_module(module_name), attr_name)
if input_json[DAGNODE_TYPE_KEY] == FunctionNode.__name__:
return FunctionNode.from_json(input_json, module)
elif input_json[DAGNODE_TYPE_KEY] == ClassNode.__name__:
return ClassNode.from_json(input_json, module)
def replace_with_handle(node):
if isinstance(node, DeploymentNode):
return RayServeLazySyncHandle(node._deployment.name)
elif isinstance(node, DeploymentExecutorNode):
return node._deployment_handle
elif isinstance(
node,
(
DeploymentMethodNode,
DeploymentMethodExecutorNode,
DeploymentFunctionNode,
DeploymentFunctionExecutorNode,
),
):
from ray.serve.pipeline.json_serde import DAGNodeEncoder
serve_dag_root_json = json.dumps(node, cls=DAGNodeEncoder)
return RayServeDAGHandle(serve_dag_root_json)
def __init__(
self,
# For serve structured deployment, deployment body can be import path
# to the class or function instead.
deployment: Deployment,
deployment_init_args: Tuple[Any],
deployment_init_kwargs: Dict[str, Any],
ray_actor_options: Dict[str, Any],
other_args_to_resolve: Optional[Dict[str, Any]] = None,
):
# Assign instance variables in base class constructor.
super().__init__(
deployment_init_args,
deployment_init_kwargs,
ray_actor_options,
other_args_to_resolve=other_args_to_resolve,
)
self._deployment = deployment
self._deployment_handle = RayServeLazySyncHandle(self._deployment.name)
class DeploymentFunctionNode(DAGNode):
"""Represents a function node decorated by @serve.deployment in a serve DAG."""
def __init__(
self,
func_body: Union[Callable, str],
deployment_name,
func_args,
func_kwargs,
func_options,
other_args_to_resolve=None,
):
self._body = func_body
self._deployment_name = deployment_name
super().__init__(
func_args,
func_kwargs,
func_options,
other_args_to_resolve=other_args_to_resolve,
)
if "deployment_schema" in self._bound_other_args_to_resolve:
deployment_schema: DeploymentSchema = self._bound_other_args_to_resolve[
"deployment_schema"]
deployment_shell = schema_to_deployment(deployment_schema)
# Prefer user specified name to override the generated one.
if (inspect.isfunction(func_body)
and deployment_shell.name != func_body.__name__):
self._deployment_name = deployment_shell.name
# Set the route prefix, prefer the one user supplied,
# otherwise set it to /deployment_name
if (deployment_shell.route_prefix is None
or deployment_shell.route_prefix !=
f"/{deployment_shell.name}"):
route_prefix = deployment_shell.route_prefix
else:
route_prefix = f"/{deployment_name}"
self._deployment = deployment_shell.options(
func_or_class=func_body,
name=self._deployment_name,
init_args=(),
init_kwargs=dict(),
route_prefix=route_prefix,
)
else:
self._deployment: Deployment = Deployment(
func_body,
deployment_name,
DeploymentConfig(),
init_args=tuple(),
init_kwargs=dict(),
ray_actor_options=func_options,
_internal=True,
)
# TODO (jiaodong): Polish with async handle support later
self._deployment_handle = RayServeLazySyncHandle(deployment_name)
def _copy_impl(
self,
new_args: List[Any],
new_kwargs: Dict[str, Any],
new_options: Dict[str, Any],
new_other_args_to_resolve: Dict[str, Any],
):
return DeploymentFunctionNode(
self._body,
self._deployment_name,
new_args,
new_kwargs,
new_options,
other_args_to_resolve=new_other_args_to_resolve,
)
def _execute_impl(self, *args, **kwargs) -> ObjectRef:
"""Executor of DeploymentFunctionNode by calling .remote() on the
deployment handle.
Deployment method always default to __call__.
"""
return self._deployment_handle.remote(*self._bound_args,
**self._bound_kwargs)
def __str__(self) -> str:
return get_dag_node_str(self, str(self._body))
def get_deployment_name(self):
return self._deployment_name
def get_import_path(self):
if ("is_from_serve_deployment" in self._bound_other_args_to_resolve
): # built by serve top level api, this is ignored for serve.run
return "dummy"
return get_deployment_import_path(self._deployment)
def to_json(self, encoder_cls) -> Dict[str, Any]:
json_dict = super().to_json_base(encoder_cls,
DeploymentFunctionNode.__name__)
json_dict["import_path"] = self.get_import_path()
json_dict["deployment_name"] = self.get_deployment_name()
return json_dict
@classmethod
def from_json(cls, input_json, object_hook=None):
assert input_json[DAGNODE_TYPE_KEY] == DeploymentFunctionNode.__name__
args_dict = super().from_json_base(input_json, object_hook=object_hook)
node = cls(
input_json["import_path"],
input_json["deployment_name"],
args_dict["args"],
args_dict["kwargs"],
args_dict["options"],
other_args_to_resolve=args_dict["other_args_to_resolve"],
)
node._stable_uuid = args_dict["uuid"]
return node
def replace_with_handle(node):
if isinstance(node, DeploymentNode):
return RayServeLazySyncHandle(node._deployment.name)
elif isinstance(node, DeploymentExecutorNode):
return node._deployment_handle
def __init__(
self,
# For serve structured deployment, deployment body can be import path
# to the class or function instead.
func_or_class: Union[Callable, str],
deployment_name: str,
deployment_init_args: Tuple[Any],
deployment_init_kwargs: Dict[str, Any],
ray_actor_options: Dict[str, Any],
other_args_to_resolve: Optional[Dict[str, Any]] = None,
):
# Assign instance variables in base class constructor.
super().__init__(
deployment_init_args,
deployment_init_kwargs,
ray_actor_options,
other_args_to_resolve=other_args_to_resolve,
)
# Deployment can be passed into other DAGNodes as init args. This is
# supported pattern in ray DAG that user can instantiate and pass class
# instances as init args to others.
# However in ray serve we send init args via .remote() that requires
# pickling, and all DAGNode types are not picklable by design.
# Thus we need convert all DeploymentNode used in init args into
# deployment handles (executable and picklable) in ray serve DAG to make
# serve DAG end to end executable.
# TODO(jiaodong): This part does some magic for DAGDriver and will throw
# error with weird pickle replace table error. Move this out.
def replace_with_handle(node):
if isinstance(node, DeploymentNode):
return RayServeLazySyncHandle(node._deployment.name)
elif isinstance(node, DeploymentExecutorNode):
return node._deployment_handle
(
replaced_deployment_init_args,
replaced_deployment_init_kwargs,
) = self.apply_functional(
[deployment_init_args, deployment_init_kwargs],
predictate_fn=lambda node: isinstance(
node,
# We need to match and replace all DAGNodes even though they
# could be None, because no DAGNode replacement should run into
# re-resolved child DAGNodes, otherwise with KeyError
(
DeploymentNode,
DeploymentMethodNode,
DeploymentFunctionNode,
DeploymentExecutorNode,
DeploymentFunctionExecutorNode,
DeploymentMethodExecutorNode,
),
),
apply_fn=replace_with_handle,
)
if "deployment_schema" in self._bound_other_args_to_resolve:
deployment_schema: DeploymentSchema = self._bound_other_args_to_resolve[
"deployment_schema"
]
deployment_shell = schema_to_deployment(deployment_schema)
# Prefer user specified name to override the generated one.
if (
inspect.isclass(func_or_class)
and deployment_shell.name != func_or_class.__name__
):
deployment_name = deployment_shell.name
# Set the route prefix, prefer the one user supplied,
# otherwise set it to /deployment_name
if (
deployment_shell.route_prefix is None
or deployment_shell.route_prefix != f"/{deployment_shell.name}"
):
route_prefix = deployment_shell.route_prefix
else:
route_prefix = f"/{deployment_name}"
self._deployment = deployment_shell.options(
func_or_class=func_or_class,
name=deployment_name,
init_args=replaced_deployment_init_args,
init_kwargs=replaced_deployment_init_kwargs,
route_prefix=route_prefix,
)
else:
self._deployment: Deployment = Deployment(
func_or_class,
deployment_name,
# TODO: (jiaodong) Support deployment config from user input
DeploymentConfig(),
init_args=replaced_deployment_init_args,
init_kwargs=replaced_deployment_init_kwargs,
ray_actor_options=ray_actor_options,
_internal=True,
)
self._deployment_handle = RayServeLazySyncHandle(self._deployment.name)