Workflows: Composing Tasks

Workflows in flytekit are executable directed acyclic graphs (DAGs) that compose tasks, sub-workflows, and other entities into a single unit of execution. flytekit provides two primary ways to define these graphs: a declarative, decorator-based approach for static definitions and an imperative, programmatic approach for dynamic construction.

Decorator-based Workflows

For most use cases, you define workflows using the @workflow decorator. This transforms a standard Python function into a PythonFunctionWorkflow object. The function body describes the data flow between tasks, and flytekit automatically infers the execution graph based on how task outputs are passed as inputs to other tasks.

Defining a Workflow

To create a workflow, decorate a function and call your tasks within it. The inputs and outputs of the function define the interface of the workflow.

from flytekit import task, workflow

@task
def get_message(name: str) -> str:
    return f"Hello, {name}!"

@task
def shout(message: str) -> str:
    return message.upper()

@workflow
def greeting_workflow(name: str) -> str:
    # The data flow here defines the graph: get_message -> shout
    msg = get_message(name=name)
    return shout(message=msg)

How it Works Internally

When you apply the @workflow decorator, flytekit instantiates a PythonFunctionWorkflow.

1. Interface Discovery: It uses transform_function_to_interface to inspect the function's type hints and docstrings, creating a TypedInterface.
2. Compilation: When the workflow is first called or serialized, the compile() method is invoked. It executes the decorated function once, replacing actual inputs with Promise objects.
3. Node Creation: Every time a task or sub-workflow is called inside the function, the flyte_entity_call_handler intercepts the call and creates a Node in the workflow's internal _nodes list.
4. Binding: The compile() method then uses binding_from_python_std to map the function's return values to the workflow's output bindings.

Imperative Workflows

When you need to build a workflow programmatically—for example, if the number of tasks depends on a configuration file or user input—use the ImperativeWorkflow class. This allows you to add inputs, outputs, and entities one by one.

Programmatic Construction

Unlike the decorator-based approach, you must explicitly define every input and output using add_workflow_input and add_workflow_output.

from flytekit.core.workflow import ImperativeWorkflow
from flytekit import task

@task
def t1(a: str) -> str:
    return a + " world"

# 1. Instantiate with a unique name
wb = ImperativeWorkflow(name="my_imperative_workflow")

# 2. Add workflow-level inputs
in1 = wb.add_workflow_input("in1", str)

# 3. Add entities (tasks, sub-workflows, or launch plans)
# add_entity returns a Node object
node = wb.add_entity(t1, a=in1)

# 4. Define workflow-level outputs
# You must reference the specific output of a node
wb.add_workflow_output("final_result", node.outputs["o0"])

Managing Nodes and Outputs

In an ImperativeWorkflow, the add_entity method (and its aliases add_task, add_subwf, add_launch_plan) returns a Node. To connect nodes, you access the outputs attribute of the returned node.

If a task returns multiple values, they are indexed as o0, o1, etc., unless a NamedTuple is used. Attempting to access a non-existent output key will raise a KeyError.

Data Flow and Promises

Both workflow types rely on Promise objects to track data dependencies. When a task is "called" during workflow compilation, it doesn't execute the Python code immediately. Instead, it returns a Promise that represents a future value.

When you pass a Promise from one task to another, flytekit records a dependency between the nodes. The WorkflowBase class maintains these relationships in its _nodes and _output_bindings attributes.

Local Execution

When you execute a workflow locally (e.g., greeting_workflow(name="Flyte")), the WorkflowBase.local_execute method handles the execution:

1. It translates Python native inputs into Flyte literals using translate_inputs_to_literals.
2. It wraps these literals in Promise objects.
3. It calls the underlying execute method, which runs the tasks in the correct order (topological sort).
4. It repacks the resulting literals back into Python types based on the workflow's interface.

Failure Handling

You can define specific logic to run if a workflow fails. This is managed via the on_failure parameter in the @workflow decorator or the add_on_failure_handler method in ImperativeWorkflow.

Failure Nodes

A failure handler can be another task or a workflow. It must accept the same inputs as the workflow itself, or a subset of them.

@task
def cleanup(name: str):
    print(f"Cleaning up for {name}")

@workflow(on_failure=cleanup)
def my_wf(name: str) -> str:
    # ... workflow logic ...
    pass

Internally, flytekit validates that the failure handler's interface is compatible with the workflow's interface. If the failure handler requires inputs that the workflow does not provide (and they aren't optional), flytekit raises a FlyteFailureNodeInputMismatchException. The failure node is assigned a special ID, typically dn or efn, and is stored in the _failure_node attribute of the workflow object.