def load_step_args(
self, step_id: StepID, workflows: List[Any], 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, 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(self, store: workflow_storage.WorkflowStorage) -> 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.
Returns:
Instances of arguments.
"""
workflow_ref_mapping = []
for r in self.workflow_refs:
if r.ref is None:
workflow_ref_mapping.append(store.load_step_output(r.task_id))
else:
workflow_ref_mapping.append(r.ref)
with serialization_context.workflow_args_resolving_context(
workflow_ref_mapping
):
# reconstruct input arguments under correct serialization context
flattened_args: List[Any] = ray.get(self.args)
# 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 resolve(self) -> 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.
Returns:
Instances of arguments.
"""
objects_mapping = []
for static_workflow_ref in self.workflow_outputs:
if static_workflow_ref._resolve_like_object_ref_in_args:
# Keep it unresolved as an ObjectRef. Then we resolve it
# later in the arguments, like how Ray does with ObjectRefs.
obj = ray.put(_SelfDereference(static_workflow_ref))
else:
obj = _resolve_static_workflow_ref(static_workflow_ref)
objects_mapping.append(obj)
workflow_ref_mapping = _resolve_dynamic_workflow_refs(
self.job_id, self.workflow_refs)
with serialization_context.workflow_args_resolving_context(
objects_mapping, workflow_ref_mapping):
# reconstruct input arguments under correct serialization context
flattened_args: List[Any] = ray.get(self.args)
# 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 _resolve_step_inputs(
step_inputs: "_BakedWorkflowInputs") -> 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 = []
for obj_ref in step_inputs.workflow_outputs:
obj, ref = _resolve_object_ref(obj_ref)
objects_mapping.append(obj)
workflow_ref_mapping = _resolve_dynamic_workflow_refs(
step_inputs.workflow_refs)
with serialization_context.workflow_args_resolving_context(
objects_mapping, workflow_ref_mapping):
# reconstruct input arguments under correct serialization context
flattened_args: List[Any] = ray.get(step_inputs.args)
# 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 test_workflow_storage(workflow_start_regular):
workflow_id = test_workflow_storage.__name__
wf_storage = workflow_storage.WorkflowStorage(workflow_id)
task_id = "some_step"
step_options = WorkflowStepRuntimeOptions(
step_type=StepType.FUNCTION,
catch_exceptions=False,
max_retries=0,
allow_inplace=False,
checkpoint=False,
ray_options={},
)
input_metadata = {
"name": "test_basic_workflows.append1",
"workflow_refs": ["some_ref"],
"step_options": step_options.to_dict(),
}
output_metadata = {
"output_step_id": "a12423",
"dynamic_output_step_id": "b1234"
}
root_output_metadata = {"output_step_id": "c123"}
flattened_args = [
signature.DUMMY_TYPE, 1, signature.DUMMY_TYPE, "2", "k", b"543"
]
args = signature.recover_args(flattened_args)
output = ["the_answer"]
object_resolved = 42
obj_ref = ray.put(object_resolved)
# test basics
wf_storage._put(wf_storage._key_step_input_metadata(task_id),
input_metadata, True)
wf_storage._put(wf_storage._key_step_function_body(task_id), some_func)
wf_storage._put(wf_storage._key_step_args(task_id), flattened_args)
wf_storage._put(wf_storage._key_obj_id(obj_ref.hex()), ray.get(obj_ref))
wf_storage._put(wf_storage._key_step_output_metadata(task_id),
output_metadata, True)
wf_storage._put(wf_storage._key_step_output_metadata(""),
root_output_metadata, True)
wf_storage._put(wf_storage._key_step_output(task_id), output)
assert wf_storage.load_step_output(task_id) == output
with serialization_context.workflow_args_resolving_context([]):
assert (signature.recover_args(
ray.get(wf_storage.load_step_args(task_id))) == args)
assert wf_storage.load_step_func_body(task_id)(33) == 34
assert ray.get(wf_storage.load_object_ref(
obj_ref.hex())) == object_resolved
# test s3 path
# here we hardcode the path to make sure s3 path is parsed correctly
from ray._private.storage import _storage_uri
if _storage_uri.startswith("s3://"):
assert wf_storage._get("steps/outputs.json",
True) == root_output_metadata
# test "inspect_step"
inspect_result = wf_storage.inspect_step(task_id)
assert inspect_result == workflow_storage.StepInspectResult(
output_object_valid=True)
assert inspect_result.is_recoverable()
task_id = "some_step2"
wf_storage._put(wf_storage._key_step_input_metadata(task_id),
input_metadata, True)
wf_storage._put(wf_storage._key_step_function_body(task_id), some_func)
wf_storage._put(wf_storage._key_step_args(task_id), args)
wf_storage._put(wf_storage._key_step_output_metadata(task_id),
output_metadata, True)
inspect_result = wf_storage.inspect_step(task_id)
assert inspect_result == workflow_storage.StepInspectResult(
output_step_id=output_metadata["dynamic_output_step_id"])
assert inspect_result.is_recoverable()
task_id = "some_step3"
wf_storage._put(wf_storage._key_step_input_metadata(task_id),
input_metadata, True)
wf_storage._put(wf_storage._key_step_function_body(task_id), some_func)
wf_storage._put(wf_storage._key_step_args(task_id), args)
inspect_result = wf_storage.inspect_step(task_id)
assert inspect_result == workflow_storage.StepInspectResult(
args_valid=True,
func_body_valid=True,
workflow_refs=input_metadata["workflow_refs"],
step_options=step_options,
)
assert inspect_result.is_recoverable()
task_id = "some_step4"
wf_storage._put(wf_storage._key_step_input_metadata(task_id),
input_metadata, True)
wf_storage._put(wf_storage._key_step_function_body(task_id), some_func)
inspect_result = wf_storage.inspect_step(task_id)
assert inspect_result == workflow_storage.StepInspectResult(
func_body_valid=True,
workflow_refs=input_metadata["workflow_refs"],
step_options=step_options,
)
assert not inspect_result.is_recoverable()
task_id = "some_step5"
wf_storage._put(wf_storage._key_step_input_metadata(task_id),
input_metadata, True)
inspect_result = wf_storage.inspect_step(task_id)
assert inspect_result == workflow_storage.StepInspectResult(
workflow_refs=input_metadata["workflow_refs"],
step_options=step_options,
)
assert not inspect_result.is_recoverable()
task_id = "some_step6"
inspect_result = wf_storage.inspect_step(task_id)
print(inspect_result)
assert inspect_result == workflow_storage.StepInspectResult()
assert not inspect_result.is_recoverable()
