def test_embedded_objectrefs(workflow_start_regular):
workflow_id = test_embedded_objectrefs.__name__
base_storage = storage.get_global_storage()
class ObjectRefsWrapper:
def __init__(self, refs):
self.refs = refs
url = base_storage.storage_url
wrapped = ObjectRefsWrapper([ray.put(1), ray.put(2)])
promise = serialization.dump_to_storage(["key"], wrapped, workflow_id,
base_storage)
workflow_storage.asyncio_run(promise)
# Be extremely explicit about shutting down. We want to make sure the
# `_get` call deserializes the full object and puts it in the object store.
# Shutting down the cluster should guarantee we don't accidently get the
# old object and pass the test.
ray.shutdown()
subprocess.check_output("ray stop --force", shell=True)
workflow.init(url)
storage2 = workflow_storage.get_workflow_storage(workflow_id)
result = workflow_storage.asyncio_run(storage2._get(["key"]))
assert ray.get(result.refs) == [1, 2]
def resume(num_records_replayed):
key = debug_store.wrapped_storage.make_key("complex_workflow")
asyncio_run(debug_store.wrapped_storage.delete_prefix(key))
async def replay():
# We need to replay one by one to avoid conflict
for i in range(num_records_replayed):
await debug_store.replay(i)
asyncio_run(replay())
return ray.get(workflow.resume(workflow_id="complex_workflow"))
def resume_all(with_failed: bool) -> List[Tuple[str, ray.ObjectRef]]:
filter_set = {WorkflowStatus.RESUMABLE}
if with_failed:
filter_set.add(WorkflowStatus.FAILED)
all_failed = list_all(filter_set)
try:
workflow_manager = get_management_actor()
except Exception as e:
raise RuntimeError("Failed to get management actor") from e
async def _resume_one(wid: str) -> Tuple[str, Optional[ray.ObjectRef]]:
try:
result: "WorkflowExecutionResult" = (
await workflow_manager.run_or_resume.remote(wid)
)
obj = flatten_workflow_output(wid, result.persisted_output)
return wid, obj
except Exception:
logger.error(f"Failed to resume workflow {wid}")
return (wid, None)
ret = workflow_storage.asyncio_run(
asyncio.gather(*[_resume_one(wid) for (wid, _) in all_failed])
)
return [(wid, obj) for (wid, obj) in ret if obj is not None]
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()
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.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
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_metadata(""),
root_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 s3 path
# here we hardcode the path to make sure s3 path is parsed correctly
if isinstance(wf_storage._storage, S3StorageImpl):
assert (asyncio_run(
wf_storage._storage.get(
"workflow/test_workflow_storage/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"
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)
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()
