def _resolve_step_inputs(
step_inputs: StepInputTupleToResolve) -> Tuple[List, Dict]:
"""
This function resolves the inputs for the code inside
a workflow step (works on the callee side). For outputs from other
workflows, we resolve them into object instances inplace.
For each ObjectRef argument, the function returns both the ObjectRef
and the object instance. If the ObjectRef is a chain of nested
ObjectRefs, then we resolve it recursively until we get the
object instance, and we return the *direct* ObjectRef of the
instance. This function does not resolve ObjectRef
inside another object (e.g. list of ObjectRefs) to give users some
flexibility.
Args:
step_inputs: Workflow step inputs.
Returns:
Instances of arguments.
"""
objects_mapping = []
input_placeholder, input_workflows, input_object_refs = step_inputs
for obj_ref in input_workflows:
obj, ref = _resolve_object_ref(obj_ref)
objects_mapping.append(obj)
with serialization_context.workflow_args_resolving_context(
objects_mapping, input_object_refs):
# reconstruct input arguments under correct serialization context
flattened_args: List[Any] = ray.get(input_placeholder)
# dereference arguments like Ray remote functions
flattened_args = [
ray.get(a) if isinstance(a, ObjectRef) else a for a in flattened_args
]
return signature.recover_args(flattened_args)
def load_step_args(
self, step_id: StepID, workflows: List[Any],
object_refs: List[ray.ObjectRef],
workflow_refs: List[WorkflowRef]) -> Tuple[List, Dict[str, Any]]:
"""Load the input arguments of the workflow step. This must be
done under a serialization context, otherwise the arguments would
not be reconstructed successfully.
Args:
step_id: ID of the workflow step.
workflows: The workflows in the original arguments,
replaced by the actual workflow outputs.
object_refs: The object refs in the original arguments.
Returns:
Args and kwargs.
"""
with serialization_context.workflow_args_resolving_context(
workflows, object_refs, workflow_refs):
flattened_args = asyncio_run(
self._get(self._key_step_args(step_id)))
# dereference arguments like Ray remote functions
flattened_args = [
ray.get(a) if isinstance(a, ray.ObjectRef) else a
for a in flattened_args
]
return signature.recover_args(flattened_args)
def _resolve_step_inputs(
step_inputs: Tuple[RRef, List[RRef], List[RRef]]
) -> Tuple[List, Dict, List[RRef]]:
"""
This function resolves the inputs for the code inside
a workflow step (works on the callee side). For outputs from other
workflows, we resolve them into object instances inplace.
For each ObjectRef argument, the function returns both the ObjectRef
and the object instance. If the ObjectRef is a chain of nested
ObjectRefs, then we resolve it recursively until we get the
object instance, and we return the *direct* ObjectRef of the
instance. This function does not resolve ObjectRef
inside another object (e.g. list of ObjectRefs) to give users some
flexibility.
Args:
step_inputs: Workflow step inputs.
Returns:
Instances of arguments and resolved object refs.
"""
resolved_object_refs = []
objects_mapping = []
input_placeholder, input_workflows, input_object_refs = step_inputs
for rref in input_workflows:
obj, ref = resolve_object_ref(rref)
objects_mapping.append(obj)
resolved_object_refs.append(ref)
with serialization_context.workflow_args_resolving_context(
objects_mapping, input_object_refs):
# reconstruct input arguments under correct serialization context
args, kwargs = ray.get(input_placeholder)
_args, _kwargs = _deref_arguments(args, kwargs)
return _args, _kwargs, resolved_object_refs
def load_step_args(
self, step_id: StepID, workflows: List[Any],
object_refs: List[ray.ObjectRef]) -> Tuple[List, Dict[str, Any]]:
"""Load the input arguments of the workflow step. This must be
done under a serialization context, otherwise the arguments would
not be reconstructed successfully.
Args:
step_id: ID of the workflow step.
workflows: The workflows in the original arguments,
replaced by the actual workflow outputs.
object_refs: The object refs in the original arguments.
Returns:
Args and kwargs.
"""
with serialization_context.workflow_args_resolving_context(
workflows, object_refs):
return self._storage.load_step_args(self._workflow_id, step_id)
