Skip to main content

Workflow composition, failure handlers, and nodes

Workflow composition

When you decorate a function with @workflow, flytekit evaluates the function body at compile time to construct a directed acyclic graph (DAG) of tasks. During this evaluation, task calls do not execute — they return Promise objects that act as references to future node outputs.

Two paradigms for building workflows

Declarative (function-based) workflows use the @workflow decorator. The function body is evaluated once during compilation to discover the DAG structure.

from flytekit import workflow, task

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

@task
def t2(a: int) -> str:
return f"result: {a}"

@workflow
def my_wf(a: int) -> str:
x = t1(a=a) # Returns a Promise, not an int
return t2(a=x) # x is a Promise pointing to node n0's output

Imperative workflows use the ImperativeWorkflow class for programmatic construction, useful when the DAG structure depends on dynamic logic.

from flytekit.core.workflow import ImperativeWorkflow

wb = ImperativeWorkflow(name="my_workflow")
wb.add_workflow_input("in1", str)
node = wb.add_entity(t1, a=wb.inputs["in1"])
wb.add_entity(t2)
wb.add_workflow_output("from_n0t1", node.outputs["o0"])

The ImperativeWorkflow.add_entity() method calls create_node() internally (defined in flytekit/core/node_creation.py), which returns a Node object during compilation.

Task outputs: Promise vs Node outputs

When you call a task inside a workflow using the regular function-call syntax, flytekit returns Promise objects via flyte_entity_call_handler() in flytekit/core/promise.py. The Promise class wraps a NodeOutput that references a node and output variable name.

x, y = t1(a=5)  # Returns (Promise, Promise) during compilation

The Promise class has two key states:

  • is_ready = True: The promise holds an actual resolved Literal value (local execution)
  • is_ready = False: The promise holds a NodeOutput reference to an upstream node
@property
def ref(self) -> NodeOutput:
"""
If the promise is NOT READY / Incomplete, then it maps to the origin node that owns the promise
"""
return self._ref

Accessing outputs via create_node() differs from regular task calls. The create_node() function in flytekit/core/node_creation.py returns a Node object, and the outputs are stored in node._outputs dict:

# During compilation, create_node returns Node with outputs attached:
t4_node = create_node(t4)
# Access via attribute:
t5(in1=t4_node.o0)
# Or via the outputs dict:
t5(in1=t4_node.outputs["o0"])

The distinction matters: calling t4() returns Promises directly, while create_node(t4) returns a Node with outputs stored as attributes.

Key difference: The Node.outputs property raises AssertionError if the node wasn't created via create_node(). Regular task calls don't produce nodes with this property accessible.

@property
def outputs(self):
if self._outputs is None:
raise AssertionError("Cannot use outputs with all Nodes, node must've been created from create_node()")
return self._outputs

Node ordering with runs_before and >>

When tasks don't produce/consume outputs, you need to manually specify dependencies. The Node class (in flytekit/core/node.py) provides two ways to establish ordering:

from flytekit.core.node_creation import create_node

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

# Two equivalent ways to make t2 run before t1:
t2_node.runs_before(t1_node)
# Or using the shift operator:
t2_node >> t1_node

The Node.__rshift__ method calls runs_before on the left operand:

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

Per-node overrides with with_overrides

The Node.with_overrides() method (defined in flytekit/core/node.py, lines 223-306) lets you set per-node execution settings. This method is also accessible via Promise.with_overrides() which delegates to the underlying node.

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

@workflow
def my_wf():
# Override timeout, retries, and interruptible for this specific node
slow_task().with_overrides(
timeout=3600, # 1 hour timeout
retries=3, # Retry up to 3 times
interruptible=True, # Allow preemption
name="custom_node_name" # Override the node ID
)

You can also override resources and caching:

from flytekit import Resources, Cache

slow_task().with_overrides(
requests=Resources(cpu="2", mem="4Gi"),
limits=Resources(cpu="4", mem="8Gi"),
cache=True,
cache_version="1.0"
)

When you call with_overrides() on a Promise, it checks if the promise is ready and delegates to the underlying node:

def with_overrides(self, *args, **kwargs):
if not self.is_ready:
self.ref.node.with_overrides(*args, **kwargs)
return self

Failure handlers with on_failure

When a workflow task fails, you can execute a cleanup handler by specifying the on_failure parameter to @workflow. The failure handler must accept all workflow inputs (they are passed automatically) plus optionally an err: Optional[FlyteError] parameter.

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

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

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

@task
def main_task(name: str):
print(f"Running main task on {name}")
raise ValueError("Something went wrong")

@workflow(on_failure=clean_up)
def wf(name: str = "flyteorg"):
c = create_cluster(name=name)
t = main_task(name=name)

Failure handler signature requirements

The validation logic in PythonFunctionWorkflow._validate_add_on_failure_handler() (lines 789-820) enforces two rules:

  1. Workflow inputs must be a subset of failure handler inputs:
if (failure_node_inputs | workflow_inputs) != failure_node_inputs:
raise FlyteFailureNodeInputMismatchException(self.on_failure, self)
  1. Additional inputs beyond workflow inputs must be Optional:
additional_keys = failure_node_inputs.keys() - workflow_inputs.keys()
for k in additional_keys:
if not is_optional_type(failure_node_inputs[k]):
raise FlyteFailureNodeInputMismatchException(self.on_failure, self)

This means a valid failure handler signature is:

# Valid: accepts all workflow inputs + optional err
@task
def handler(name: str, err: typing.Optional[FlyteError] = None):
...

# Valid: accepts all workflow inputs only
@task
def handler(name: str):
...

# Invalid: additional non-optional parameter
@task
def handler(name: str, code: int): # WRONG - code is not Optional
...

Imperative workflow failure handlers

For ImperativeWorkflow, use add_on_failure_handler():

wb = ImperativeWorkflow(name="my_workflow")
wb.add_workflow_input("in1", str)
wb.add_on_failure_handler(clean_up) # Pass the failure handler task
node = wb.add_entity(t1, a=wb.inputs["in1"])

The failure node is assigned the ID 'efn' (defined as DEFAULT_FAILURE_NODE_ID in flytekit/core/constants.py).

Promise attribute access

Promise objects support dot notation and bracket notation for accessing nested attributes:

import typing
@task
def t1() -> typing.NamedTuple("Output", [("a", int), ("b", str)]):
return (1, "hello")

@workflow
def wf():
o = t1()
# Access named outputs:
t2(x=o.a, y=o.b)
# Or bracket notation for indexing:
t3(items=o["a"][0])

The Promise.__getattr__() and Promise.__getitem__() methods append the attribute path and return a new promise:

def _append_attr(self, key) -> Promise:
new_promise = self.deepcopy()
new_promise._attr_path.append(key)
if new_promise.ref is not None:
new_promise._ref = new_promise.ref.with_attr(key)
return new_promise

What you cannot do with Promises

Promises are not Python values — they cannot be used in places that expect actual runtime values:

# This will raise ValueError:
if promise: # Cannot test truth value
...

# This will raise ValueError:
for x in promise: # Cannot iterate
...

# This will raise ValueError:
if a and b: # Cannot use 'and' operator
...

# Use bitwise operators instead:
if a.is_true() & b.is_true():
...

The Promise.__bool__() raises a ValueError to prevent accidental misuse:

def __bool__(self):
raise ValueError(
"Flytekit does not support Unary expressions or performing truth value testing,"
" This is a limitation in python. For Logical `and\\or` use `&\\|` (bitwise) instead"
)

Comparison and logical expressions

Promise objects support comparison operators that return ComparisonExpression objects:

@workflow
def wf(x: int):
# Comparisons:
is_large = x > 10
is_valid = (x >= 5) & (x <= 100)

The ComparisonExpression class (in flytekit/core/promise.py) handles primitive type comparisons, and ConjunctionExpression combines them with & (AND) and | (OR) operators.