Workflow composition, failure handlers, and nodes

Workflow Composition and Promises

When you call a task inside a @workflow function, flytekit does not run the task's Python body. Instead, it compiles the call into a Node and hands you back a Promise (or a tuple of them). Promises are stand-ins for future outputs; they let you wire outputs from one task into the inputs of another without executing anything yet.

from flytekit import task, workflow

@task
def add_5(a: int) -> int:
    a = a + 5
    return a

@workflow
def my_wf(a: int) -> int:
    x = add_5(a=a)
    z = add_5(a=x)
    return z

When my_wf is compiled, flytekit evaluates the function body. add_5(a=a) returns a Promise pointing to the newly created Node in flytekit/core/node.py. When you pass that Promise as a=x into the next add_5 call, flytekit records a data dependency: the second node binds its a input to the first node's output.

This compilation behavior is why the docstring in flytekit/core/workflow.py warns:

Even though you may have a task t1() -> int, when a = t1() is called, a will not be an integer so if you try to range(a) you'll get an error.

create_node and .outputs

Tasks with no outputs, or tasks whose outputs you don't need to thread forward, still need a way to express ordering. Calling t1() returns a VoidPromise (defined in flytekit/core/promise.py) which blocks comparison and arithmetic operators. To order such tasks, or to get a handle to a node directly, use create_node from flytekit/core/node_creation.py.

from flytekit.core.node_creation import create_node

t1_node = create_node(t1)
t2_node = create_node(t2)

t2_node.runs_before(t1_node)
# or
t2_node >> t1_node

create_node accepts a task, workflow, or launch plan plus keyword arguments for inputs. Under compilation it calls the entity, pulls the last node out of ctx.compilation_state.nodes, and attaches the outputs to it:

node._outputs = {}
...
if isinstance(outputs, tuple):
    for output_name in entity.interface.outputs.keys():
        attr = getattr(outputs, output_name)
        setattr(node, output_name, attr)
        node.outputs[output_name] = attr

This is why node.outputs exists and why it raises AssertionError("Cannot use outputs with all Nodes, node must've been created from create_node()") for nodes produced by an ordinary task call. Promise objects returned by t1() are thin wrappers around a NodeOutput; they expose .ref, .is_ready, and comparison helpers. create_node(...).outputs is a dict[str, Promise] populated by node_creation.create_node, letting you look up outputs by name string, which is useful in imperative workflows:

node = wb.add_entity(t1, a=wb.inputs["in1"])
wb.add_workflow_output("from_n0t1", node.outputs["o0"])

Per-Node Overrides

Once you have a Node (via create_node or as the backing object of a Promise), you can customize its execution with with_overrides. It lives on Node in flytekit/core/node.py and is also forwarded by Promise.with_overrides and VoidPromise.with_overrides.

from flytekit import Resources

@workflow
def wf(a: int) -> int:
    return add_5(a=a).with_overrides(
        node_name="renamed-node",
        timeout=60,
        retries=3,
        requests=Resources(cpu="1"),
        limits=Resources(mem="2Gi"),
        cache=True,
        cache_version="1",
        container_image="ghcr.io/flyteorg/flytekit:latest",
    )

with_overrides mutates the node and returns it. Supported arguments include:

• node_name / name — overrides the node's DNS-compliant ID (via _dnsify)
• aliases — maps output names to new variable names as a dict[str, str]
• requests, limits, resources — Resources objects; resources cannot be combined with requests or limits
• timeout — int seconds or datetime.timedelta
• retries, interruptible
• cache and cache_version — cache=True with no version raises ValueError("must specify cache version when overriding"); Cache objects are also supported
• container_image, pod_template, accelerator, shared_memory
• task_config — must match the type of the existing task config

_override_node_metadata is the internal helper that applies metadata changes. For ArrayNodeMapTask instances it overrides the sub-node's metadata rather than the map task's.

Ordering Nodes Explicitly

When outputs are not enough to establish ordering, use runs_before or the >> operator. Both append the upstream node to other._upstream_nodes:

def runs_before(self, other: Node):
    if self not in other._upstream_nodes:
        other._upstream_nodes.append(self)

def __rshift__(self, other: Node):
    self.runs_before(other)
    return other

Promise and VoidPromise also implement __rshift__, which delegates to Node.runs_before when both promises have a ref.

Failure Policies and On-Failure Handlers

By default, a workflow fails as soon as any node fails. Pass failure_policy=WorkflowFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE to let already-running nodes finish before the workflow transitions to a failed state. WorkflowFailurePolicy is an Enum in flytekit/core/workflow.py that maps to the Flyte IDL integer values (0 for immediate, 1 for after-executable-nodes-complete).

from flytekit import WorkflowFailurePolicy

@workflow(failure_policy=WorkflowFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE)
def wf(a: int) -> typing.Tuple[str, str]:
    ...

For cleanup logic, pass an on_failure handler. The handler can be a @task or @workflow that accepts every input of the parent workflow, plus optional additional inputs.

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

@task
def clean_up(name: str, err: typing.Optional[FlyteError] = None):
    print(f"Deleting cluster {name} due to {err}")

@task
def create_cluster(name: str):
    print(f"Creating cluster: {name}")

@task
def t1(a: int, b: str):
    raise ValueError("boom")

@workflow(on_failure=clean_up)
def wf(name: str = "flyteorg"):
    c = create_cluster(name=name)
    t = t1(a=1, b="2")
    c >> t

Input Mismatch Validation

PythonFunctionWorkflow._validate_add_on_failure_handler (in flytekit/core/workflow.py) and ImperativeWorkflow.add_on_failure_handler both enforce two rules:

# Workflow inputs should be a subset of failure node inputs.
if (failure_node_inputs | workflow_inputs) != failure_node_inputs:
    raise FlyteFailureNodeInputMismatchException(self.on_failure, self)
additional_keys = failure_node_inputs.keys() - workflow_inputs.keys()
# Raising an error if the additional inputs in the failure node are not optional.
for k in additional_keys:
    if not is_optional_type(failure_node_inputs[k]):
        raise FlyteFailureNodeInputMismatchException(self.on_failure, self)

FlyteFailureNodeInputMismatchException is imported from flytekit.exceptions.user. is_optional_type comes from typing_inspect. The handler must accept every workflow input, and any extra input it declares must be Optional. A common extra input is err.

Error Injection at Local Execution Time

When a workflow runs locally and raises, WorkflowBase.__call__ catches the exception, invokes the on_failure entity with the workflow's inputs, and re-raises:

try:
    return flyte_entity_call_handler(self, *args, **input_kwargs)
except Exception as exc:
    if self.on_failure:
        if self.on_failure.python_interface and "err" in self.on_failure.python_interface.inputs:
            id = self.failure_node.id if self.failure_node else ""
            input_kwargs["err"] = FlyteError(failed_node_id=id, message=str(exc))
        self.on_failure(**input_kwargs)
    raise exc

If the handler declares an err input, flytekit constructs a FlyteError(failed_node_id=..., message=str(exc)) and passes it in. On the Flyte backend, the failure node is compiled with ID DEFAULT_FAILURE_NODE_ID ("efn" from flytekit/core/constants.py) and the error from the failed node is wired into err automatically.

Imperative Failure Handlers

ImperativeWorkflow exposes add_on_failure_handler, which validates the handler's interface, creates a node, pops it off compilation_state.nodes so it isn't part of the main DAG, and stores it on self._failure_node with n._id = _common_constants.DEFAULT_FAILURE_NODE_ID.

n = create_node(entity=entity, **self._inputs)
ctx.compilation_state.nodes.pop(-1)
self._failure_node = n
n._id = _common_constants.DEFAULT_FAILURE_NODE_ID

This mirrors how PythonFunctionWorkflow._validate_add_on_failure_handler compiles the handler inside an inner CompilationState with prefix prefix + "f", requiring exactly one node:

if not inner_nodes or len(inner_nodes) > 1:
    raise AssertionError("Unable to compile failure node, only either a task or a workflow can be used")
self._failure_node = inner_nodes[0]

Promises, VoidPromises, and Node Outputs

A Promise in flytekit/core/promise.py wraps either a ready Literal (local execution) or a NodeOutput reference (compilation). Promise.is_ready distinguishes the two. NodeOutput is a subclass of OutputReference that points back to its owning Node and a variable name.

VoidPromise is returned by tasks with no declared outputs. It disables comparison, arithmetic, and truth testing with AssertionError("Task ... returns nothing, NoneType return cannot be used"), but it still supports >> and with_overrides by delegating to the underlying node. This lets you chain side-effect-only tasks without flytekit silently treating them as having a value.

Promise also supports attribute and item access (o.a.b or o["a"][0]) by appending to _attr_path and returning a new Promise whose NodeOutput shares the attribute path. This is how flytekit threads nested outputs (e.g., fields of a dataclass or elements of a dict) between tasks without resolving them at compile time.