Workflow composition, failure handlers, and nodes

Flytekit workflows are declarative structures that define a Directed Acyclic Graph (DAG) of tasks. While the @workflow decorator makes these look like standard Python functions, they are actually evaluated at compile-time to build a graph of Node objects connected by Promise references.

Workflow Composition and Promises

In a flytekit workflow, calling a task does not return the actual result of the computation. Instead, it returns a Promise (or a tuple of promises). These promises act as placeholders that represent the future output of a node.

When you pass a promise from one task to another, flytekit automatically creates a dependency between the underlying nodes.

from flytekit import task, workflow

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

@task
def greet(greeting: str):
    print(greeting)

@workflow
def my_workflow(name: str):
    # 'greeting' is a Promise object, not a string
    greeting = get_greeting(name=name)
    # Passing the promise to greet() creates a dependency
    greet(greeting=greeting)

Internally, the Promise class (found in flytekit/core/promise.py) tracks the Node that produces the value. If a task returns multiple values, flytekit returns a NamedTuple of Promise objects, allowing you to access specific outputs by name (e.g., output.o0).

Explicit Node Creation and Dependencies

Sometimes you need to define execution order between tasks that do not share data. For example, you might want a cleanup task to run only after a processing task completes, even if the cleanup task doesn't take the processing task's output as an input.

The create_node function in flytekit/core/node_creation.py allows you to explicitly instantiate a Node for a task or sub-workflow. Once you have a Node object, you can use the >> operator or the runs_before method to enforce ordering.

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

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

@task
def work():
    print("Working...")

@workflow
def ordered_workflow():
    setup_node = create_node(setup)
    work_node = create_node(work)
    
    # Ensure setup runs before work
    setup_node >> work_node

Accessing Outputs via create_node

A key difference between calling a task directly and using create_node is how you access outputs. While a task call returns a Promise, create_node returns a Node object. To get a promise from a Node, you access its outputs attribute.

@task
def produce_value() -> int:
    return 42

@workflow
def output_access_wf():
    node = create_node(produce_value)
    
    # Accessing the output 'o0' from the node's outputs dictionary
    # This returns a Promise that can be passed to other tasks
    val_promise = node.outputs["o0"]
    
    # Alternatively, attributes are dynamically added to the node object
    # val_promise = node.o0

In flytekit/core/node_creation.py, the create_node implementation ensures that for every output defined in the entity's interface, a corresponding attribute and dictionary entry is added to the Node object during compilation.

Per-Node Overrides

You can customize the execution behavior of specific nodes within a workflow using the with_overrides method. This is available on both Promise objects and Node objects.

Common overrides include:

• Resources: requests and limits for CPU, memory, and GPU.
• Retries: The number of times to retry a failed node.
• Timeout: A datetime.timedelta or integer seconds.
• Interruptible: Whether the node can run on spot/preemptible instances.

from datetime import timedelta
from flytekit import Resources

@workflow
def override_wf(val: int):
    # Overriding via the Promise returned by a task call
    promise = get_greeting(name="Flyte").with_overrides(
        retries=3,
        timeout=timedelta(minutes=5)
    )
    
    # Overriding via an explicit Node
    node = create_node(greet, greeting=promise)
    node.with_overrides(
        requests=Resources(cpu="2", mem="500Mi"),
        node_name="custom-greet-node"
    )

The Node.with_overrides method in flytekit/core/node.py handles these updates by modifying the NodeMetadata and Resources models associated with the node.

Failure Handlers (on_failure)

Flytekit allows you to define a specific task or workflow to execute if any node in your workflow fails. This is configured using the on_failure parameter in the @workflow decorator.

Signature Requirements

A failure handler must accept all the inputs of the workflow it is protecting. Additionally, it can optionally accept a parameter named err of type flytekit.types.error.error.FlyteError to receive details about the failure.

from typing import Optional
from flytekit import task, workflow
from flytekit.types.error.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}")
    else:
        print(f"Cleaning up for {name}")

@workflow(on_failure=clean_up)
def wf_with_handler(name: str):
    # If this task fails, clean_up(name=name, err=...) is invoked
    get_greeting(name=name)

When a failure occurs, Flyte executes the on_failure entity, passing the original workflow inputs to it. The err parameter, if present in the signature, is automatically populated with a FlyteError object containing the failed_node_id and the error message.

Imperative Workflows

For scenarios where the workflow structure is dynamic or programmatically generated, flytekit provides the ImperativeWorkflow class. This allows you to build a workflow by explicitly adding inputs, nodes, and outputs.

from flytekit.core.workflow import ImperativeWorkflow

# Create the workflow object
wb = ImperativeWorkflow(name="dynamic_wf")

# Add inputs
input_promise = wb.add_workflow_input("val", int)

# Add tasks as nodes
node = wb.add_task(get_greeting, name=input_promise)

# Add an on_failure handler
wb.add_on_failure_handler(clean_up)

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

The ImperativeWorkflow methods like add_task and add_launch_plan internally call create_node and manage the registration of nodes within the workflow's internal list.