Workflow composition, failure handlers, and nodes

Workflows in flytekit are the primary mechanism for composing tasks and other workflows into complex data pipelines. While the @workflow decorator provides a declarative way to define these pipelines, the underlying implementation relies on a graph of Node objects, data-passing Promise objects, and failure handling mechanisms.

Workflow Composition and Promises

When you define a workflow using the @workflow decorator, flytekit tracks the execution of tasks within the function. Each task call does not execute the task immediately; instead, it returns one or more Promise objects.

A Promise (defined in flytekit.core.promise) acts as a placeholder for a future value. It contains a reference to the Node that will produce the value.

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 welcome_wf(name: str):
    # 'greeting' is a Promise object, not a string
    greeting = get_greeting(name=name)
    greet(greeting=greeting)

Promise Gotchas: Truth Value Testing

Because a Promise is a placeholder for a value that doesn't exist yet during compilation, you cannot use it in standard Python boolean contexts. Attempting to use if my_promise: or while my_promise: will raise a ValueError.

For logical operations in conditionals, flytekit provides bitwise operator overrides (& for AND, | for OR) which produce ConjunctionExpression objects instead of evaluating to a boolean.

Explicit Node Creation

In most cases, flytekit automatically creates nodes when you call a task. However, you may need to use create_node from flytekit.core.node_creation to define execution dependencies that are not based on data flow (i.e., when a task doesn't take the output of another task as input).

Defining Execution Order

You can use the >> operator (or the runs_before method) to enforce that one node runs after another.

from flytekit import task, workflow, create_node

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

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

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

Accessing Node Outputs

Nodes created via create_node expose their outputs differently than standard task calls. While a task call returns a Promise (or a tuple of them), create_node returns a Node object. You access the outputs of this node using the .o0, .o1, etc., attributes or the .outputs dictionary.

@task
def compute() -> (int, str):
    return 1, "done"

@workflow
def output_access_wf():
    node = create_node(compute)
    
    # Access by attribute
    use_int(val=node.o0)
    # Access by dictionary key
    use_str(val=node.outputs["o1"])

Note that node.outputs is only populated for nodes created via create_node(). Calling .outputs on a node generated by a standard task call will result in 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 the Node object (returned by create_node) and the Promise object (returned by a task call).

Common overrides include:

• requests and limits: Resource requirements using flytekit.Resources.
• timeout: A datetime.timedelta or integer seconds.
• retries: Number of times to retry on failure.
• interruptible: Boolean indicating if the node can run on low-priority/spot instances.

from flytekit import Resources

@workflow
def override_wf(val: int):
    # Applying overrides to a Promise
    promise = task_a(val=val).with_overrides(
        retries=3,
        requests=Resources(cpu="2", mem="500Mi")
    )
    
    # Applying overrides to a Node
    node = create_node(task_b).with_overrides(timeout=600)

Internally, with_overrides modifies the NodeMetadata and resource specifications stored within the Node class in flytekit.core.node.

Failure Handlers

Flytekit allows you to define a cleanup or notification task that runs if a workflow fails. This is configured via the on_failure parameter in the @workflow decorator.

The on_failure Contract

A failure handler must be a @task or another @workflow. It must satisfy two requirements:

1. It must accept all inputs that the parent workflow accepts.
2. It can optionally accept an err argument of type flytekit.types.error.FlyteError to receive details about the failure.

from flytekit.types.error import FlyteError

@task
def cleanup_task(name: str, err: FlyteError):
    print(f"Workflow for {name} failed!")
    print(f"Error: {err.message} in node {err.failed_node_id}")

@workflow(on_failure=cleanup_task)
def my_wf(name: str):
    # If any task here fails, cleanup_task is invoked with 'name' and the error
    risky_task(name=name)

When a failure occurs, flytekit's WorkflowBase.__call__ catches the exception and, if an on_failure handler is present, injects the workflow's original input_kwargs and the FlyteError into the handler before re-raising the exception.