Skip to main content

Workflow composition, failure handlers, and nodes

How @workflow builds a DAG

When you write a Flyte workflow, the function body is not executed for its side-effects at runtime—it is evaluated at compile time to construct a directed acyclic graph. The @workflow decorator in flytekit.core.workflow captures the tasks you call, the data flowing between them, and any explicit control logic, then serializes the result as a WorkflowTemplate.

You declare a workflow just like a Python function, but with Flyte-specific metadata attached:

from flytekit import workflow, task
from flytekit.core.workflow import WorkflowFailurePolicy
import typing

@task
def t1(a: int) -> typing.Tuple[str, str]:
return "one", "two"

@workflow(interruptible=True, failure_policy=WorkflowFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE)
def wf(a: int) -> typing.Tuple[str, str]:
x, y = t1(a=a)
_, v = t1(a=x)
return y, v

failure_policy controls what happens when a node fails. The default is WorkflowFailurePolicy.FAIL_IMMEDIATELY, which stops the whole workflow as soon as one node errors. Setting it to FAIL_AFTER_EXECUTABLE_NODES_COMPLETE lets other runnable branches finish first. interruptible tells Flyte that tasks launched from this workflow may be interrupted by the scheduler.

Task outputs are Promises, not values

Inside a workflow body, calling t1(a=a) does not return the real string tuple. During compilation, flytekit.core.python_function_task.PythonFunctionTask.__call__ delegates to create_and_link_node, which produces a flytekit.core.promise.Promise object that points to a future NodeOutput. The Promise class (defined in flytekit.core.promise) wraps either a NodeOutput (when compiling) or a concrete Literal (when running locally):

class Promise(object):
def __init__(
self,
var: str,
val: Union[NodeOutput, _literals_models.Literal],
type: typing.Optional[_type_models.LiteralType] = None,
):
self._var = var
self._promise_ready = True
self._val = val
self._ref = None
self._attr_path: List[Union[str, int]] = []
self._type = type
if val and isinstance(val, NodeOutput):
self._ref = val
self._promise_ready = False
self._val = None

If a task returns nothing, you get a VoidPromise instead. You cannot compare, index, or print a Promise during compilation—it exists only to be passed into downstream tasks or used for conditional expressions.

Explicit dependencies with create_node

Not every task produces an output that the next task consumes. If t1 and t2 both return None, how do you force t2 to finish before t1? Ordinary data flow won't help, because there is no data.

flytekit.core.node_creation.create_node exists for exactly this scenario. It creates a Node object (or a VoidPromise/tuple during local execution) and lets you wire explicit dependencies:

from flytekit.core.node_creation import create_node

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

# Two ways to say the same thing
t2_node.runs_before(t1_node)
# or
t2_node >> t1_node

create_node only accepts keyword arguments for inputs. If you pass positional arguments, it raises FlyteAssertion immediately:

def create_node(entity, *args, **kwargs):
if len(args) > 0:
raise _user_exceptions.FlyteAssertion(
f"Only keyword args are supported to pass inputs to workflows and tasks."
f"Aborting execution as detected {len(args)} positional args {args}"
)

During compilation, create_node invokes the entity, grabs the last node from the compilation state, and returns that Node. During local execution it returns the actual task result, so you should not rely on getting a Node object when running locally.

Accessing node outputs

Here is the distinction you must keep in mind:

  • A normal task call returns a Promise (or tuple of Promises).
  • create_node(...) returns a Node during compilation, and that Node exposes its outputs both as named attributes and through a dictionary.

For example, if t4() -> (int, str):

t4_node = create_node(t4)

# Access by attribute
t5(in1=t4_node.o0)

# Access by dictionary key
t5(in1=t4_node.outputs["o0"])

The Node.outputs property is only safe after create_node(). On an ordinary node that was not produced this way, it raises an explicit error:

@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

This design is heavily used by ImperativeWorkflow. When you build a workflow programmatically with flytekit.core.workflow.Workflow, add_entity delegates to create_node, and you bind workflow outputs via node.outputs:

from flytekit.core.workflow import Workflow

wb = Workflow(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"])

Internally, python_function_task.py uses the same pattern to wrap a task as a workflow:

node = wb.add_entity(self, **input_kwargs)
for output_name, output_python_type in self.python_interface.outputs.items():
wb.add_workflow_output(output_name, node.outputs[output_name])

Per-node overrides

Sometimes a single task in your workflow needs more memory, a different container image, or a longer timeout than its default definition. You can call with_overrides on either a Node or a Promise. If you call it on a Promise, it forwards the call to the underlying Node via self.ref.node.with_overrides(...):

# On a Promise returned by a normal task call
t1(a=1).with_overrides(requests=Resources(cpu="2", mem="1Gi"), timeout=300)

# On a Node returned by create_node
node = create_node(t2)
node.with_overrides(node_name="special-t2", retries=2, interruptible=True)

Node.with_overrides in flytekit.core.node accepts a long list of parameters:

def with_overrides(
self,
node_name: Optional[str] = None,
aliases: Optional[Dict[str, str]] = None,
requests: Optional[Resources] = None,
limits: Optional[Resources] = None,
timeout: Optional[Union[int, datetime.timedelta, object]] = TIMEOUT_OVERRIDE_SENTINEL,
retries: Optional[int] = None,
interruptible: Optional[bool] = None,
name: Optional[str] = None,
task_config: Optional[Any] = None,
container_image: Optional[str] = None,
accelerator: Optional[BaseAccelerator] = None,
cache: Optional[Union[bool, Cache]] = None,
shared_memory: Optional[Union[L[True], str]] = None,
pod_template: Optional[PodTemplate] = None,
resources: Optional[Resources] = None,
*args,
**kwargs,
):
pass

Be careful with resource arguments: resources is a newer consolidated parameter, and you cannot combine it with the older limits or requests parameters. Doing so raises ValueError:

if resources is not None:
if limits is not None or requests is not None:
msg = "`resource` can not be used together with the `limits` or `requests`. Please only set `resource`."
raise ValueError(msg)

Any node_name you provide is automatically DNSified via _dnsify so it is valid as a Kubernetes subdomain name.

Failure handlers (on_failure)

Workflows can declare a cleanup or notification task that runs if the workflow fails. You pass it to the @workflow decorator as on_failure. The handler has a strict signature contract: it must accept every workflow input, and any additional inputs beyond those must be Optional. If you violate this, PythonFunctionWorkflow.compile() raises FlyteFailureNodeInputMismatchException.

The canonical pattern is to accept all workflow inputs plus an optional err: typing.Optional[FlyteError] = None, which Flyte populates with the failure details:

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

@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 delete_cluster(name: str, err: typing.Optional[FlyteError] = None):
print(f"Deleting cluster {name}")

@task
def t1(a: int, b: str):
print(f"{a} {b}")
raise ValueError("something went wrong")

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

Internally, PythonFunctionWorkflow._validate_add_on_failure_handler builds a temporary CompilationState, invokes the failure handler with copies of the workflow input promises, and validates the interface:

workflow_inputs = self.python_interface.inputs
failure_node_inputs = self.on_failure.python_interface.inputs

# 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()
for k in additional_keys:
if not is_optional_type(failure_node_inputs[k]):
raise FlyteFailureNodeInputMismatchException(self.on_failure, self)

It also asserts that the failure handler compiles to exactly one node (len(inner_nodes) == 1), because only a single task or single workflow may serve as the handler. That node is then popped from the regular workflow node list and stored separately as _failure_node with the ID "efn" (DEFAULT_FAILURE_NODE_ID).

If you are building workflows imperatively, you can attach the same handler via ImperativeWorkflow.add_on_failure_handler, which performs the same validation and node-pop behavior:

def add_on_failure_handler(self, entity):
from flytekit.core.node_creation import create_node

ctx = FlyteContext.current_context()
if ctx.compilation_state is not None:
raise RuntimeError("Can't already be compiling")
with FlyteContextManager.with_context(ctx.with_compilation_state(self.compilation_state)) as ctx:
if entity.python_interface and self.python_interface:
workflow_inputs = self.python_interface.inputs
failure_node_inputs = entity.python_interface.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()
for k in additional_keys:
if not is_optional_type(failure_node_inputs[k]):
raise FlyteFailureNodeInputMismatchException(self.on_failure, self)

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
return n

In local execution, the failure handler is actually invoked when an exception is raised inside the workflow, making it useful for testing cleanup logic before you deploy.