Example #1
0
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()
Example #2
0
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()