Skip to main content

Workflow composition and nodes

When you call a Flyte task inside a function decorated with @workflow, flytekit does not execute the task immediately. Instead, it records the call as a Node in a Directed Acyclic Graph (DAG) and returns Promise objects that represent future values. This process, known as compilation, allows flytekit to build a complete execution plan that can be serialized and sent to the Flyte backend.

Workflow Compilation and Nodes​

A Flyte workflow is essentially a blueprint. When the @workflow function is evaluated, flytekit enters a compilation state where every interaction with a Flyte entity (task, subworkflow, or launch plan) is intercepted by the flyte_entity_call_handler in flytekit/core/promise.py.

How Nodes are Created​

The core of workflow composition is create_and_link_node. When you call a task like t1(a=x), flytekit performs the following steps:

  1. Interface Validation: It checks the inputs against the task's python_interface. If you miss an input or provide an extra one, it raises a FlyteAssertion.
  2. Binding Discovery: It inspects the input values. If an input is a Promise (an output from a previous task), flytekit identifies the Node that produced it.
  3. Node Instantiation: It creates a Node object (defined in flytekit/core/node.py). This node stores the underlying Flyte entity, the input bindings, and a list of upstream_nodes derived from the promises.
  4. ID Assignment: It assigns a unique ID to the node, typically following the pattern n0, n1, etc., based on the order of creation within the CompilationState.
  5. Promise Return: It returns new Promise objects that point to this specific node's outputs.
@workflow
def my_wf(a: int) -> int:
    # t1(a=a) creates Node 'n0'. It returns a Promise.
    x = t1(a=a) 
    # t2(b=x) creates Node 'n1'. It discovers 'n0' is upstream because of 'x'.
    return t2(b=x)

The Promise Gotcha​

Inside a workflow, variables like x in the example above are not native Python types (like int or str); they are flytekit.core.promise.Promise objects. You cannot use them in standard Python logic that requires concrete values:

  • NO: if x > 10: (Raises an error because a Promise cannot be evaluated as a boolean).
  • NO: range(x) (Raises a TypeError).
  • YES: Pass them as arguments to other tasks or return them as workflow outputs.

Explicit Ordering and Overrides​

While flytekit usually infers the graph structure from data dependencies, you can manually control node behavior and execution order.

Manual Dependencies​

If a task has side effects and no data dependency exists, use the >> operator or the runs_before method to enforce order. These operations are available on both Node objects and the Promise objects returned by tasks.

@workflow
def ordering_wf():
    a = setup_task()
    b = business_logic_task()
    c = cleanup_task()
    
    # Ensure setup runs before logic, and logic runs before cleanup
    a >> b >> c

Internally, node_a >> node_b calls node_a.runs_before(node_b), which appends node_a to the _upstream_nodes list of node_b.

Node Overrides​

You can customize the runtime requirements of a specific node using .with_overrides(). This is commonly used for resource allocation or retry policies.

from flytekit import Resources

@workflow
def resource_wf(n: int):
    # Apply overrides to the promise returned by the task
    x = heavy_task(n=n).with_overrides(
        requests=Resources(cpu="2", mem="500Mi"),
        retries=3,
        node_name="my-custom-heavy-node"
    )
    return x

The Node.with_overrides method in flytekit/core/node.py mutates the node's metadata. Note that you cannot use Promises for override values (e.g., you cannot pass a task output as the number of retries); these must be static values known at compile time.

Imperative Workflows​

For scenarios where the DAG structure is dynamic or generated programmatically, flytekit provides the ImperativeWorkflow class. This bypasses the decorator-based function scanning.

from flytekit import Workflow

# 1. Define the workflow container
wb = Workflow(name="programmatic_workflow")

# 2. Add inputs
in1 = wb.add_workflow_input("in1", int)

# 3. Add entities (tasks/subworkflows)
# add_entity returns a Node object
node_t1 = wb.add_entity(t1, a=in1)

# 4. Define outputs using node.outputs
wb.add_workflow_output("out1", node_t1.outputs["o0"])

In imperative workflows, node.outputs provides access to the promises generated by that node. This differs from function-based workflows where you typically use the return value of the task call.

Workflow Metadata and Failure Handling​

Workflows support high-level configuration via the @workflow decorator parameters, which are stored in WorkflowMetadata and WorkflowMetadataDefaults.

Failure Policies​

You can control how the engine reacts when a node fails using failure_policy:

  • WorkflowFailurePolicy.FAIL_IMMEDIATELY (Default): Stop the workflow as soon as any node fails.
  • WorkflowFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE: Continue executing other branches of the DAG that do not depend on the failed node.

On-Failure Handlers​

The on_failure parameter allows you to specify a task or subworkflow to run if the workflow fails. This is useful for cleanup operations.

@task
def cleanup(err: FlyteError):
    print(f"Workflow failed with error: {err.message}")

@workflow(on_failure=cleanup)
def my_wf():
    ...

Implementation Note: When flytekit executes the on_failure entity, it attempts to inject the failure details into an input named err. Ensure your cleanup task includes an input with this name if you need access to the error message.

Conditionals and Subworkflows​

Flytekit treats complex structures like conditionals and nested workflows as specialized nodes.

  • Subworkflows: When you call one @workflow inside another, the compiler creates a node where the flyte_entity is the subworkflow. The backend treats this as a single node that expands into its own DAG at runtime.
  • Conditionals: A conditional block results in a BranchNode. This node contains the logic for multiple branches. A critical requirement for compilation is that every conditional used as a workflow output must terminate with an .else_() branch to ensure a value is always returned.
@workflow
def conditional_wf(n: int) -> int:
    return (
        conditional("check_val")
        .if_(n > 0)
        .then(positive_task(n=n))
        .else_()
        .then(negative_task(n=n))
    )