Skip to main content

Workflow composition and nodes

Compile a workflow into a node graph​

A Python function decorated with @workflow does not become a backend program that re-runs its Python body for every execution. Flytekit evaluates the body while it compiles or serializes the workflow, using Promise objects in place of workflow inputs and task outputs. That evaluation discovers the nodes, their bindings, and the workflow outputs; the resulting PythonFunctionWorkflow is what Flytekit can serialize as an executable workflow definition. The decorator's docstring in flytekit/core/workflow.py explicitly cautions that non-Flyte code in the body runs only once for the platform (local execution is a separate path).

Keep ordinary Python computation inside a @task when it must happen as part of execution. Within a workflow body, a task result is a reference, not a native value:

from flytekit import task, workflow

@task(cache=True, cache_version="1", retries=3)
def sum(x: int, y: int) -> int:
    return x + y

@task(cache=True, cache_version="1", retries=3)
def square(z: int) -> int:
    return z * z

@workflow
def my_workflow(x: int, y: int) -> int:
    return sum(x=square(z=x), y=square(z=y))

This is the data-flow form used in README.md. square(z=x) and square(z=y) produce promises. Passing those promises to sum creates bindings from sum's inputs to the two upstream node outputs, so the data dependencies also determine execution order.

The workflow object and compilation​

from flytekit import workflow is the public decorator. With no arguments, @workflow uses WorkflowFailurePolicy.FAIL_IMMEDIATELY and interruptible=False. The decorator also accepts the options implemented in flytekit/core/workflow.py:

@workflow(
    failure_policy=WorkflowFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE,
    interruptible=True,
    on_failure=clean_up,
    pickle_untyped=False,
    docs=workflow_docs,
    default_options=options,
)
def wf(a: int) -> str:
    return some_task(a=a)

failure_policy must be one of the two WorkflowFailurePolicy values. WorkflowMetadata converts that policy into the workflow model, while WorkflowMetadataDefaults stores the interruptible default passed down to tasks. docs, pickle_untyped, and default_options are forwarded when the workflow object is constructed. pickle_untyped=True is the decorator's opt-in path for treating unannotated values as pickled Any values.

The decorator creates a PythonFunctionWorkflow, which inherits from WorkflowBase and ClassStorageTaskResolver. It retains the original function as function and stores the function name in _lhs. The common WorkflowBase state includes _nodes, _output_bindings, the typed interface, input promises, and optional failure-node information. Accessing wf.nodes or wf.output_bindings calls compile() lazily.

PythonFunctionWorkflow.compile() is guarded by self.compiled: it returns immediately after the first compilation. Otherwise it creates a fresh CompilationState, constructs input promises with construct_input_promises(), invokes the original function with those promises, and collects the nodes from the compilation state. It then converts the returned promises into output Binding objects and stores them in _output_bindings alongside _nodes. Consequently, changing a workflow function or its dependencies after compilation is not picked up unless you explicitly reset wf.compiled; the source marks that behavior as something not recommended.

At call time, WorkflowBase.__call__() fills in default inputs, compiles the workflow when appropriate, and delegates to flyte_entity_call_handler() in flytekit/core/promise.py. The call handler has separate compilation and local-execution paths. In compilation mode it calls create_and_link_node() and returns promises; during local workflow execution it invokes local_execute() and wraps results so the same workflow body can be evaluated locally.

A sub-workflow is composed the same way as a task. Its call becomes a workflow node in the containing DAG, and named outputs remain addressable by name. The tests in test_workflows.py use this pattern:

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

@workflow
def simple_wf() -> int:
    return add_5(a=1)

@workflow
def my_wf_example(a: int) -> typing.Tuple[int, int]:
    x = add_5(a=a)
    z = add_5(a=x)
    d = simple_wf()
    e = conditional("bool").if_(a == 5).then(add_5(a=d)).else_().then(add_5(a=z))
    return x, e

The x to z binding makes the second add_5 depend on the first. simple_wf() is represented as a sub-workflow node. Conditional construction is also part of the compiled structure; it is not an ordinary Python if over a task result.

How calls become Node objects​

The Node class in flytekit/core/node.py is the DAG-level record for one task, workflow, or launch-plan invocation. It stores:

  • an ID and NodeMetadata;
  • input Binding objects;
  • explicit upstream_nodes; and
  • the wrapped flyte_entity.

During compilation, flyte_entity_call_handler() detects ctx.compilation_state.mode == 1 and calls create_and_link_node(). That function in flytekit/core/promise.py resolves each input through binding_from_python_std(). A promise reference contributes both a binding and the referenced node to the dependency set. The global workflow inputs are represented by GLOBAL_START_NODE in workflow.py; references to that sentinel are excluded from upstream_nodes, because they are workflow inputs rather than executable predecessor nodes.

The node ID is generated from the active CompilationState:

f"{ctx.compilation_state.prefix}n{len(ctx.compilation_state.nodes)}"

Thus an ordinary compilation produces IDs such as n0, n1, and so on. create_and_link_node() constructs the Node, adds it to CompilationState.nodes, and creates Promise objects whose NodeOutput points back to that node. Those promises are what subsequent calls consume and what the workflow's final bindings reference.

Explicit ordering when there is no data flow​

Data dependencies are usually enough, but two side-effecting tasks may have no input or output to connect. Use create_node() from flytekit/core/node_creation.py when you need the Node itself:

from flytekit import task, workflow
from flytekit.core.node_creation import create_node

@task
def t2():
    ...

@task
def t3():
    ...

@workflow
def empty_wf():
    t2_node = create_node(t2)
    t3_node = create_node(t3)
    t3_node.runs_before(t2_node)

@workflow
def empty_wf2():
    t2_node = create_node(t2)
    t3_node = create_node(t3)
    t3_node >> t2_node

Node.runs_before() appends the left-hand node to the other node's _upstream_nodes if it is not already present. Node.__rshift__() calls runs_before() and returns the right-hand node, so t3_node >> t2_node is equivalent and can be chained. The unit tests verify that the serialized n0 node has n1 in its upstream IDs.

create_node() calls the entity during compilation, takes the newly added node from the compilation state, and attaches output promises to it. A node created this way exposes both named attributes and an outputs dictionary:

@task
def t1(a: int) -> typing.NamedTuple("OneOutput", [("t1_str_output", str)]):
    return (str(a + 2),)

@workflow
def my_wf(a: int, b: str) -> typing.Tuple[str, str]:
    t1_node = create_node(t1, a=a)
    t2_node = create_node(t2, a=[t1_node.t1_str_output, b])
    return t1_node.t1_str_output, t2_node.o0

The project tests use the same attribute form for named outputs and node.outputs["o0"] for default output names. Do not assume every Node has outputs: direct task calls return promises and their automatically created nodes do not initialize _outputs; Node.outputs raises an AssertionError unless the node came through create_node().

Only keyword inputs are supported by create_node() and by the task/workflow composition path. Passing a positional argument raises FlyteAssertion with the message that only keyword arguments are supported.

Rename and configure a node​

Node.with_overrides() mutates the node and returns self, so it can be called directly on a create_node() result or on the promise returned by a task call. For example:

@workflow
def configured(a: str) -> str:
    return t1(a=a).with_overrides(
        node_name="stable-node",
        timeout=60,
        retries=2,
        interruptible=True,
        container_image="python:3.11",
    )

The complete public override surface in Node.with_overrides() includes:

  • Identity and output metadata: node_name, name, and aliases.
  • Resources and runtime placement: requests, limits, or the combined resources form; container_image; accelerator; shared_memory; and pod_template.
  • Execution behavior: timeout as integer seconds or datetime.timedelta, retries, interruptible, and task_config.
  • Caching: cache, plus the compatibility arguments cache_version, cache_serialize, and cache_ignore_input_vars.

Use either requests/limits or resources, not both. Node.with_overrides() rejects that combination, converts resource values into the resource model, and checks that resource values do not contain promises. A node_name is passed through _dnsify(), so the serialized ID is DNS-compliant; for example, an override containing an underscore is normalized rather than preserved verbatim.

For cache overrides, give a Cache object a version:

@workflow
def cached(a: str) -> str:
    return t1(a=a).with_overrides(
        cache=Cache(version="foo", serialize=True)
    )

Node._override_node_metadata() copies the cache version and serialization setting into node metadata. A Cache without version raises ValueError("must specify cache version when overriding"). The older cache_version and cache_serialize form is still handled for compatibility, while supplying deprecated cache arguments together with a Cache object is rejected.

Build the graph imperatively​

When a workflow must be assembled by an application rather than expressed as a Python function, use ImperativeWorkflow (also exported in the tests as Workflow). It is the programmatic analogue of a decorated workflow. Unlike PythonFunctionWorkflow, it maintains a persistent CompilationState and adds one entity at a time.

from flytekit.core.workflow import ImperativeWorkflow as 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"])

add_workflow_input() adds the input to the workflow interface and returns a promise whose NodeOutput points to GLOBAL_START_NODE. add_entity() enters the workflow's compilation context and delegates node creation to create_node(). The convenience methods add_task(), add_subwf(), and add_launch_plan() all delegate to add_entity().

Use a node output as the next entity's input to create a data binding explicitly:

wb = ImperativeWorkflow(name="my.workflow.a")
t1_node = wb.add_entity(t1, a=2)
t2_node = wb.add_entity(t2, a=t1_node.outputs["o0"])
wb.add_workflow_output("from_n0t2", t2_node.outputs["o0"])

add_workflow_output() creates the output binding and updates the workflow interface. For a single promise, Flytekit can infer its Python type; for a list or dictionary of promises, pass python_type explicitly.

Local imperative execution walks self.compilation_state.nodes in insertion order, resolves each node's bindings from previously collected outputs, invokes the entity, and then resolves the workflow output bindings. Therefore, add entities in topological order: add a producer before an entity that consumes its output. ready() is true only when the workflow has at least one node and every declared input has been consumed. The test fixture demonstrates that this imperative workflow is equivalent to:

nt = typing.NamedTuple("wf_output", [("from_n0t1", str)])

@workflow
def my_workflow(in1: str) -> nt:
    x = t1(a=in1)
    t2()
    return nt(x)

Failure handling and workflow defaults​

Attach a failure task with on_failure when the workflow should invoke cleanup after an execution error:

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

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

WorkflowBase.__call__() catches an exception, and when the configured handler declares an err input it supplies a FlyteError containing the failed node ID and exception message before re-raising the exception. During compilation, PythonFunctionWorkflow._validate_add_on_failure_handler() compiles the handler in a separate CompilationState prefix and requires exactly one compiled handler node.

The handler's inputs must be a superset of the workflow's inputs. Any handler input not present on the workflow must be optional; otherwise Flytekit raises FlyteFailureNodeInputMismatchException. For an imperative workflow, call add_on_failure_handler(entity). It performs the same interface validation, creates the handler node temporarily, removes it from the main node list, and stores it as the workflow's failure node.

Choose WorkflowFailurePolicy.FAIL_IMMEDIATELY when a node failure should put the workflow into a failed state immediately. Choose FAIL_AFTER_EXECUTABLE_NODES_COMPLETE when remaining runnable nodes should proceed first. The test in test_workflows.py combines the latter with interruptible=True and verifies that the serialized workflow metadata contains both settings.

Composition constraints to keep visible​

  • Promises are references. An output such as a = task_a(x=x) is not an int; using it as an ordinary Python value in expressions such as range(a) or a native if a > 5 does not evaluate the task. Use Flytekit's conditional construction for workflow branching.
  • Use keyword arguments. The composition helpers reject positional task/workflow arguments with FlyteAssertion.
  • Avoid mutable defaults. create_and_link_node() rejects list, dictionary, and set default values for inputs with FlyteAssertion.
  • Compilation is cached. PythonFunctionWorkflow.compile() runs once per workflow object unless its compiled flag is reset.
  • Use create_node() for node references. Direct calls expose promise outputs; create_node() is the path that exposes Node.outputs and enables explicit runs_before()/>> ordering.
  • Keep imperative nodes topological. ImperativeWorkflow.execute() iterates insertion order, so a consumer must be added after its producer.
  • Version cache overrides. Cache(version=...) is required when overriding node caching.

Together, WorkflowBase's input/output bindings, CompilationState, Node dependencies, and promise-valued outputs give Flytekit enough information to turn either a decorated PythonFunctionWorkflow or an ImperativeWorkflow into a serialized, executable workflow graph.