def compile_into_workflow(
self, ctx: FlyteContext, task_function: Callable, **kwargs
) -> Union[_dynamic_job.DynamicJobSpec, _literal_models.LiteralMap]:
with ctx.new_compilation_context(prefix="dynamic"):
# TODO: Resolve circular import
from flytekit.common.translator import get_serializable
workflow_metadata = WorkflowMetadata(
on_failure=WorkflowFailurePolicy.FAIL_IMMEDIATELY)
defaults = WorkflowMetadataDefaults(interruptible=False)
self._wf = Workflow(task_function,
metadata=workflow_metadata,
default_metadata=defaults)
self._wf.compile(**kwargs)
wf = self._wf
sdk_workflow = get_serializable(ctx.serialization_settings, wf)
# If no nodes were produced, let's just return the strict outputs
if len(sdk_workflow.nodes) == 0:
return _literal_models.LiteralMap(
literals={
binding.var: binding.binding.to_literal_model()
for binding in sdk_workflow._outputs
})
# Gather underlying tasks/workflows that get referenced. Launch plans are handled by propeller.
tasks = set()
sub_workflows = set()
for n in sdk_workflow.nodes:
self.aggregate(tasks, sub_workflows, n)
dj_spec = _dynamic_job.DynamicJobSpec(
min_successes=len(sdk_workflow.nodes),
tasks=list(tasks),
nodes=sdk_workflow.nodes,
outputs=sdk_workflow._outputs,
subworkflows=list(sub_workflows),
)
return dj_spec
class PythonFunctionTask(PythonAutoContainerTask[T]):
"""
A Python Function task should be used as the base for all extensions that have a python function. It will
automatically detect interface of the python function and also, create the write execution command to execute the
function
It is advised this task is used using the @task decorator as follows
.. code-block: python
@task
def my_func(a: int) -> str:
...
In the above code, the name of the function, the module, and the interface (inputs = int and outputs = str) will be
auto detected.
"""
class ExecutionBehavior(Enum):
DEFAULT = 1
DYNAMIC = 2
def __init__(
self,
task_config: T,
task_function: Callable,
task_type="python-task",
ignore_input_vars: Optional[List[str]] = None,
execution_mode: Optional[ExecutionBehavior] = ExecutionBehavior.
DEFAULT,
**kwargs,
):
"""
:param task_config: Configuration object for Task. Should be a unique type for that specific Task
:param task_function: Python function that has type annotations and works for the task
:param ignore_input_vars: When supplied, these input variables will be removed from the interface. This
can be used to inject some client side variables only. Prefer using ExecutionParams
:param task_type: String task type to be associated with this Task
"""
if task_function is None:
raise ValueError(
"TaskFunction is a required parameter for PythonFunctionTask")
if not istestfunction(func=task_function) and isnested(
func=task_function):
raise ValueError(
"TaskFunction cannot be a nested/inner or local function. "
"It should be accessible at a module level for Flyte to execute it. Test modules with "
"names begining with `test_` are allowed to have nested tasks")
self._native_interface = transform_signature_to_interface(
inspect.signature(task_function))
mutated_interface = self._native_interface.remove_inputs(
ignore_input_vars)
super().__init__(
task_type=task_type,
name=f"{task_function.__module__}.{task_function.__name__}",
interface=mutated_interface,
task_config=task_config,
**kwargs,
)
self._task_function = task_function
self._execution_mode = execution_mode
@property
def execution_mode(self) -> ExecutionBehavior:
return self._execution_mode
@property
def task_function(self):
return self._task_function
def execute(self, **kwargs) -> Any:
"""
This method will be invoked to execute the task. If you do decide to override this method you must also
handle dynamic tasks or you will no longer be able to use the task as a dynamic task generator.
"""
if self.execution_mode == self.ExecutionBehavior.DEFAULT:
return self._task_function(**kwargs)
elif self.execution_mode == self.ExecutionBehavior.DYNAMIC:
return self.dynamic_execute(self._task_function, **kwargs)
def get_command(self, settings: SerializationSettings) -> List[str]:
return [
"pyflyte-execute",
"--task-module",
self._task_function.__module__,
"--task-name",
self._task_function.__name__,
"--inputs",
"{{.input}}",
"--output-prefix",
"{{.outputPrefix}}",
"--raw-output-data-prefix",
"{{.rawOutputDataPrefix}}",
]
def compile_into_workflow(
self, ctx: FlyteContext, task_function: Callable, **kwargs
) -> Union[_dynamic_job.DynamicJobSpec, _literal_models.LiteralMap]:
with ctx.new_compilation_context(prefix="dynamic"):
# TODO: Resolve circular import
from flytekit.common.translator import get_serializable
workflow_metadata = WorkflowMetadata(
on_failure=WorkflowFailurePolicy.FAIL_IMMEDIATELY)
defaults = WorkflowMetadataDefaults(interruptible=False)
self._wf = Workflow(task_function,
metadata=workflow_metadata,
default_metadata=defaults)
self._wf.compile(**kwargs)
wf = self._wf
sdk_workflow = get_serializable(ctx.serialization_settings, wf)
# If no nodes were produced, let's just return the strict outputs
if len(sdk_workflow.nodes) == 0:
return _literal_models.LiteralMap(
literals={
binding.var: binding.binding.to_literal_model()
for binding in sdk_workflow._outputs
})
# Gather underlying tasks/workflows that get referenced. Launch plans are handled by propeller.
tasks = set()
sub_workflows = set()
for n in sdk_workflow.nodes:
self.aggregate(tasks, sub_workflows, n)
dj_spec = _dynamic_job.DynamicJobSpec(
min_successes=len(sdk_workflow.nodes),
tasks=list(tasks),
nodes=sdk_workflow.nodes,
outputs=sdk_workflow._outputs,
subworkflows=list(sub_workflows),
)
return dj_spec
@staticmethod
def aggregate(tasks, workflows, node) -> None:
if node.task_node is not None:
tasks.add(node.task_node.sdk_task)
if node.workflow_node is not None:
if node.workflow_node.sdk_workflow is not None:
workflows.add(node.workflow_node.sdk_workflow)
for sub_node in node.workflow_node.sdk_workflow.nodes:
PythonFunctionTask.aggregate(tasks, workflows, sub_node)
if node.branch_node is not None:
if node.branch_node.if_else.case.then_node is not None:
PythonFunctionTask.aggregate(
tasks, workflows, node.branch_node.if_else.case.then_node)
if node.branch_node.if_else.other:
for oth in node.branch_node.if_else.other:
if oth.then_node:
PythonFunctionTask.aggregate(tasks, workflows,
oth.then_node)
if node.branch_node.if_else.else_node is not None:
PythonFunctionTask.aggregate(
tasks, workflows, node.branch_node.if_else.else_node)
def dynamic_execute(self, task_function: Callable, **kwargs) -> Any:
"""
By the time this function is invoked, the _local_execute function should have unwrapped the Promises and Flyte
literal wrappers so that the kwargs we are working with here are now Python native literal values. This
function is also expected to return Python native literal values.
Since the user code within a dynamic task constitute a workflow, we have to first compile the workflow, and
then execute that workflow.
When running for real in production, the task would stop after the compilation step, and then create a file
representing that newly generated workflow, instead of executing it.
"""
ctx = FlyteContext.current_context()
if ctx.execution_state and ctx.execution_state.mode == ExecutionState.Mode.LOCAL_WORKFLOW_EXECUTION:
with ctx.new_execution_context(ExecutionState.Mode.TASK_EXECUTION):
logger.info("Executing Dynamic workflow, using raw inputs")
return task_function(**kwargs)
if ctx.execution_state and ctx.execution_state.mode == ExecutionState.Mode.TASK_EXECUTION:
return self.compile_into_workflow(ctx, task_function, **kwargs)
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"