def test_workflow_storage(ray_start_regular, raw_storage):
workflow_id = test_workflow_storage.__name__
step_id = "some_step"
input_metadata = {
"name": "test_basic_workflows.append1",
"object_refs": ["abc"],
"workflows": ["def"]
}
output_metadata = {
"output_step_id": "a12423",
"dynamic_output_step_id": "b1234"
}
args = ([1, "2"], {"k": b"543"})
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, 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)
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(
args_valid=True,
func_body_valid=True,
object_refs=input_metadata["object_refs"],
workflows=input_metadata["workflows"])
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(
func_body_valid=True,
object_refs=input_metadata["object_refs"],
workflows=input_metadata["workflows"])
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(
object_refs=input_metadata["object_refs"],
workflows=input_metadata["workflows"])
assert not inspect_result.is_recoverable()
step_id = "some_step6"
inspect_result = wf_storage.inspect_step(step_id)
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()