Workflow composition, failure handlers, and nodes

Flytekit workflows are defined using the @workflow decorator, which transforms a Python function into a structured execution graph. Within this graph, every task call or sub-workflow invocation creates a Node. These nodes represent the fundamental units of execution and can be customized with specific resource requirements, retry policies, and failure handlers.

Workflow Composition and Promises

When you call a task inside a @workflow function, flytekit does not execute the task immediately. Instead, it returns a Promise object. This object acts as a placeholder for a value that will be computed at runtime.

@workflow
def my_workflow(val: int) -> int:
    # t1 returns a Promise[int] during compilation
    result = t1(a=val)
    # The promise is passed to t2, establishing a data dependency
    return t2(b=result)

Internally, the Promise class (found in flytekit/core/promise.py) manages the duality between local execution and remote compilation. During compilation, a Promise holds a NodeOutput reference, which points to the Node that will produce the value.

Accessing Task Outputs

Tasks can return single values, tuples, or NamedTuple objects. Flytekit handles these by wrapping them in Promise objects that support attribute access and indexing.

@task
def compute_stats(data: list[int]) -> typing.NamedTuple("Stats", mean=float, std=float):
    ...

@workflow
def stats_workflow(data: list[int]) -> float:
    stats = compute_stats(data=data)
    # Accessing an attribute on a Promise returns a new Promise with an updated attribute path
    return stats.mean

The Promise.__getattr__ and Promise.__getitem__ methods in flytekit/core/promise.py implement this by appending the key to the _attr_path. This allows Flyte to resolve the specific field from the task's output at runtime.

Explicit Node Creation

While data dependencies (passing a Promise from one task to another) implicitly define the execution order, you may sometimes need to enforce order between tasks that do not share data. The create_node function in flytekit/core/node_creation.py allows you to explicitly instantiate a Node.

Ordering with the Shift Operator

You can use the >> operator (right shift) to define execution order between nodes.

from flytekit.core.node_creation import create_node

@workflow
def ordered_workflow():
    n1 = create_node(task_a)
    n2 = create_node(task_b)
    
    # task_a will run before task_b
    n1 >> n2

The Node.__rshift__ method in flytekit/core/node.py calls runs_before, which appends the upstream node to the _upstream_nodes list of the downstream node.

Accessing Outputs from create_node

Unlike standard task calls that return Promise objects directly, create_node returns a Node object (or a VoidPromise if the task has no outputs). To access the outputs of a node created this way, use the .outputs attribute or named attributes like .o0, .o1, etc.

@workflow
def explicit_output_wf(a: int) -> int:
    node = create_node(t1, a=a)
    # Accessing output via attribute (o0 is the first output)
    return t2(b=node.o0)
    # Or via the outputs dictionary
    # return t2(b=node.outputs["o0"])

Per-Node Overrides

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

from flytekit import Resources

@workflow
def resource_wf(val: int) -> int:
    return t1(a=val).with_overrides(
        requests=Resources(cpu="2", mem="500Mi"),
        limits=Resources(cpu="4", mem="1Gi"),
        retries=3,
        timeout=3600  # seconds
    )

The Node.with_overrides method in flytekit/core/node.py updates the NodeMetadata and resource specifications. It supports overriding:

• requests and limits: Using flytekit.Resources.
• retries: Number of retry attempts.
• timeout: A datetime.timedelta or integer seconds.
• container_image: A specific image for this node.
• interruptible: Boolean flag for spot/preemptible instance usage.

Workflow 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 of the @workflow decorator.

Signature Requirements

A failure handler must be a task or workflow that accepts:

1. All inputs defined in the parent workflow's signature.
2. An optional err argument of type typing.Optional[FlyteError].

from flytekit.types.error.error import FlyteError

@task
def cleanup_task(name: str, err: typing.Optional[FlyteError] = None):
    print(f"Workflow for {name} failed with error: {err}")

@workflow(on_failure=cleanup_task)
def main_wf(name: str):
    t1(a=name)

If the workflow fails, Flyte will invoke the on_failure entity. The FlyteError object provides details about the failure, including the message and the failed_node_id. The PythonFunctionWorkflow._validate_add_on_failure_handler method in flytekit/core/workflow.py ensures that the handler's interface is compatible with the workflow's inputs.