Workflow composition, failure handlers, and nodes

Flytekit workflows are defined using the @workflow decorator, which transforms a Python function into a domain-specific language (DSL) for orchestrating tasks. While the code inside a workflow looks like standard Python, it actually constructs a directed acyclic graph (DAG) where each task call creates a node and each return value is a promise of future data.

Workflow Composition and Promises

When you call a task inside a @workflow, flytekit does not execute the task immediately. Instead, it returns a Promise object (defined in flytekit.core.promise.Promise). This promise acts as a proxy for the task's output.

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, not an int
    result = add_one(x=val)
    return result

Because these are promises, you cannot perform standard Python operations on them, such as if result > 0: or for i in range(result):. These operations require concrete values, which are not available during the workflow's compilation phase. Instead, data flow is established by passing the Promise from one task call to another.

Accessing Multiple Outputs

If a task returns multiple values (e.g., a typing.NamedTuple or typing.Tuple), the task call returns a collection of promises. You can access individual outputs by name or index:

import typing
from flytekit import task, workflow

@task
def compute(x: int) -> typing.Tuple[int, str]:
    return x + 1, "success"

@workflow
def multi_output_wf(val: int) -> str:
    res_int, res_str = compute(x=val)
    return res_str

Explicit Node Creation

In most cases, flytekit automatically creates nodes when you call a task. However, if you need to enforce an execution order between tasks that do not share data, you must use create_node from flytekit.core.node_creation.

Establishing Dependencies

You can use the >> operator or the runs_before method on a Node object to define execution order.

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 manual_dependencies_wf():
    setup_node = create_node(setup)
    work_node = create_node(work)
    
    # Enforce setup runs before work
    setup_node >> work_node

Accessing Outputs from create_node

Unlike a direct task call which returns a Promise, create_node returns a Node object (or a VoidPromise if the task has no outputs). To use the outputs of a node created this way, you must access them via the .outputs attribute or named attributes like .o0, .o1, etc.

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

@workflow
def node_output_wf() -> int:
    node = create_node(get_val)
    # Accessing the first output of the node
    return node.o0 

Per-Node Overrides

You can customize the execution behavior of specific nodes within a workflow using the with_overrides method. This method is available on both Node objects and Promise objects (which forward the call to their underlying node).

Common overrides include:

• Resources: Specify requests and limits using flytekit.Resources.
• Retries: Set the number of retries for a failing task.
• Timeout: Define a datetime.timedelta for the maximum execution time.
• Caching: Enable caching by passing a Cache object or boolean.

from datetime import timedelta
from flytekit import task, workflow, Resources

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

@workflow
def override_wf(val: int) -> int:
    promise = heavy_task(x=val).with_overrides(
        requests=Resources(cpu="2", mem="200Mi"),
        retries=3,
        timeout=timedelta(minutes=5),
        node_name="custom-heavy-node"
    )
    return promise

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

Failure Handlers

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

Failure Handler Signature

A failure handler must be a task or another workflow. It must accept all of the original workflow's input arguments by name. It can also optionally accept an err argument of type flytekit.models.core.errors.FlyteError (or typing.Optional[FlyteError]), which contains details about the failure.

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

@task
def clean_up(name: str, val: int, err: typing.Optional[FlyteError] = None):
    print(f"Cleaning up for {name} with input {val}. Error: {err}")

@task
def failing_task(a: int):
    raise ValueError("Something went wrong")

@workflow(on_failure=clean_up)
def failure_wf(name: str, val: int):
    failing_task(a=val)

When failure_wf fails, flytekit automatically invokes clean_up, passing the name and val provided to the workflow, along with the error that caused the failure. This mechanism is implemented in flytekit/core/workflow.py, where the on_failure entity is stored in the WorkflowMetadata.