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 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()