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()
def _node_visitor(node: Any) -> Any:
if isinstance(node, FunctionNode):
bound_options = node._bound_options.copy()
num_returns = bound_options.get("num_returns", 1)
if num_returns is None: # ray could use `None` as default value
num_returns = 1
if num_returns > 1:
raise ValueError("Workflow steps can only have one return.")
workflow_options = bound_options.pop("_metadata",
{}).get(WORKFLOW_OPTIONS, {})
# If checkpoint option is not specified, inherit checkpoint
# options from context (i.e. checkpoint options of the outer
# step). If it is still not specified, it's True by default.
checkpoint = workflow_options.get("checkpoint", None)
if checkpoint is None:
checkpoint = context.checkpoint if context is not None else True
# When it returns a nested workflow, catch_exception
# should be passed recursively.
catch_exceptions = workflow_options.get("catch_exceptions", None)
if catch_exceptions is None:
# TODO(suquark): should we also handle exceptions from a "leaf node"
# in the continuation? For example, we have a workflow
# > @ray.remote
# > def A(): pass
# > @ray.remote
# > def B(x): return x
# > @ray.remote
# > def C(x): return workflow.continuation(B.bind(A.bind()))
# > dag = C.options(**workflow.options(catch_exceptions=True)).bind()
# Should C catches exceptions of A?
if node.get_stable_uuid() == dag_node.get_stable_uuid():
# 'catch_exception' context should be passed down to
# its direct continuation task.
# In this case, the direct continuation is the output node.
catch_exceptions = (context.catch_exceptions
if context is not None else False)
else:
catch_exceptions = False
max_retries = bound_options.get("max_retries", 3)
if not isinstance(max_retries, int) or max_retries < -1:
raise ValueError(
"'max_retries' only accepts 0, -1 or a positive integer.")
step_options = WorkflowStepRuntimeOptions(
step_type=StepType.FUNCTION,
catch_exceptions=catch_exceptions,
max_retries=max_retries,
allow_inplace=False,
checkpoint=checkpoint,
ray_options=bound_options,
)
workflow_refs: List[WorkflowRef] = []
with serialization_context.workflow_args_serialization_context(
workflow_refs):
_func_signature = signature.extract_signature(node._body)
flattened_args = signature.flatten_args(
_func_signature, node._bound_args, node._bound_kwargs)
# NOTE: When calling 'ray.put', we trigger python object
# serialization. Under our serialization context,
# Workflows are separated from the arguments,
# leaving a placeholder object with all other python objects.
# Then we put the placeholder object to object store,
# so it won't be mutated later. This guarantees correct
# semantics. See "tests/test_variable_mutable.py" as
# an example.
input_placeholder: ray.ObjectRef = ray.put(flattened_args)
name = workflow_options.get("name")
if name is None:
name = f"{get_module(node._body)}.{slugify(get_qualname(node._body))}"
task_id = ray.get(mgr.gen_step_id.remote(workflow_id, name))
state.add_dependencies(task_id, [s.task_id for s in workflow_refs])
state.task_input_args[task_id] = input_placeholder
user_metadata = workflow_options.pop("metadata", {})
validate_user_metadata(user_metadata)
state.tasks[task_id] = Task(
name=name,
options=step_options,
user_metadata=user_metadata,
func_body=node._body,
)
return WorkflowRef(task_id)
if isinstance(node, InputAttributeNode):
return node._execute_impl() # get data from input node
if isinstance(node, InputNode):
return input_context # replace input node with input data
if not isinstance(node, DAGNode):
return node # return normal objects
raise TypeError(f"Unsupported DAG node: {node}")
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()
