Workflow composition, failure handlers, and nodes

Flytekit workflows are declarative structures that define a directed acyclic graph (DAG) of tasks. While data dependencies usually determine the execution order, flytekit provides explicit controls for node creation, execution overrides, and failure handling.

Workflow Composition and Promises

In flytekit, calling a task inside a @workflow function does not execute the task immediately. Instead, it returns a Promise object (defined in flytekit.core.promise.Promise). This object represents a future value that will be available when the task completes on the Flyte platform.

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

from flytekit import task, workflow

@task
def get_value() -> int:
    return 10

@task
def process_value(val: int) -> int:
    return val * 2

@workflow
def my_workflow() -> int:
    # get_value() returns a Promise
    val_promise = get_value()
    # Passing the promise to process_value creates a dependency
    return process_value(val=val_promise)

Internally, the Promise class tracks the NodeOutput (the specific output of a node) and the attr_path if you are accessing attributes of a complex object.

Explicit Node Creation

Sometimes you need to define dependencies between tasks that do not share data, or you need fine-grained control over the node itself. The create_node function in flytekit.core.node_creation allows you to explicitly wrap a task, workflow, or launch plan in a Node.

Handling Non-Data Dependencies

If task_b must run after task_a, but task_b does not consume any outputs from task_a, use the >> operator or the runs_before method on the Node objects.

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

@task
def setup():
    ...

@task
def compute():
    ...

@workflow
def manual_dependency_wf():
    setup_node = create_node(setup)
    compute_node = create_node(compute)
    
    # Explicitly order the nodes
    setup_node >> compute_node
    # Equivalent to: setup_node.runs_before(compute_node)

Accessing Node Outputs

When using create_node, the return value is a Node object, not the task's output. To access the outputs for downstream tasks, use the .outputs property or the auto-generated attributes (e.g., .o0, .o1).

@task
def multi_output() -> (int, str):
    return 1, "hello"

@workflow
def output_access_wf():
    node = create_node(multi_output)
    
    # Access outputs via attributes or the outputs dictionary
    val1 = node.o0
    val2 = node.outputs["o1"]
    ...

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. This is useful for adjusting resources, timeouts, or retries without modifying the task definition itself.

You can call with_overrides on the Node returned by create_node or directly on the Promise returned by a standard task call.

from datetime import timedelta
from flytekit import Resources

@workflow
def override_wf(val: int):
    # Overriding on a Promise
    promise = process_value(val=val).with_overrides(
        node_name="custom-process-node",
        requests=Resources(cpu="2", mem="4Gi"),
        timeout=timedelta(minutes=5),
        retries=3
    )
    
    # Overriding on a Node
    node = create_node(get_value).with_overrides(interruptible=True)

The Node.with_overrides method (in flytekit.core.node) updates the NodeMetadata and resource requirements. Supported overrides include:

• node_name: Changes the ID of the node in the graph.
• requests / limits: Sets flytekit.Resources.
• timeout: Accepts an int (seconds) or datetime.timedelta.
• retries: Sets the RetryStrategy.
• interruptible: Boolean for spot instance usage.

Failure Handlers

Flytekit allows you to define a cleanup or recovery strategy using the on_failure parameter in the @workflow decorator. The failure handler can be another task or a workflow.

Signature Requirements

The failure handler must accept:

1. All inputs defined in the workflow's signature.
2. An optional error parameter of type flytekit.types.error.error.FlyteError.

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}")

@task
def failing_task(name: str):
    raise ValueError("Simulated failure")

@workflow(on_failure=clean_up)
def failure_wf(name: str = "flyte-user"):
    failing_task(name=name)

When the workflow fails, Flyte executes the on_failure entity, passing the original workflow inputs and the error details. This is implemented in flytekit.core.workflow.PythonFunctionWorkflow by capturing the failure entity and including it in the serialized workflow model's metadata.