Workflow composition, failure handlers, and nodes

Flytekit workflows are declarative structures that define a Directed Acyclic Graph (DAG) of tasks. While most workflows are constructed implicitly by passing outputs from one task to the inputs of another, flytekit provides lower-level primitives for explicit node management, execution overrides, and failure handling.

Workflow Composition and Promises

When you call a task inside a function decorated with @workflow, flytekit does not execute the task immediately. Instead, it returns a Promise object.

A Promise (defined in flytekit.core.promise.Promise) acts as a placeholder for a future value. These objects are used to track data dependencies. When you pass a Promise from one task call to another, flytekit automatically creates an edge in the workflow DAG.

Accessing Task Outputs

If a task returns multiple values, the call returns a tuple-like object of Promise instances. You can access individual outputs using attribute access or indexing:

from flytekit import task, workflow

@task
def compute_values() -> (int, str):
    return 1, "hello"

@task
def consume_value(v: int):
    print(v)

@workflow
def my_wf():
    # compute_values returns a tuple of Promises
    val_int, val_str = compute_values()
    consume_value(v=val_int)

Internally, Promise objects handle the duality between compilation and local execution. During compilation, they point to a NodeOutput (which references a Node). During local execution, they may wrap actual Python literals.

Explicit Node Creation

In some scenarios, you may need to define execution order without a direct data dependency, or you may want to apply overrides to a specific execution step. The create_node function in flytekit.core.node_creation allows you to explicitly wrap a task, workflow, or launch plan in a Node.

Defining Execution Order

If you have two tasks, t1 and t2, and you want t1 to run before t2 even though t2 doesn't use t1's output, use the >> operator:

from flytekit import task, workflow, create_node

@task
def t1(): ...

@task
def t2(): ...

@workflow
def ordered_wf():
    node1 = create_node(t1)
    node2 = create_node(t2)
    
    # node1 must complete before node2 starts
    node1 >> node2

The >> operator is a shorthand for node1.runs_before(node2). This modifies the _upstream_nodes list of the downstream Node instance.

Accessing Node Outputs

A significant difference between a standard task call and create_node is how outputs are accessed. While a task call returns Promise objects directly, create_node returns a Node object. You access the promises for that node's outputs via the .outputs property or as attributes (e.g., .o0, .o1):

@task
def t1(a: int) -> str:
    return str(a)

@workflow
def wf(a: int):
    n1 = create_node(t1, a=a)
    # Accessing the first output of the node
    t2(val=n1.o0) 
    # Or via the outputs dictionary
    t2(val=n1.outputs["o0"])

Note: The .outputs property is only available on Node objects created via create_node. Accessing it on nodes created implicitly by task calls will raise an AssertionError.

Per-Node Overrides

You can customize the execution parameters of a specific node using the .with_overrides() method. This is available on both Node objects and Promise objects.

from flytekit import Resources

@workflow
def override_wf(a: int):
    # Overriding on a Promise (returned by task call)
    promise = t1(a=a).with_overrides(
        requests=Resources(cpu="2", mem="4Gi"),
        retries=3
    )
    
    # Overriding on an explicit Node
    node = create_node(t2).with_overrides(timeout=3600)

The with_overrides method (implemented in flytekit.core.node.Node) updates the NodeMetadata. Supported overrides include:

• requests and limits: Resource specifications using flytekit.Resources.
• timeout: A datetime.timedelta or integer seconds.
• retries: Number of retry attempts.
• interruptible: Boolean indicating if the node can be run on spot/interruptible instances.
• container_image: Override the image used for this specific task execution.

Workflow Failure Handlers

The @workflow decorator accepts an on_failure parameter to define a "clean-up" or notification task that runs if the workflow fails.

Signature Requirements

The task or workflow passed to on_failure must follow strict signature rules:

1. It must accept all inputs that the main workflow accepts.
2. It can optionally accept a special err argument of type FlyteError (from flytekit.models.core.errors).
3. Any additional arguments must be Optional (have default values).

from typing import Optional
from flytekit import task, workflow
from flytekit.models.core.errors import FlyteError

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

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

In this example, if t1 fails, clean_up is invoked. Flytekit ensures that the name input provided to my_wf is passed to clean_up, and the err parameter is populated with the failure details.

Implementation Detail

When a workflow is defined with on_failure, the PythonFunctionWorkflow object stores this reference. During serialization, flytekit configures the WorkflowMetadata with the failure entity. If the workflow fails, the Flyte platform handles the execution of the failure node, ensuring it receives the original workflow inputs.