Workflow composition, failure handlers, and nodes

Flytekit workflows are declarative structures that define a Directed Acyclic Graph (DAG) of tasks. While simple workflows rely on data flow to determine execution order, flytekit provides advanced composition tools for explicit node management, per-node configuration overrides, and robust failure handling.

Workflow Composition and Promises

In flytekit, calling a @task or @workflow inside another workflow does not execute the function immediately. Instead, it returns a Promise (defined in flytekit.core.promise.Promise). A Promise represents a future value that will be available during execution.

When you pass a Promise from one task as an input to another, flytekit implicitly creates a dependency between the corresponding nodes in the DAG.

from flytekit import task, workflow

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

@task
def process_data(val: int) -> str:
    return f"Processed {val}"

@workflow
def my_workflow() -> str:
    # data_promise is a Promise object
    data_promise = get_data()
    # Passing the promise creates a dependency: get_data -> process_data
    return process_data(val=data_promise)

Accessing Promise Attributes

If a task returns a complex type like a dataclass or a dict, you can access its attributes or keys directly on the Promise. Flytekit tracks these accesses in the attr_path of the Promise and resolves them at runtime.

@task
def get_map() -> dict:
    return {"a": 1, "b": 2}

@workflow
def attr_workflow() -> int:
    m = get_map()
    # Accessing "a" creates a new Promise with an updated attr_path
    return m["a"] 

Explicit Node Creation

While implicit dependencies via data flow are common, you may need to define execution order for tasks that do not share data, or apply specific overrides to a single invocation. The create_node function in flytekit.core.node_creation allows for explicit node management.

Dependency Management without Data Flow

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

from flytekit.core.node_creation import create_node

@workflow
def explicit_wf():
    n1 = create_node(task_a)
    n2 = create_node(task_b)
    
    # Ensure task_a runs before task_b even without data dependency
    n1 >> n2

Accessing Node Outputs

Unlike regular task calls that return Promise objects directly, create_node returns a Node object. To access the outputs of the underlying entity, you must use the .outputs attribute or named attributes like .o0, .o1, etc.

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

@workflow
def node_output_wf():
    node = create_node(multi_output)
    
    # Accessing outputs via attributes (o0, o1, ...)
    # or via the outputs dictionary
    use_task(val=node.o0, msg=node.outputs["o1"])

Per-Node Overrides

The Node class in flytekit.core.node provides a with_overrides method to customize execution parameters for a specific instance of a task or workflow. This is useful for adjusting resources, retries, or timeouts without changing the task definition itself.

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

from flytekit import Resources

@workflow
def override_wf(val: int):
    # Applying overrides to a standard task call promise
    t1_promise = task_a(val=val).with_overrides(
        requests=Resources(cpu="2", mem="4Gi"),
        retries=3,
        node_name="custom-node-name"
    )
    
    # Applying overrides via create_node
    node = create_node(task_b, val=t1_promise).with_overrides(
        timeout=3600,
        interruptible=True
    )

Key overrides supported by Node.with_overrides:

• requests / limits: Specify flytekit.Resources for CPU, memory, and storage.
• timeout: A datetime.timedelta or integer seconds.
• retries: Number of times to retry the node on failure.
• interruptible: Boolean indicating if the node can be run on spot/preemptible instances.
• container_image: Override the image used for this specific node.

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.

Input Requirements for Failure Handlers

The on_failure handler must be a @task or @workflow that satisfies specific signature requirements:

1. It must accept all inputs defined in the parent workflow.
2. It may accept an optional err argument of type flytekit.types.error.FlyteError to receive failure details.
3. Any additional arguments must be Optional (defaulting to None).

Flytekit enforces these rules in WorkflowBase._validate_add_on_failure_handler during compilation.

from typing import Optional
from flytekit import task, workflow
from flytekit.types.error import FlyteError

@task
def clean_up(wf_input_a: int, err: Optional[FlyteError] = None, extra: Optional[str] = None):
    if err:
        print(f"Workflow failed with error: {err.message}")
    print(f"Cleaning up for input: {wf_input_a}")

@workflow(on_failure=clean_up)
def failure_wf(wf_input_a: int) -> int:
    return task_that_might_fail(a=wf_input_a)

Failure Policies

You can further control workflow behavior on failure using the failure_policy argument:

• WorkflowFailurePolicy.FAIL_IMMEDIATELY (Default): The workflow stops as soon as any node fails.
• WorkflowFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE: The workflow continues to run other nodes that do not depend on the failed node before failing.

from flytekit.core.workflow import WorkflowFailurePolicy

@workflow(failure_policy=WorkflowFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE)
def robust_wf():
    ...