wehub-resource-sync
337 lines
12 KiB
Python
337 lines
12 KiB
Python
from typing import Any, Dict, List, Optional, Union
|
|
|
|
from ray.dag import DAGNode
|
|
from ray.dag.format_utils import get_dag_node_str
|
|
from ray.experimental.gradio_utils import type_to_string
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
IN_CONTEXT_MANAGER = "__in_context_manager__"
|
|
|
|
|
|
@DeveloperAPI
|
|
class InputNode(DAGNode):
|
|
r"""Ray dag node used in DAG building API to mark entrypoints of a DAG.
|
|
|
|
Should only be function or class method. A DAG can have multiple
|
|
entrypoints, but only one instance of InputNode exists per DAG, shared
|
|
among all DAGNodes.
|
|
|
|
Example:
|
|
|
|
.. code-block::
|
|
|
|
m1.forward
|
|
/ \
|
|
dag_input ensemble -> dag_output
|
|
\ /
|
|
m2.forward
|
|
|
|
In this pipeline, each user input is broadcasted to both m1.forward and
|
|
m2.forward as first stop of the DAG, and authored like
|
|
|
|
.. code-block:: python
|
|
|
|
import ray
|
|
|
|
@ray.remote
|
|
class Model:
|
|
def __init__(self, val):
|
|
self.val = val
|
|
def forward(self, input):
|
|
return self.val * input
|
|
|
|
@ray.remote
|
|
def combine(a, b):
|
|
return a + b
|
|
|
|
with InputNode() as dag_input:
|
|
m1 = Model.bind(1)
|
|
m2 = Model.bind(2)
|
|
m1_output = m1.forward.bind(dag_input[0])
|
|
m2_output = m2.forward.bind(dag_input.x)
|
|
ray_dag = combine.bind(m1_output, m2_output)
|
|
|
|
# Pass mix of args and kwargs as input.
|
|
ray_dag.execute(1, x=2) # 1 sent to m1, 2 sent to m2
|
|
|
|
# Alternatively user can also pass single data object, list or dict
|
|
# and access them via list index, object attribute or dict key str.
|
|
ray_dag.execute(UserDataObject(m1=1, m2=2))
|
|
# dag_input.m1, dag_input.m2
|
|
ray_dag.execute([1, 2])
|
|
# dag_input[0], dag_input[1]
|
|
ray_dag.execute({"m1": 1, "m2": 2})
|
|
# dag_input["m1"], dag_input["m2"]
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*args: Any,
|
|
input_type: Optional[Union[type, Dict[Union[int, str], type]]] = None,
|
|
_other_args_to_resolve: Optional[Dict[str, Any]] = None,
|
|
**kwargs: Any,
|
|
):
|
|
"""InputNode should only take attributes of validating and converting
|
|
input data rather than the input data itself. User input should be
|
|
provided via `ray_dag.execute(user_input)`.
|
|
|
|
Args:
|
|
*args: Reserved; passing any positional argument raises ``ValueError``.
|
|
input_type: Describes the data type of inputs user will be giving.
|
|
- if given through singular InputNode: type of InputNode
|
|
- if given through InputAttributeNodes: map of key -> type
|
|
Used when deciding what Gradio block to represent the input nodes with.
|
|
_other_args_to_resolve: Internal only to keep InputNode's execution
|
|
context throughput pickling, replacement and serialization.
|
|
User should not use or pass this field.
|
|
**kwargs: Reserved; passing any keyword argument raises ``ValueError``.
|
|
"""
|
|
if len(args) != 0 or len(kwargs) != 0:
|
|
raise ValueError("InputNode should not take any args or kwargs.")
|
|
|
|
self.input_attribute_nodes = {}
|
|
|
|
self.input_type = input_type
|
|
if input_type is not None and isinstance(input_type, type):
|
|
if _other_args_to_resolve is None:
|
|
_other_args_to_resolve = {}
|
|
_other_args_to_resolve["result_type_string"] = type_to_string(input_type)
|
|
|
|
super().__init__([], {}, {}, other_args_to_resolve=_other_args_to_resolve)
|
|
|
|
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 InputNode(_other_args_to_resolve=new_other_args_to_resolve)
|
|
|
|
def _execute_impl(self, *args, **kwargs):
|
|
"""Executor of InputNode."""
|
|
# Catch and assert singleton context at dag execution time.
|
|
assert self._in_context_manager(), (
|
|
"InputNode is a singleton instance that should be only used in "
|
|
"context manager for dag building and execution. See the docstring "
|
|
"of class InputNode for examples."
|
|
)
|
|
# If user only passed in one value, for simplicity we just return it.
|
|
if len(args) == 1 and len(kwargs) == 0:
|
|
return args[0]
|
|
|
|
return DAGInputData(*args, **kwargs)
|
|
|
|
def _in_context_manager(self) -> bool:
|
|
"""Return if InputNode is created in context manager."""
|
|
if (
|
|
not self._bound_other_args_to_resolve
|
|
or IN_CONTEXT_MANAGER not in self._bound_other_args_to_resolve
|
|
):
|
|
return False
|
|
else:
|
|
return self._bound_other_args_to_resolve[IN_CONTEXT_MANAGER]
|
|
|
|
def set_context(self, key: str, val: Any):
|
|
"""Set field in parent DAGNode attribute that can be resolved in both
|
|
pickle and JSON serialization
|
|
"""
|
|
self._bound_other_args_to_resolve[key] = val
|
|
|
|
def __str__(self) -> str:
|
|
return get_dag_node_str(self, "__InputNode__")
|
|
|
|
def __getattr__(self, key: str):
|
|
assert isinstance(
|
|
key, str
|
|
), "Please only access dag input attributes with str key."
|
|
if key not in self.input_attribute_nodes:
|
|
self.input_attribute_nodes[key] = InputAttributeNode(
|
|
self, key, "__getattr__"
|
|
)
|
|
return self.input_attribute_nodes[key]
|
|
|
|
def __getitem__(self, key: Union[int, str]) -> Any:
|
|
assert isinstance(key, (str, int)), (
|
|
"Please only use int index or str as first-level key to "
|
|
"access fields of dag input."
|
|
)
|
|
|
|
input_type = None
|
|
if self.input_type is not None and key in self.input_type:
|
|
input_type = type_to_string(self.input_type[key])
|
|
|
|
if key not in self.input_attribute_nodes:
|
|
self.input_attribute_nodes[key] = InputAttributeNode(
|
|
self, key, "__getitem__", input_type
|
|
)
|
|
return self.input_attribute_nodes[key]
|
|
|
|
def __enter__(self):
|
|
self.set_context(IN_CONTEXT_MANAGER, True)
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
pass
|
|
|
|
def get_result_type(self) -> str:
|
|
"""Get type of the output of this DAGNode.
|
|
|
|
Generated by ray.experimental.gradio_utils.type_to_string().
|
|
"""
|
|
if "result_type_string" in self._bound_other_args_to_resolve:
|
|
return self._bound_other_args_to_resolve["result_type_string"]
|
|
|
|
|
|
@DeveloperAPI
|
|
class InputAttributeNode(DAGNode):
|
|
"""Represents partial access of user input based on an index (int),
|
|
object attribute or dict key (str).
|
|
|
|
Examples:
|
|
|
|
.. code-block:: python
|
|
|
|
with InputNode() as dag_input:
|
|
a = dag_input[0]
|
|
b = dag_input.x
|
|
ray_dag = add.bind(a, b)
|
|
|
|
# This makes a = 1 and b = 2
|
|
ray_dag.execute(1, x=2)
|
|
|
|
with InputNode() as dag_input:
|
|
a = dag_input[0]
|
|
b = dag_input[1]
|
|
ray_dag = add.bind(a, b)
|
|
|
|
# This makes a = 2 and b = 3
|
|
ray_dag.execute(2, 3)
|
|
|
|
# Alternatively, you can input a single object
|
|
# and the inputs are automatically indexed from the object:
|
|
# This makes a = 2 and b = 3
|
|
ray_dag.execute([2, 3])
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dag_input_node: InputNode,
|
|
key: Union[int, str],
|
|
accessor_method: str,
|
|
input_type: str = None,
|
|
):
|
|
"""Initialize an InputAttributeNode.
|
|
|
|
Args:
|
|
dag_input_node: The parent ``InputNode`` this attribute access
|
|
derives from.
|
|
key: The index, attribute name, or dict key used to access a
|
|
value of the user input.
|
|
accessor_method: The accessor method used to extract the value
|
|
from the user input (e.g., ``"__getitem__"`` or
|
|
``"__getattr__"``).
|
|
input_type: Type hint for the extracted value, used by the
|
|
Gradio visualizer to pick a UI component.
|
|
"""
|
|
self._dag_input_node = dag_input_node
|
|
self._key = key
|
|
self._accessor_method = accessor_method
|
|
super().__init__(
|
|
[],
|
|
{},
|
|
{},
|
|
{
|
|
"dag_input_node": dag_input_node,
|
|
"key": key,
|
|
"accessor_method": accessor_method,
|
|
# Type of the input tied to this node. Used by
|
|
# gradio_visualize_graph.GraphVisualizer to determine which Gradio
|
|
# component should be used for this node.
|
|
"result_type_string": input_type,
|
|
},
|
|
)
|
|
|
|
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 InputAttributeNode(
|
|
new_other_args_to_resolve["dag_input_node"],
|
|
new_other_args_to_resolve["key"],
|
|
new_other_args_to_resolve["accessor_method"],
|
|
new_other_args_to_resolve["result_type_string"],
|
|
)
|
|
|
|
def _execute_impl(self, *args, **kwargs):
|
|
"""Executor of InputAttributeNode.
|
|
|
|
Args and kwargs are to match base class signature, but not in the
|
|
implementation. All args and kwargs should be resolved and replaced
|
|
with value in bound_args and bound_kwargs via bottom-up recursion when
|
|
current node is executed.
|
|
"""
|
|
|
|
if isinstance(self._dag_input_node, DAGInputData):
|
|
return self._dag_input_node[self._key]
|
|
else:
|
|
# dag.execute() is called with only one arg, thus when an
|
|
# InputAttributeNode is executed, its dependent InputNode is
|
|
# resolved with original user input python object.
|
|
user_input_python_object = self._dag_input_node
|
|
if isinstance(self._key, str):
|
|
if self._accessor_method == "__getitem__":
|
|
return user_input_python_object[self._key]
|
|
elif self._accessor_method == "__getattr__":
|
|
return getattr(user_input_python_object, self._key)
|
|
elif isinstance(self._key, int):
|
|
return user_input_python_object[self._key]
|
|
else:
|
|
raise ValueError(
|
|
"Please only use int index or str as first-level key to "
|
|
"access fields of dag input."
|
|
)
|
|
|
|
def __str__(self) -> str:
|
|
return get_dag_node_str(self, f'["{self._key}"]')
|
|
|
|
def get_result_type(self) -> str:
|
|
"""Get type of the output of this DAGNode.
|
|
|
|
Generated by ray.experimental.gradio_utils.type_to_string().
|
|
"""
|
|
if "result_type_string" in self._bound_other_args_to_resolve:
|
|
return self._bound_other_args_to_resolve["result_type_string"]
|
|
|
|
@property
|
|
def key(self) -> Union[int, str]:
|
|
return self._key
|
|
|
|
|
|
@DeveloperAPI
|
|
class DAGInputData:
|
|
"""If user passed multiple args and kwargs directly to dag.execute(), we
|
|
generate this wrapper for all user inputs as one object, accessible via
|
|
list index or object attribute key.
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
self._args = list(args)
|
|
self._kwargs = kwargs
|
|
|
|
def __getitem__(self, key: Union[int, str]) -> Any:
|
|
if isinstance(key, int):
|
|
# Access list args by index.
|
|
return self._args[key]
|
|
elif isinstance(key, str):
|
|
# Access kwarg by key.
|
|
return self._kwargs[key]
|
|
else:
|
|
raise ValueError(
|
|
"Please only use int index or str as first-level key to "
|
|
"access fields of dag input."
|
|
)
|