def test_imperative_call_from_normal():
@task
def t1(a: str) -> str:
return a + " world"
wb = ImperativeWorkflow(name="my.workflow")
wb.add_workflow_input("in1", str)
node = wb.add_entity(t1, a=wb.inputs["in1"])
wb.add_workflow_output("from_n0t1", node.outputs["o0"])
assert wb(in1="hello") == "hello world"
@workflow
def my_functional_wf(a: str) -> str:
x = wb(in1=a)
return x
assert my_functional_wf(a="hello") == "hello world"
# Create launch plan from wf
lp = LaunchPlan.create("test_wb_2", wb, fixed_inputs={"in1": "hello"})
@workflow
def my_functional_wf_lp() -> str:
x = lp()
return x
assert my_functional_wf_lp() == "hello world"
def test_imperative():
@task
def t1(a: str) -> str:
return a + " world"
@task
def t2():
print("side effect")
wb = ImperativeWorkflow(name="my.workflow")
wb.add_workflow_input("in1", str)
node = wb.add_entity(t1, a=wb.inputs["in1"])
wb.add_entity(t2)
wb.add_workflow_output("from_n0t1", node.outputs["o0"])
assert wb(in1="hello") == "hello world"
srz_wf = get_serializable(OrderedDict(), serialization_settings, wb)
assert len(srz_wf.nodes) == 2
assert srz_wf.nodes[0].task_node is not None
assert len(srz_wf.outputs) == 1
assert srz_wf.outputs[0].var == "from_n0t1"
assert len(srz_wf.interface.inputs) == 1
assert len(srz_wf.interface.outputs) == 1
# Create launch plan from wf, that can also be serialized.
lp = LaunchPlan.create("test_wb", wb)
srz_lp = get_serializable(OrderedDict(), serialization_settings, lp)
assert srz_lp.workflow_id.name == "my.workflow"
def test_basics():
@task
def t1(a: int) -> typing.NamedTuple("OutputsBC", t1_int_output=int, c=str):
return a + 2, "world"
@task
def t2(a: str, b: str) -> str:
return b + a
@workflow
def my_wf(a: int, b: str) -> (int, str):
x, y = t1(a=a)
d = t2(a=y, b=b)
return x, d
wf_spec = get_serializable(OrderedDict(), serialization_settings, my_wf,
False)
assert len(wf_spec.template.interface.inputs) == 2
assert len(wf_spec.template.interface.outputs) == 2
assert len(wf_spec.template.nodes) == 2
assert wf_spec.template.id.resource_type == identifier_models.ResourceType.WORKFLOW
# Gets cached the first time around so it's not actually fast.
task_spec = get_serializable(OrderedDict(), serialization_settings, t1,
True)
assert "pyflyte-execute" in task_spec.template.container.args
lp = LaunchPlan.create(
"testlp",
my_wf,
)
lp_model = get_serializable(OrderedDict(), serialization_settings, lp)
assert lp_model.id.name == "testlp"
def test_call_normal():
@task
def t1(a: int) -> (int, str):
return a + 2, "world"
@workflow
def my_functional_wf(a: int) -> (int, str):
return t1(a=a)
my_functional_lp = LaunchPlan.create("my_functional_wf.lp0",
my_functional_wf,
default_inputs={"a": 3})
wb = ImperativeWorkflow(name="imperio")
node = wb.add_entity(my_functional_wf, a=3)
wb.add_workflow_output("from_n0_1", node.outputs["o0"])
wb.add_workflow_output("from_n0_2", node.outputs["o1"])
assert wb() == (5, "world")
wb_lp = ImperativeWorkflow(name="imperio")
node = wb_lp.add_entity(my_functional_lp)
wb_lp.add_workflow_output("from_n0_1", node.outputs["o0"])
wb_lp.add_workflow_output("from_n0_2", node.outputs["o1"])
assert wb_lp() == (5, "world")
def test_launch_plan_with_fixed_input():
@task
def greet(day_of_week: str, number: int, am: bool) -> str:
greeting = "Have a great " + day_of_week + " "
greeting += "morning" if am else "evening"
return greeting + "!" * number
@workflow
def go_greet(day_of_week: str, number: int, am: bool = False) -> str:
return greet(day_of_week=day_of_week, number=number, am=am)
morning_greeting = LaunchPlan.create(
"morning_greeting",
go_greet,
fixed_inputs={"am": True},
default_inputs={"number": 1},
)
@workflow
def morning_greeter_caller(day_of_week: str) -> str:
greeting = morning_greeting(day_of_week=day_of_week)
return greeting
settings = (
serialization_settings.new_builder().with_fast_serialization_settings(
FastSerializationSettings(enabled=True)).build())
task_spec = get_serializable(OrderedDict(), settings,
morning_greeter_caller)
assert len(task_spec.template.interface.inputs) == 1
assert len(task_spec.template.interface.outputs) == 1
assert len(task_spec.template.nodes) == 1
assert len(task_spec.template.nodes[0].inputs) == 2
def test_basics():
@task
def t1(a: int) -> typing.NamedTuple("OutputsBC", t1_int_output=int, c=str):
return a + 2, "world"
@task
def t2(a: str, b: str) -> str:
return b + a
@workflow
def my_wf(a: int, b: str) -> (int, str):
x, y = t1(a=a)
d = t2(a=y, b=b)
return x, d
sdk_wf = get_serializable(OrderedDict(), serialization_settings, my_wf, False)
assert len(sdk_wf.interface.inputs) == 2
assert len(sdk_wf.interface.outputs) == 2
assert len(sdk_wf.nodes) == 2
# Gets cached the first time around so it's not actually fast.
sdk_task = get_serializable(OrderedDict(), serialization_settings, t1, True)
assert "pyflyte-execute" in sdk_task.container.args
lp = LaunchPlan.create(
"testlp",
my_wf,
)
sdk_lp = get_serializable(OrderedDict(), serialization_settings, lp)
assert sdk_lp.id.name == "testlp"
def test_imperative():
@task
def t1(a: str) -> str:
return a + " world"
@task
def t2():
print("side effect")
wb = ImperativeWorkflow(name="my.workflow")
wb.add_workflow_input("in1", str)
node = wb.add_entity(t1, a=wb.inputs["in1"])
wb.add_entity(t2)
wb.add_workflow_output("from_n0t1", node.outputs["o0"])
assert wb(in1="hello") == "hello world"
wf_spec = get_serializable(OrderedDict(), serialization_settings, wb)
assert len(wf_spec.template.nodes) == 2
assert wf_spec.template.nodes[0].task_node is not None
assert len(wf_spec.template.outputs) == 1
assert wf_spec.template.outputs[0].var == "from_n0t1"
assert len(wf_spec.template.interface.inputs) == 1
assert len(wf_spec.template.interface.outputs) == 1
# Create launch plan from wf, that can also be serialized.
lp = LaunchPlan.create("test_wb", wb)
lp_model = get_serializable(OrderedDict(), serialization_settings, lp)
assert lp_model.spec.workflow_id.name == "my.workflow"
wb2 = ImperativeWorkflow(name="parent.imperative")
p_in1 = wb2.add_workflow_input("p_in1", str)
p_node0 = wb2.add_subwf(wb, in1=p_in1)
wb2.add_workflow_output("parent_wf_output", p_node0.from_n0t1, str)
wb2_spec = get_serializable(OrderedDict(), serialization_settings, wb2)
assert len(wb2_spec.template.nodes) == 1
assert len(wb2_spec.template.interface.inputs) == 1
assert wb2_spec.template.interface.inputs["p_in1"].type.simple is not None
assert len(wb2_spec.template.interface.outputs) == 1
assert wb2_spec.template.interface.outputs[
"parent_wf_output"].type.simple is not None
assert wb2_spec.template.nodes[
0].workflow_node.sub_workflow_ref.name == "my.workflow"
assert len(wb2_spec.sub_workflows) == 1
wb3 = ImperativeWorkflow(name="parent.imperative")
p_in1 = wb3.add_workflow_input("p_in1", str)
p_node0 = wb3.add_launch_plan(lp, in1=p_in1)
wb3.add_workflow_output("parent_wf_output", p_node0.from_n0t1, str)
wb3_spec = get_serializable(OrderedDict(), serialization_settings, wb3)
assert len(wb3_spec.template.nodes) == 1
assert len(wb3_spec.template.interface.inputs) == 1
assert wb3_spec.template.interface.inputs["p_in1"].type.simple is not None
assert len(wb3_spec.template.interface.outputs) == 1
assert wb3_spec.template.interface.outputs[
"parent_wf_output"].type.simple is not None
assert wb3_spec.template.nodes[
0].workflow_node.launchplan_ref.name == "test_wb"
def test_lp_from_ref_wf():
@reference_workflow(project="project", domain="domain", name="name", version="version")
def ref_wf1(p1: str, p2: str) -> None:
...
lp = LaunchPlan.create("reference-wf-12345", ref_wf1, fixed_inputs={"p1": "p1-value", "p2": "p2-value"})
assert lp.name == "reference-wf-12345"
assert lp.workflow == ref_wf1
assert lp.workflow.id.name == "name"
assert lp.workflow.id.project == "project"
assert lp.workflow.id.domain == "domain"
assert lp.workflow.id.version == "version"
def test_file_handling_remote_default_wf_input():
SAMPLE_DATA = "https://raw.githubusercontent.com/jbrownlee/Datasets/master/pima-indians-diabetes.data.csv"
@task
def t1(fname: os.PathLike) -> int:
with open(fname, "r") as fh:
x = len(fh.readlines())
return x
@workflow
def my_wf(fname: os.PathLike = SAMPLE_DATA) -> int:
length = t1(fname=fname)
return length
assert my_wf.python_interface.inputs_with_defaults["fname"][1] == SAMPLE_DATA
sample_lp = LaunchPlan.create("test_launch_plan", my_wf)
assert sample_lp.parameters.parameters["fname"].default.scalar.blob.uri == SAMPLE_DATA
def test_schedule_with_lp():
@task
def double(a: int) -> int:
return a * 2
@workflow
def quadruple(a: int) -> int:
b = double(a=a)
c = double(a=b)
return c
lp = LaunchPlan.create(
"schedule_test",
quadruple,
schedule=FixedRate(_datetime.timedelta(hours=12), "kickoff_input"),
)
assert lp.schedule == _schedule_models.Schedule(
"kickoff_input",
rate=_schedule_models.Schedule.FixedRate(
12, _schedule_models.Schedule.FixedRateUnit.HOUR))
def test_with_launch_plan():
@task
def double(a: int) -> int:
return a * 2
@workflow
def quadruple(a: int) -> int:
b = double(a=a)
c = double(a=b)
return c
lp = LaunchPlan.create(
"notif_test",
quadruple,
notifications=[
notification.Email(phases=[_workflow_execution_succeeded],
recipients_email=["*****@*****.**"])
],
)
assert lp.notifications == [
notification.Email(phases=[_workflow_execution_succeeded],
recipients_email=["*****@*****.**"])
]
def test_imperative():
# Re import with alias
from flytekit.core.workflow import ImperativeWorkflow as Workflow # noqa
# docs_tasks_start
@task
def t1(a: str) -> str:
return a + " world"
@task
def t2():
print("side effect")
# docs_tasks_end
# docs_start
# Create the workflow with a name. This needs to be unique within the project and takes the place of the function
# name that's used for regular decorated function-based workflows.
wb = Workflow(name="my_workflow")
# Adds a top level input to the workflow. This is like an input to a workflow function.
wb.add_workflow_input("in1", str)
# Call your tasks.
node = wb.add_entity(t1, a=wb.inputs["in1"])
wb.add_entity(t2)
# This is analogous to a return statement
wb.add_workflow_output("from_n0t1", node.outputs["o0"])
# docs_end
assert wb(in1="hello") == "hello world"
wf_spec = get_serializable(OrderedDict(), serialization_settings, wb)
assert len(wf_spec.template.nodes) == 2
assert wf_spec.template.nodes[0].task_node is not None
assert len(wf_spec.template.outputs) == 1
assert wf_spec.template.outputs[0].var == "from_n0t1"
assert len(wf_spec.template.interface.inputs) == 1
assert len(wf_spec.template.interface.outputs) == 1
# docs_equivalent_start
nt = typing.NamedTuple("wf_output", from_n0t1=str)
@workflow
def my_workflow(in1: str) -> nt:
x = t1(a=in1)
t2()
return nt(
x,
)
# docs_equivalent_end
# Create launch plan from wf, that can also be serialized.
lp = LaunchPlan.create("test_wb", wb)
lp_model = get_serializable(OrderedDict(), serialization_settings, lp)
assert lp_model.spec.workflow_id.name == "my_workflow"
wb2 = ImperativeWorkflow(name="parent.imperative")
p_in1 = wb2.add_workflow_input("p_in1", str)
p_node0 = wb2.add_subwf(wb, in1=p_in1)
wb2.add_workflow_output("parent_wf_output", p_node0.from_n0t1, str)
wb2_spec = get_serializable(OrderedDict(), serialization_settings, wb2)
assert len(wb2_spec.template.nodes) == 1
assert len(wb2_spec.template.interface.inputs) == 1
assert wb2_spec.template.interface.inputs["p_in1"].type.simple is not None
assert len(wb2_spec.template.interface.outputs) == 1
assert wb2_spec.template.interface.outputs["parent_wf_output"].type.simple is not None
assert wb2_spec.template.nodes[0].workflow_node.sub_workflow_ref.name == "my_workflow"
assert len(wb2_spec.sub_workflows) == 1
wb3 = ImperativeWorkflow(name="parent.imperative")
p_in1 = wb3.add_workflow_input("p_in1", str)
p_node0 = wb3.add_launch_plan(lp, in1=p_in1)
wb3.add_workflow_output("parent_wf_output", p_node0.from_n0t1, str)
wb3_spec = get_serializable(OrderedDict(), serialization_settings, wb3)
assert len(wb3_spec.template.nodes) == 1
assert len(wb3_spec.template.interface.inputs) == 1
assert wb3_spec.template.interface.inputs["p_in1"].type.simple is not None
assert len(wb3_spec.template.interface.outputs) == 1
assert wb3_spec.template.interface.outputs["parent_wf_output"].type.simple is not None
assert wb3_spec.template.nodes[0].workflow_node.launchplan_ref.name == "test_wb"