Skip to main content

Workflow composition, failure handlers, and nodes

Workflows in flytekit are declarative graphs where each step is represented by a Node. While the @workflow decorator is the most common way to compose these graphs, flytekit provides lower-level primitives like create_node and ImperativeWorkflow for programmatic composition and fine-grained control over execution parameters.

Workflow Composition​

Flytekit supports two primary styles of workflow composition: functional (using decorators) and imperative (using a class-based API).

Functional Composition​

In the functional style, you define a workflow by decorating a Python function with @workflow. Inside this function, calling a task returns a Promise object.

from flytekit import task, workflow

@task
def add_one(x: int) -> int:
    return x + 1

@workflow
def my_workflow(val: int) -> int:
    # result is a Promise object
    result = add_one(x=val)
    return result

Imperative Composition​

For scenarios where the workflow structure is determined at runtime (e.g., based on a configuration file), use the ImperativeWorkflow class. This allows you to programmatically add inputs, tasks, and outputs.

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

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

# Create the workflow programmatically
wb = ImperativeWorkflow(name="my_imperative_wf")

# Add a top-level input
wf_input = wb.add_workflow_input("in1", str)

# Add a task and bind its input to the workflow input
node = wb.add_entity(t1, a=wf_input)

# Define the workflow output
wb.add_workflow_output("final_result", node.outputs["o0"])

Working with Nodes​

A Node is the fundamental unit of a flytekit workflow graph. When you call a task inside a @workflow, flytekit automatically creates a Node behind the scenes.

Explicit Node Creation​

You can use create_node from flytekit.core.node_creation to explicitly instantiate a Node. This is particularly useful for establishing execution dependencies between tasks that do not share data.

from flytekit import task, workflow
from flytekit.core.node_creation import create_node

@task
def setup():
    print("Setting up...")

@task
def compute(x: int) -> int:
    return x * 2

@workflow
def manual_node_wf(val: int) -> int:
    setup_node = create_node(setup)
    compute_node = create_node(compute, x=val)
    
    # Ensure setup runs before compute using the shift operator
    setup_node >> compute_node
    
    return compute_node.o0

Accessing Node Outputs​

There is a critical distinction between how you access outputs from a standard task call versus a node created via create_node:

  1. Standard Task Call: Returns a Promise (or a tuple of Promise objects).
  2. create_node: Returns a Node object. You access its outputs via the .outputs dictionary or as attributes (e.g., node.o0, node.o1).

Internally, create_node populates the _outputs attribute of the Node class by mapping the task's interface names to the resulting Promise objects.

Per-Node Overrides​

The Node class provides a with_overrides method to customize execution parameters for a specific instance of a task within a workflow. You can override resources, timeouts, and retries.

@workflow
def override_wf(val: int) -> int:
    # Standard task call returns a Promise, which also supports with_overrides
    promise = add_one(x=val).with_overrides(
        node_name="custom-node-name",
        retries=3,
        timeout=3600
    )
    return promise

The Promise.with_overrides method is a proxy that calls Node.with_overrides on the underlying node reference (self.ref.node).

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 the add_on_failure_handler method in ImperativeWorkflow.

Signature Constraints​

The failure handler must follow strict signature rules enforced by _validate_add_on_failure_handler in workflow.py:

  1. It must accept all inputs defined in the parent workflow.
  2. Any additional inputs it requires must be marked as Optional.
from typing import Optional
from flytekit import task, workflow

@task
def cleanup(err: str, val: int, extra: Optional[str] = None):
    print(f"Workflow failed with error: {err} for input {val}")

@workflow(on_failure=cleanup)
def wf_with_failure(val: int) -> int:
    return add_one(x=val)

During compilation, flytekit verifies these constraints by comparing the python_interface.inputs of the workflow and the failure handler. If the workflow inputs are not a subset of the failure handler's inputs, it raises a FlyteFailureNodeInputMismatchException.

Promises and Data Flow​

The Promise class in flytekit.core.promise acts as a bridge between compilation and execution:

  • At Compile Time: It holds a NodeOutput reference, which tracks which node and which specific output variable the data originates from.
  • At Local Execution: It holds the actual Python value in its _val attribute.

Attribute Access and Indexing​

Promises support attribute access (.) and indexing ([]) to handle complex data types like dataclass or List. When you access an attribute on a Promise, flytekit creates a new Promise with an updated attr_path.

@workflow
def complex_wf(data: MyDataclass) -> int:
    # Accessing an attribute on a workflow input Promise
    result = task_taking_int(x=data.some_integer_field)
    return result

This works because Promise.__getattr__ and Promise.__getitem__ call _append_attr, which deep-copies the promise and updates the _attr_path used for resolving the value during execution.