Skip to main content

Workflow composition and nodes

Workflows build a graph at compile time​

When you write a Flyte workflow, the function body is not the code that Flyte runs for each execution. It is the code flytekit evaluates while compiling or serializing the workflow so it can discover the directed acyclic graph (DAG). The workflow decorator's docstring states that the body is evaluated at serialization time and is not evaluated again when the workflow runs on Flyte.

Use the decorator for the usual function-based form:

from flytekit import task, workflow

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

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

workflow is defined in flytekit/core/workflow.py. It accepts either the bare @workflow form or configuration such as interruptible, failure_policy, on_failure, docs, pickle_untyped, and default_options:

from flytekit.core.workflow import WorkflowFailurePolicy

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

The decorator creates a PythonFunctionWorkflow, a WorkflowBase implementation. PythonFunctionWorkflow.compile() compiles lazily: it creates input Promise stubs, invokes the wrapped function with those stubs, collects nodes from the active CompilationState, and converts the returned values into output bindings. It sets compiled before doing this work, so subsequent compilation requests do not rebuild the graph. The resulting _nodes and _output_bindings are what flytekit uses to produce the executable workflow definition.

A workflow body should therefore contain Flyte entities and graph-building expressions. Ordinary Python statements in the body execute during compilation. In particular, a task result in that context is normally a Promise, not the task's eventual Python value; operations such as using a task result as the argument to range() or len() cannot determine the runtime value while the graph is being built.

Composing tasks, workflows, and conditions​

Every task call made while the compilation context is active creates a node. You connect nodes most commonly by passing one node's output to another node's input:

@workflow
def my_wf_example(a: int) -> 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

This is the example in tests/flytekit/unit/core/test_workflows.py. x, z, and d are references to outputs produced by graph nodes during compilation. Passing x to add_5(a=x) creates the data-flow relationship; calling simple_wf() creates a sub-workflow node. The conditional expression creates a conditional section whose selected branch supplies e.

Workflow inputs are represented similarly. construct_input_promises() creates a Promise for each declared input, with a NodeOutput referring to the module-level GLOBAL_START_NODE. The serializer treats that start node as the workflow boundary rather than as an executable node. When compile() processes the function's return value, binding_from_python_std() creates bindings from those output promises, so workflow outputs must come from values returned by prior nodes.

For multiple outputs, return a tuple matching the declared interface. For named outputs, use a typing.NamedTuple return annotation; otherwise output names default to o0, o1, and so on. Named outputs can then be used as attributes on a node or its result. The imperative workflow test uses NamedTuple("wf_output", [("from_n0t1", str)]) so the workflow output has the name from_n0t1 instead of o0.

Nodes are the executable DAG units​

flytekit/core/node.py defines Node. A node stores the information needed to represent one execution step:

  • a DNS-compliant id (the constructor applies _dnsify()),
  • metadata,
  • input bindings,
  • explicit upstream_nodes, and
  • the wrapped flyte_entity (such as a task, workflow, or launch plan).

At the top level, node IDs are generated as n0, n1, and so forth by node creation. Nested compilation states add prefixes to avoid collisions. The node's bindings describe values supplied to the wrapped entity; its upstream_nodes describe ordering dependencies that are not necessarily expressed by an input value.

The central dispatch point is flyte_entity_call_handler() in flytekit/core/promise.py. In compilation mode it delegates to node creation and returns promises for the new node's outputs. In local workflow execution it instead invokes the entity's local execution path. This distinction lets the same workflow call produce a graph during serialization and concrete values during a local run.

Connect data flow and ordering​

Data-flow dependencies​

Prefer passing an output directly to the next entity when a value is consumed:

@workflow
def chained(a: int) -> int:
    first = add_5(a=a)
    second = add_5(a=first)
    return second

The output Promise carries a reference to the producing node. flytekit uses that reference when it builds the consumer's input binding and dependency information.

Explicit ordering for side-effect nodes​

Use runs_before() or >> when the ordering matters but no output is passed between the entities. Node.runs_before() appends the current node to the other node's _upstream_nodes if it is not already present. Node.__rshift__() calls that method and returns the right-hand node, which supports chaining:

@task
def t2():
    ...

@task
def t3():
    ...

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

The equivalent spelling is t3_node.runs_before(t2_node). The check in runs_before() makes repeated calls idempotent. Promise objects also expose the ordering behavior, so direct task-call results can be chained when the result is a single value:

@workflow
def ordered(x: int):
    task_a(x=x) >> task_b(x=x) >> task_c(x=x)

The explicit edge is serialized as an upstream-node relationship; it is separate from the value bindings used for data flow.

Use create_node() when you need a node handle​

A direct task call is convenient when you want its output promise. Use create_node() from flytekit/core/node_creation.py when you need a Node object—for example, to order void tasks or to address outputs by names held in strings:

@workflow
def my_wf(a: str) -> tuple[str, list[str]]:
    t1_node = create_node(t1, a=a)
    dyn_node = create_node(my_subwf, a=3)
    return t1_node.o0, dyn_node.o0

During compilation, create_node() calls the entity, obtains the newly registered node, and attaches its output promises both as attributes and in the node's outputs dictionary. Thus these two forms refer to the same output:

value = t1_node.o0
value_again = t1_node.outputs["o0"]

The dictionary form is useful for dynamic code, while the attribute form is convenient for statically named outputs. A task with no outputs returns the node itself, allowing ordering with runs_before() or >>:

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

create_node() accepts only keyword inputs. Passing positional arguments raises FlyteAssertion. It also has a local-execution path: while compiling it returns the graph node, while local execution calls the entity and returns the corresponding local result.

Customize one node with with_overrides()​

A task-call result or a node returned by create_node() can be customized for that invocation:

@workflow
def my_wf(a: str) -> str:
    s = t1(a=a).with_overrides(timeout=timeout)
    s1 = t1(a=s).with_overrides()
    s2 = t2(a=s1).with_overrides(timeout=timeout)
    s3 = t2(a=s2).with_overrides()
    return s3

Node.with_overrides() mutates the node and returns it, so the call can remain in the data-flow expression. The method supports node names, aliases, requests and limits (or a combined resources value), timeout, retries, interruptible, cache, task_config, container_image, accelerators, shared memory, and a pod_template.

Timeout handling has a deliberate distinction between omission and reset. Node.TIMEOUT_OVERRIDE_SENTINEL means “no timeout argument was supplied”; timeout=None explicitly sets an empty timeout. An integer is interpreted as seconds, and a datetime.timedelta is accepted directly. Calling .with_overrides() without a timeout therefore does not automatically erase an existing task timeout.

Names are converted through the same DNS-normalization path used by the constructor. For example, a node_name supplied to with_overrides() is converted to a DNS-compliant ID. Resource values must be static: node creation calls assert_no_promises_in_resources(), so a resource request or limit cannot depend on a task output. When using resources, do not also provide requests or limits; Node.with_overrides() rejects that combination.

For cache overrides, a Cache object must include a cache version. The method retains deprecated cache arguments for compatibility, but raises an error when deprecated version or serialization arguments are combined with a Cache object. The task_config override is marked beta in the source and must have the same type as the task's existing configuration.

Build workflows programmatically with Workflow​

Function-based workflows are not the only composition API. ImperativeWorkflow is exported as flytekit.Workflow and lets you construct the graph explicitly. This is useful when code must create entities dynamically; flytekit plugins use it for generated pipelines.

The basic API is Workflow(name=...), add_workflow_input(), add_entity(), and add_workflow_output():

from flytekit import task, Workflow

@task
def t1(a: str) -> str:
    return a + " world"

@task
def t2():
    print("side effect")

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"])

assert wb(in1="hello") == "hello world"

The equivalent decorator form from tests/flytekit/unit/core/test_imperative.py is:

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

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

ImperativeWorkflow also provides add_task(), add_launch_plan(), add_subwf(), and add_on_failure_handler() convenience methods. Its add_entity() result is a Node, so plugin code can chain values and apply node overrides:

node_1 = wf.add_entity(upload_jsonl_file_task_obj, jsonl_in=wf.inputs["jsonl_in"])
node_2 = wf.add_entity(batch_endpoint_task_obj, input_file_id=node_1.outputs["result"])
node_3 = wf.add_entity(download_json_files_task_obj, batch_endpoint_result=node_2.outputs["result"])

node_1.with_overrides(requests=Resources(mem=file_upload_mem), limits=Resources(mem=file_upload_mem))
node_3.with_overrides(requests=Resources(mem=file_download_mem), limits=Resources(mem=file_download_mem))
wf.add_workflow_output("batch_output", node_3.outputs["result"], BatchResult)

Imperative entities must be added in topological order: add a producer before adding an entity that consumes its output. ready() checks whether declared inputs are bound and whether the workflow can execute locally. Imperative local execution resolves bindings with the get_promise_map() / get_promise() helpers rather than relying on the Flyte execution engine.

Failure handling and failure policy​

Attach a cleanup task with on_failure:

@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 = t1(a=1, b="2")
    d = delete_cluster(name=name)
    c >> t >> d

The failure handler's inputs must cover the workflow's inputs; additional inputs must be optional. err: Optional[FlyteError] is the special optional input that receives failure information. The imperative equivalent is add_on_failure_handler(entity).

WorkflowFailurePolicy controls what happens after a node fails. FAIL_IMMEDIATELY is the default. FAIL_AFTER_EXECUTABLE_NODES_COMPLETE allows nodes that do not depend on the failed node to finish. Select it on the decorator, or pass the corresponding option when constructing an imperative workflow:

@workflow(failure_policy=WorkflowFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE)
def wf(a: int):
    return t1(a=a)

Referencing existing workflows​

Use the reference_workflow decorator when you need a typed stub for an existing Flyte workflow rather than a workflow function whose source is available. It takes project, domain, name, and version. ReferenceWorkflow represents that external entity. flytekit's serializer does not support using a reference workflow as a sub-workflow inside another workflow; use a reference launch plan for that composition case.

Composition checks to keep in mind

  • Treat the decorated function body as compile-time DAG construction. Non-Flyte Python code in it runs during compilation, not as a task step on Flyte.
  • Treat task results in the body as Promise objects. Pass them to Flyte entities or return them; do not use them as already-computed Python values.
  • Use typing.NamedTuple when output names matter. Without it, output names are generally o0, o1, and so on.
  • Call with_overrides() while the workflow is being compiled. It mutates the node created in that compilation context.
  • Use create_node() for explicit node handles, especially for void-task ordering; use ordinary task calls when direct promise-based data flow is sufficient.