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], 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 _readonly_method_executor(actor_id: str, storage: Storage, cls: type,
method_name: str, flattened_args: List):
instance = cls.__new__(cls)
try:
state = get_latest_output(actor_id, storage)
except Exception as e:
raise VirtualActorNotInitializedError(
f"Virtual actor '{actor_id}' has not been initialized. "
"We cannot get the latest state for the "
"readonly virtual actor.") from e
instance.__setstate__(state)
args, kwargs = signature.recover_args(flattened_args)
method = getattr(instance, method_name)
return method(*args, **kwargs)
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 _batch_args_kwargs(
list_of_flattened_args: List[List[Any]],
) -> Tuple[Tuple[Any], Dict[Any, Any]]:
"""Batch a list of flatten args and returns regular args and kwargs"""
# Ray's flatten arg format is a list with alternating key and values
# e.g. args=(1, 2), kwargs={"key": "val"} got turned into
# [None, 1, None, 2, "key", "val"]
arg_lengths = {len(args) for args in list_of_flattened_args}
assert (len(arg_lengths) == 1
), "All batch requests should have the same number of parameters."
arg_length = arg_lengths.pop()
batched_flattened_args = []
for idx in range(arg_length):
if idx % 2 == 0:
batched_flattened_args.append(list_of_flattened_args[0][idx])
else:
batched_flattened_args.append(
[item[idx] for item in list_of_flattened_args])
return recover_args(batched_flattened_args)
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):
flattened_args = asyncio_run(
self._storage.load_step_args(self._workflow_id, step_id))
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)
step_id = "some_step"
step_options = WorkflowStepRuntimeOptions.make(step_type=StepType.FUNCTION)
input_metadata = {
"name": "test_basic_workflows.append1",
"workflows": ["def"],
"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(step_id),
input_metadata, True)
wf_storage._put(wf_storage._key_step_function_body(step_id), some_func)
wf_storage._put(wf_storage._key_step_args(step_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(step_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(step_id), output)
assert wf_storage.load_step_output(step_id) == output
assert wf_storage.load_step_args(step_id, [], []) == args
assert wf_storage.load_step_func_body(step_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(step_id)
assert inspect_result == workflow_storage.StepInspectResult(
output_object_valid=True)
assert inspect_result.is_recoverable()
step_id = "some_step2"
wf_storage._put(wf_storage._key_step_input_metadata(step_id),
input_metadata, True)
wf_storage._put(wf_storage._key_step_function_body(step_id), some_func)
wf_storage._put(wf_storage._key_step_args(step_id), args)
wf_storage._put(wf_storage._key_step_output_metadata(step_id),
output_metadata, True)
inspect_result = wf_storage.inspect_step(step_id)
assert inspect_result == workflow_storage.StepInspectResult(
output_step_id=output_metadata["dynamic_output_step_id"])
assert inspect_result.is_recoverable()
step_id = "some_step3"
wf_storage._put(wf_storage._key_step_input_metadata(step_id),
input_metadata, True)
wf_storage._put(wf_storage._key_step_function_body(step_id), some_func)
wf_storage._put(wf_storage._key_step_args(step_id), args)
inspect_result = wf_storage.inspect_step(step_id)
assert inspect_result == workflow_storage.StepInspectResult(
args_valid=True,
func_body_valid=True,
workflows=input_metadata["workflows"],
workflow_refs=input_metadata["workflow_refs"],
step_options=step_options,
)
assert inspect_result.is_recoverable()
step_id = "some_step4"
wf_storage._put(wf_storage._key_step_input_metadata(step_id),
input_metadata, True)
wf_storage._put(wf_storage._key_step_function_body(step_id), some_func)
inspect_result = wf_storage.inspect_step(step_id)
assert inspect_result == workflow_storage.StepInspectResult(
func_body_valid=True,
workflows=input_metadata["workflows"],
workflow_refs=input_metadata["workflow_refs"],
step_options=step_options,
)
assert not inspect_result.is_recoverable()
step_id = "some_step5"
wf_storage._put(wf_storage._key_step_input_metadata(step_id),
input_metadata, True)
inspect_result = wf_storage.inspect_step(step_id)
assert inspect_result == workflow_storage.StepInspectResult(
workflows=input_metadata["workflows"],
workflow_refs=input_metadata["workflow_refs"],
step_options=step_options,
)
assert not inspect_result.is_recoverable()
step_id = "some_step6"
inspect_result = wf_storage.inspect_step(step_id)
print(inspect_result)
assert inspect_result == workflow_storage.StepInspectResult()
assert not inspect_result.is_recoverable()
def test_workflow_storage(workflow_start_regular):
raw_storage = workflow_storage._StorageImpl(storage.get_global_storage())
workflow_id = test_workflow_storage.__name__
step_id = "some_step"
input_metadata = {
"name": "test_basic_workflows.append1",
"step_type": StepType.FUNCTION,
"object_refs": ["abc"],
"workflows": ["def"],
"workflow_refs": ["some_ref"],
"max_retries": 1,
"catch_exceptions": False,
"ray_options": {},
}
output_metadata = {
"output_step_id": "a12423",
"dynamic_output_step_id": "b1234"
}
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
asyncio_run(
raw_storage.save_step_input_metadata(workflow_id, step_id,
input_metadata))
asyncio_run(
raw_storage.save_step_func_body(workflow_id, step_id, some_func))
asyncio_run(
raw_storage.save_step_args(workflow_id, step_id, flattened_args))
asyncio_run(raw_storage.save_object_ref(workflow_id, obj_ref))
asyncio_run(
raw_storage.save_step_output_metadata(workflow_id, step_id,
output_metadata))
asyncio_run(raw_storage.save_step_output(workflow_id, step_id, output))
wf_storage = workflow_storage.WorkflowStorage(workflow_id,
storage.get_global_storage())
assert wf_storage.load_step_output(step_id) == output
assert wf_storage.load_step_args(step_id, [], [], []) == args
assert wf_storage.load_step_func_body(step_id)(33) == 34
assert ray.get(wf_storage.load_object_ref(
obj_ref.hex())) == object_resolved
# test "inspect_step"
inspect_result = wf_storage.inspect_step(step_id)
assert inspect_result == workflow_storage.StepInspectResult(
output_object_valid=True)
assert inspect_result.is_recoverable()
step_id = "some_step2"
asyncio_run(
raw_storage.save_step_input_metadata(workflow_id, step_id,
input_metadata))
asyncio_run(
raw_storage.save_step_func_body(workflow_id, step_id, some_func))
asyncio_run(raw_storage.save_step_args(workflow_id, step_id, args))
asyncio_run(
raw_storage.save_step_output_metadata(workflow_id, step_id,
output_metadata))
inspect_result = wf_storage.inspect_step(step_id)
assert inspect_result == workflow_storage.StepInspectResult(
output_step_id=output_metadata["dynamic_output_step_id"])
assert inspect_result.is_recoverable()
step_id = "some_step3"
asyncio_run(
raw_storage.save_step_input_metadata(workflow_id, step_id,
input_metadata))
asyncio_run(
raw_storage.save_step_func_body(workflow_id, step_id, some_func))
asyncio_run(raw_storage.save_step_args(workflow_id, step_id, args))
inspect_result = wf_storage.inspect_step(step_id)
assert inspect_result == workflow_storage.StepInspectResult(
step_type=StepType.FUNCTION,
args_valid=True,
func_body_valid=True,
object_refs=input_metadata["object_refs"],
workflows=input_metadata["workflows"],
workflow_refs=input_metadata["workflow_refs"],
ray_options={})
assert inspect_result.is_recoverable()
step_id = "some_step4"
asyncio_run(
raw_storage.save_step_input_metadata(workflow_id, step_id,
input_metadata))
asyncio_run(
raw_storage.save_step_func_body(workflow_id, step_id, some_func))
inspect_result = wf_storage.inspect_step(step_id)
assert inspect_result == workflow_storage.StepInspectResult(
step_type=StepType.FUNCTION,
func_body_valid=True,
object_refs=input_metadata["object_refs"],
workflows=input_metadata["workflows"],
workflow_refs=input_metadata["workflow_refs"],
ray_options={})
assert not inspect_result.is_recoverable()
step_id = "some_step5"
asyncio_run(
raw_storage.save_step_input_metadata(workflow_id, step_id,
input_metadata))
inspect_result = wf_storage.inspect_step(step_id)
assert inspect_result == workflow_storage.StepInspectResult(
step_type=StepType.FUNCTION,
object_refs=input_metadata["object_refs"],
workflows=input_metadata["workflows"],
workflow_refs=input_metadata["workflow_refs"],
ray_options={})
assert not inspect_result.is_recoverable()
step_id = "some_step6"
inspect_result = wf_storage.inspect_step(step_id)
print(inspect_result)
assert inspect_result == workflow_storage.StepInspectResult()
assert not inspect_result.is_recoverable()
def _virtual_actor_init(cls: type, flattened_args: List) -> Any:
instance = cls.__new__(cls)
args, kwargs = signature.recover_args(flattened_args)
instance.__init__(*args, **kwargs)
return instance.__getstate__()
def test_workflow_storage(workflow_start_regular):
workflow_id = test_workflow_storage.__name__
wf_storage = workflow_storage.WorkflowStorage(workflow_id,
storage.get_global_storage())
step_id = "some_step"
step_options = WorkflowStepRuntimeOptions(
step_type=StepType.FUNCTION,
catch_exceptions=False,
max_retries=1,
ray_options={})
input_metadata = {
"name": "test_basic_workflows.append1",
"workflows": ["def"],
"workflow_refs": ["some_ref"],
"step_options": step_options.to_dict(),
}
output_metadata = {
"output_step_id": "a12423",
"dynamic_output_step_id": "b1234"
}
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
asyncio_run(
wf_storage._put(
wf_storage._key_step_input_metadata(step_id), input_metadata,
True))
asyncio_run(
wf_storage._put(
wf_storage._key_step_function_body(step_id), some_func))
asyncio_run(
wf_storage._put(wf_storage._key_step_args(step_id), flattened_args))
asyncio_run(
wf_storage._put(
wf_storage._key_obj_id(obj_ref.hex()), ray.get(obj_ref)))
asyncio_run(
wf_storage._put(
wf_storage._key_step_output_metadata(step_id), output_metadata,
True))
asyncio_run(wf_storage._put(wf_storage._key_step_output(step_id), output))
assert wf_storage.load_step_output(step_id) == output
assert wf_storage.load_step_args(step_id, [], []) == args
assert wf_storage.load_step_func_body(step_id)(33) == 34
assert ray.get(wf_storage.load_object_ref(
obj_ref.hex())) == object_resolved
# test "inspect_step"
inspect_result = wf_storage.inspect_step(step_id)
assert inspect_result == workflow_storage.StepInspectResult(
output_object_valid=True)
assert inspect_result.is_recoverable()
step_id = "some_step2"
asyncio_run(
wf_storage._put(
wf_storage._key_step_input_metadata(step_id), input_metadata,
True))
asyncio_run(
wf_storage._put(
wf_storage._key_step_function_body(step_id), some_func))
asyncio_run(wf_storage._put(wf_storage._key_step_args(step_id), args))
asyncio_run(
wf_storage._put(
wf_storage._key_step_output_metadata(step_id), output_metadata,
True))
inspect_result = wf_storage.inspect_step(step_id)
assert inspect_result == workflow_storage.StepInspectResult(
output_step_id=output_metadata["dynamic_output_step_id"])
assert inspect_result.is_recoverable()
step_id = "some_step3"
asyncio_run(
wf_storage._put(
wf_storage._key_step_input_metadata(step_id), input_metadata,
True))
asyncio_run(
wf_storage._put(
wf_storage._key_step_function_body(step_id), some_func))
asyncio_run(wf_storage._put(wf_storage._key_step_args(step_id), args))
inspect_result = wf_storage.inspect_step(step_id)
step_options = WorkflowStepRuntimeOptions(
step_type=StepType.FUNCTION,
catch_exceptions=False,
max_retries=1,
ray_options={})
assert inspect_result == workflow_storage.StepInspectResult(
args_valid=True,
func_body_valid=True,
workflows=input_metadata["workflows"],
workflow_refs=input_metadata["workflow_refs"],
step_options=step_options)
assert inspect_result.is_recoverable()
step_id = "some_step4"
asyncio_run(
wf_storage._put(
wf_storage._key_step_input_metadata(step_id), input_metadata,
True))
asyncio_run(
wf_storage._put(
wf_storage._key_step_function_body(step_id), some_func))
inspect_result = wf_storage.inspect_step(step_id)
assert inspect_result == workflow_storage.StepInspectResult(
func_body_valid=True,
workflows=input_metadata["workflows"],
workflow_refs=input_metadata["workflow_refs"],
step_options=step_options)
assert not inspect_result.is_recoverable()
step_id = "some_step5"
asyncio_run(
wf_storage._put(
wf_storage._key_step_input_metadata(step_id), input_metadata,
True))
inspect_result = wf_storage.inspect_step(step_id)
assert inspect_result == workflow_storage.StepInspectResult(
workflows=input_metadata["workflows"],
workflow_refs=input_metadata["workflow_refs"],
step_options=step_options)
assert not inspect_result.is_recoverable()
step_id = "some_step6"
inspect_result = wf_storage.inspect_step(step_id)
print(inspect_result)
assert inspect_result == workflow_storage.StepInspectResult()
assert not inspect_result.is_recoverable()
