wehub-resource-sync
156 lines
5.3 KiB
Python
156 lines
5.3 KiB
Python
from ray.dag import DAGNode
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
|
|
@DeveloperAPI
|
|
def get_dag_node_str(
|
|
dag_node: DAGNode,
|
|
body_line,
|
|
):
|
|
indent = _get_indentation()
|
|
other_args_to_resolve_lines = _get_other_args_to_resolve_lines(
|
|
dag_node._bound_other_args_to_resolve
|
|
)
|
|
return (
|
|
f"({dag_node.__class__.__name__}, {dag_node._stable_uuid})(\n"
|
|
f"{indent}body={body_line}\n"
|
|
f"{indent}args={_get_args_lines(dag_node._bound_args)}\n"
|
|
f"{indent}kwargs={_get_kwargs_lines(dag_node._bound_kwargs)}\n"
|
|
f"{indent}options={_get_options_lines(dag_node._bound_options)}\n"
|
|
f"{indent}other_args_to_resolve={other_args_to_resolve_lines}\n"
|
|
f")"
|
|
)
|
|
|
|
|
|
def _get_indentation(num_spaces=4):
|
|
return " " * num_spaces
|
|
|
|
|
|
def _get_args_lines(bound_args):
|
|
"""Pretty prints bounded args of a DAGNode, and recursively handle
|
|
DAGNode in list / dict containers.
|
|
"""
|
|
indent = _get_indentation()
|
|
lines = []
|
|
for arg in bound_args:
|
|
if isinstance(arg, DAGNode):
|
|
node_repr_lines = str(arg).split("\n")
|
|
for node_repr_line in node_repr_lines:
|
|
lines.append(f"{indent}" + node_repr_line)
|
|
elif isinstance(arg, list):
|
|
for ele in arg:
|
|
node_repr_lines = str(ele).split("\n")
|
|
for node_repr_line in node_repr_lines:
|
|
lines.append(f"{indent}" + node_repr_line)
|
|
elif isinstance(arg, dict):
|
|
for _, val in arg.items():
|
|
node_repr_lines = str(val).split("\n")
|
|
for node_repr_line in node_repr_lines:
|
|
lines.append(f"{indent}" + node_repr_line)
|
|
# TODO: (jiaodong) Handle nested containers and other obj types
|
|
else:
|
|
lines.append(f"{indent}" + str(arg) + ", ")
|
|
|
|
if len(lines) == 0:
|
|
args_line = "[]"
|
|
else:
|
|
args_line = "["
|
|
for args in lines:
|
|
args_line += f"\n{indent}{args}"
|
|
args_line += f"\n{indent}]"
|
|
|
|
return args_line
|
|
|
|
|
|
def _get_kwargs_lines(bound_kwargs):
|
|
"""Pretty prints bounded kwargs of a DAGNode, and recursively handle
|
|
DAGNode in list / dict containers.
|
|
"""
|
|
# TODO: (jiaodong) Nits, we're missing keys and indentation was a bit off.
|
|
if not bound_kwargs:
|
|
return "{}"
|
|
indent = _get_indentation()
|
|
kwargs_lines = []
|
|
for key, val in bound_kwargs.items():
|
|
if isinstance(val, DAGNode):
|
|
node_repr_lines = str(val).split("\n")
|
|
for index, node_repr_line in enumerate(node_repr_lines):
|
|
if index == 0:
|
|
kwargs_lines.append(
|
|
f"{indent}{key}:" + f"{indent}" + node_repr_line
|
|
)
|
|
else:
|
|
kwargs_lines.append(f"{indent}{indent}" + node_repr_line)
|
|
|
|
elif isinstance(val, list):
|
|
for ele in val:
|
|
node_repr_lines = str(ele).split("\n")
|
|
for node_repr_line in node_repr_lines:
|
|
kwargs_lines.append(f"{indent}" + node_repr_line)
|
|
elif isinstance(val, dict):
|
|
for _, inner_val in val.items():
|
|
node_repr_lines = str(inner_val).split("\n")
|
|
for node_repr_line in node_repr_lines:
|
|
kwargs_lines.append(f"{indent}" + node_repr_line)
|
|
# TODO: (jiaodong) Handle nested containers and other obj types
|
|
else:
|
|
kwargs_lines.append(val)
|
|
|
|
if len(kwargs_lines) > 0:
|
|
kwargs_line = "{"
|
|
for line in kwargs_lines:
|
|
kwargs_line += f"\n{indent}{line}"
|
|
kwargs_line += f"\n{indent}}}"
|
|
else:
|
|
kwargs_line = "{}"
|
|
|
|
return kwargs_line
|
|
|
|
|
|
def _get_options_lines(bound_options):
|
|
"""Pretty prints .options() in DAGNode. Only prints non-empty values."""
|
|
if not bound_options:
|
|
return "{}"
|
|
indent = _get_indentation()
|
|
options_lines = []
|
|
for key, val in bound_options.items():
|
|
if val:
|
|
options_lines.append(f"{indent}{key}: " + str(val))
|
|
|
|
options_line = "{"
|
|
for line in options_lines:
|
|
options_line += f"\n{indent}{line}"
|
|
options_line += f"\n{indent}}}"
|
|
return options_line
|
|
|
|
|
|
def _get_other_args_to_resolve_lines(other_args_to_resolve):
|
|
if not other_args_to_resolve:
|
|
return "{}"
|
|
indent = _get_indentation()
|
|
other_args_to_resolve_lines = []
|
|
for key, val in other_args_to_resolve.items():
|
|
if isinstance(val, DAGNode):
|
|
node_repr_lines = str(val).split("\n")
|
|
for index, node_repr_line in enumerate(node_repr_lines):
|
|
if index == 0:
|
|
other_args_to_resolve_lines.append(
|
|
f"{indent}{key}:"
|
|
+ f"{indent}"
|
|
+ "\n"
|
|
+ f"{indent}{indent}{indent}"
|
|
+ node_repr_line
|
|
)
|
|
else:
|
|
other_args_to_resolve_lines.append(
|
|
f"{indent}{indent}" + node_repr_line
|
|
)
|
|
else:
|
|
other_args_to_resolve_lines.append(f"{indent}{key}: " + str(val))
|
|
|
|
other_args_to_resolve_line = "{"
|
|
for line in other_args_to_resolve_lines:
|
|
other_args_to_resolve_line += f"\n{indent}{line}"
|
|
other_args_to_resolve_line += f"\n{indent}}}"
|
|
return other_args_to_resolve_line
|