Workflow composition, failure handlers, and nodes

Workflow composition in flytekit allows you to define data dependencies and execution order between tasks, sub-workflows, and launch plans. While most users use the declarative @workflow decorator, flytekit also provides an imperative API for programmatic construction and manual node management.

Workflow Composition Styles

You can compose workflows using two primary patterns: declarative functions and imperative builders.

Declarative Workflows

The @workflow decorator transforms a Python function into a PythonFunctionWorkflow. When you call a task inside this function, flytekit doesn't execute the task immediately. Instead, it records the call as a node in a Directed Acyclic Graph (DAG).

import typing
from flytekit import task, workflow

@task
def t1(a: int) -> typing.NamedTuple("Outputs", [("val", int), ("msg", str)]):
    return a + 2, f"result: {a + 2}"

@workflow
def my_workflow(a: int) -> str:
    # Task call returns a Promise or a collection of Promises
    res = t1(a=a)
    return res.msg

Imperative Workflows

For scenarios where the workflow structure is dynamic or programmatically generated, use ImperativeWorkflow. This class allows you to manually add inputs, entities, and outputs.

from flytekit import ImperativeWorkflow

wb = ImperativeWorkflow(name="my_imperative_wf")
# Returns a Promise backed by the global start node
wf_input = wb.add_workflow_input("in1", int)

# add_entity creates a node and binds inputs
node = wb.add_entity(t1, a=wf_input)

# Expose a specific node output as a workflow output
wb.add_workflow_output("final_msg", node.outputs["msg"])

Nodes and Explicit Creation

A Node represents a unit of execution in the DAG. While ordinary task calls return Promise objects, you can use create_node to obtain the underlying Node object directly.

Using create_node

The create_node function is useful for specifying dependencies between tasks that do not share data, or when you need to access outputs by string keys.

from flytekit.core.node_creation import create_node

@workflow
def manual_node_wf(a: int):
    # Explicitly create nodes
    t1_node = create_node(t1, a=a)
    t2_node = create_node(t2)

    # Specify execution order without data dependency
    t1_node >> t2_node  # or t1_node.runs_before(t2_node)

Accessing Node Outputs

There is a critical distinction between the result of a standard task call and the result of create_node:

• Standard Task Call: Returns a Promise, a NamedTuple of Promises, or a VoidPromise. These objects represent the future value of the task.
• create_node: Returns a Node object during compilation. The outputs are accessible via the .outputs attribute (a dictionary) or as direct attributes on the node (e.g., node.o0).

# Standard call: returns a Promise
p = t1(a=10)

# create_node: returns a Node
n = create_node(t1, a=10)
# Accessing outputs via the node
val_promise = n.outputs["val"] 
msg_promise = n.val  # Attributes are also dynamically set

Note that Node.outputs is only available on nodes created via create_node. Accessing .outputs on a standard node will raise an AssertionError if it wasn't initialized by the creation utility.

Promises and Output Handling

Promise objects are placeholders for values that will be computed at runtime. They handle the duality between compilation (where they are NodeOutput references) and local execution (where they wrap actual Literal values).

Output Shapes

The create_task_output utility normalizes task results based on the task's interface:

• Zero outputs: Returns a VoidPromise. This object supports dependency chaining (>>) but rejects value operations like comparisons.
• Single output: Returns a single Promise object.
• Multiple outputs: Returns a NamedTuple where each field is a Promise.

Attribute and Index Access

You can access attributes or indexes on a Promise to pass specific parts of a complex structure (like a dict or dataclass) to downstream tasks. This creates a new Promise with an updated attr_path.

@workflow
def attribute_wf(data: dict):
    # Accessing a key in a dictionary promise
    t2(val=data["some_key"])
    
    # Accessing an attribute on a dataclass promise
    res = t1(a=5)
    t2(val=res.val)

Per-Node Overrides

You can customize the execution behavior of individual nodes using the .with_overrides() method. This method is available on both Node objects and the Promise (or NamedTuple) returned by task calls.

from flytekit import Resources

@workflow
def override_wf(a: int):
    # Override on a task call result
    t1(a=a).with_overrides(
        requests=Resources(cpu="2", mem="500Mi"),
        retries=3,
        node_name="custom-t1-node"
    )

Internally, Promise.with_overrides delegates to the underlying Node.with_overrides. Supported overrides include:

• Resources: requests and limits (using flytekit.Resources). Note that you cannot use the resources parameter simultaneously with requests or limits.
• Metadata: timeout (as datetime.timedelta or int seconds), retries, and interruptible.
• Caching: cache (boolean or Cache object) and cache_version.
• Infrastructure: container_image, pod_template, and accelerator.

Failure Handlers

Flytekit allows you to define a specific task or workflow to run if a workflow fails. This is configured via the on_failure parameter in the @workflow decorator or add_on_failure_handler in imperative workflows.

Defining a Handler

A failure handler must accept all inputs of the workflow it protects. It can also optionally accept an error payload by including an input named err of type Optional[FlyteError].

from typing import Optional
from flytekit.types.error import FlyteError

@task
def clean_up(name: str, err: Optional[FlyteError] = None):
    if err:
        print(f"Workflow failed for {name} with error: {err.message}")

@workflow(on_failure=clean_up)
def my_wf(name: str):
    t1(a=10)

Critical Rules for Handlers

1. Input Superset: The handler's signature must be a superset of the workflow's signature.
2. Optional Extras: Any inputs in the handler that are not present in the workflow must be Optional.
3. Error Injection: If the handler has an input named exactly err, flytekit injects a FlyteError containing the failed node ID and the exception message.
4. No Swallowing: The failure handler does not "catch" the exception in the sense of stopping it. After the handler runs, the original exception is re-raised.

Failure Policies

You can control how the workflow behaves when a node fails using WorkflowFailurePolicy:

• FAIL_IMMEDIATELY (Default): Stops the workflow as soon as any node fails.
• FAIL_AFTER_EXECUTABLE_NODES_COMPLETE: Continues executing other nodes in the DAG that do not depend on the failed node before failing the workflow.

from flytekit import WorkflowFailurePolicy

@workflow(failure_policy=WorkflowFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE)
def robust_wf():
    ...