Workflow composition, failure handlers, and nodes

Flytekit allows you to compose complex workflows by chaining tasks, defining explicit execution orders, and handling failures gracefully. While standard task calls return Promise objects that represent future values, you can also use create_node for fine-grained control over node metadata and dependencies.

Workflow Composition and Promises

When you call a task within a @workflow function, flytekit does not execute the task immediately. Instead, it creates a Node in the workflow graph and returns one or more Promise objects. These promises act as placeholders for the task's future outputs.

from flytekit import task, workflow

@task
def get_data() -> str:
    return "flyte"

@task
def process_data(val: str) -> int:
    return len(val)

@workflow
def my_workflow() -> int:
    # data is a Promise object
    data = get_data()
    # Passing the promise to another task creates a data dependency
    return process_data(val=data)

Internally, the Promise class (found in flytekit/core/promise.py) wraps a NodeOutput. When a promise is passed as an input to another task, flytekit's compilation state records the dependency between the upstream node that produces the output and the downstream node that consumes it.

Accessing Nested Outputs

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 records these accesses in the attr_path of the Promise and resolves them during execution.

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

@workflow
def nested_wf() -> int:
    m = get_map()
    # Accessing nested keys returns a new Promise with an updated attr_path
    return m["a"]["b"]

Explicit Node Creation

In some scenarios, you may need to define dependencies between tasks that do not share data. For example, if task_b must run after task_a even though task_b doesn't consume task_a's output, you can use create_node from flytekit/core/node_creation.py.

The create_node function returns a Node object (or a VoidPromise if the entity has no outputs). You can use the >> operator or the runs_before method to enforce execution order.

from flytekit import task, workflow, create_node

@task
def setup():
    ...

@task
def compute():
    ...

@workflow
def ordered_wf():
    setup_node = create_node(setup)
    compute_node = create_node(compute)
    
    # Enforce that setup runs before compute
    setup_node >> compute_node

Accessing Outputs from create_node

Unlike standard task calls that return promises directly, create_node returns a Node object. Its outputs are accessed via the .outputs dictionary or as attributes on the node itself (e.g., node.o0).

@task
def t1(a: int) -> str:
    return str(a)

@workflow
def node_output_wf(val: int) -> str:
    n1 = create_node(t1, a=val)
    # Access output by name (default is o0 for single output)
    return n1.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. Overrides allow you to specify resources, retries, timeouts, and more for a specific instance of a task in a workflow.

from flytekit import Resources

@workflow
def override_wf(val: int):
    # Applying overrides to a Promise
    t1(a=val).with_overrides(
        requests=Resources(cpu="2", mem="4Gi"),
        retries=3,
        timeout=3600  # seconds
    )

The Node.with_overrides method in flytekit/core/node.py updates the NodeMetadata and resource requirements for that specific node. Supported overrides include:

• node_name: Customizes the ID of the node in the graph.
• requests / limits: Sets flytekit.Resources.
• timeout: A datetime.timedelta or integer seconds.
• retries: Number of retry attempts.
• interruptible: Boolean flag for spot/preemptible instance usage.

Failure Handlers

Flytekit supports on_failure handlers at the workflow level. A failure handler is a task or workflow that executes only if the primary workflow fails. This is useful for cleaning up resources or sending notifications.

The handler must accept all inputs of the original workflow. It can optionally accept an additional err parameter of type FlyteError to inspect the failure details.

from flytekit import task, workflow, FlyteError
from typing import Optional

@task
def clean_up(name: str, err: Optional[FlyteError] = None):
    print(f"Cleaning up for {name}")
    if err:
        print(f"Caught error: {err.message}")

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

@workflow(on_failure=clean_up)
def wf_with_handler(name: str):
    failing_task(name=name)

When wf_with_handler fails, Flyte invokes clean_up with the original name input and the error details. Note that the err parameter in the handler must be named exactly err or error (as per flytekit/core/workflow.py logic) and should typically be Optional to maintain compatibility with local execution or manual triggers.