Workflow composition and nodes

Flyte workflows are the primary mechanism for composing tasks and other workflows into complex data pipelines. In flytekit, this composition is achieved by defining a Python function that describes the flow of data between entities.

Workflow Composition with Decorators

The most common way to define a workflow is using the @workflow decorator. When you decorate a function with @workflow, flytekit transforms it into a PythonFunctionWorkflow.

from flytekit import task, workflow

@task
def add_one(x: int) -> int:
    return x + 1

@workflow
def my_workflow(val: int) -> int:
    # Calling a task returns a Promise, not a native value
    result = add_one(x=val)
    return result

The Compilation Model

A critical distinction in flytekit is that the body of a @workflow function is evaluated at compile time (or serialization time), not during execution on the Flyte platform.

When PythonFunctionWorkflow.compile() is called, flytekit:

1. Constructs Promise objects for each workflow input using construct_input_promises. These promises are anchored to a GLOBAL_START_NODE.
2. Executes the function body once.
3. Intercepts calls to tasks and sub-workflows via flyte_entity_call_handler.
4. Records each call as a Node and tracks the data dependencies between them.

Because the body runs at compile time, you cannot use native Python logic that depends on the values of task outputs. For example, if result > 0: will fail because result is a Promise, not an integer. You must use Flyte-specific constructs like conditional for branching.

Nodes: The Building Blocks of Workflows

Every task call, sub-workflow call, or launch plan invocation within a workflow creates a Node. A Node (defined in flytekit.core.node) encapsulates the execution of a Flyte entity and its position in the workflow graph.

Node Creation and Linking

The engine behind node creation is create_and_link_node in flytekit.core.promise. When an entity is called during compilation:

1. Binding: It uses binding_from_python_std to create bindings for each input. If an input is a Promise from a previous task, flytekit identifies that task's node as an upstream dependency.
2. Node Construction: A new Node is instantiated with a unique ID (e.g., n0, n1).
3. Registration: The node is added to the current CompilationState.

Explicit Dependency Ordering

While most dependencies are inferred from data flow (passing a Promise from one task to another), you can enforce execution order for tasks that do not share data using the >> operator or the runs_before method.

@workflow
def ordering_wf():
    node_a = create_node(task_a)
    node_b = create_node(task_b)
    
    # task_a will run before task_b
    node_a >> node_b 

Per-Node Overrides

You can customize the execution parameters of a specific node without changing the underlying task definition. This is done using the .with_overrides() method, which is available on Promise, VoidPromise, and Node objects.

Node.with_overrides supports several parameters:

• Resources: requests and limits using flytekit.Resources. Note that you cannot provide a resources object alongside individual requests or limits.
• Retries: Number of retries for the specific node.
• Timeout: A datetime.timedelta or integer seconds.
• Caching: Enable or disable caching for this specific invocation using a Cache object.
• Container Image: Override the image used for this node.

@workflow
def override_wf(val: int) -> int:
    return add_one(x=val).with_overrides(
        retries=3,
        requests=Resources(cpu="2", mem="500Mi"),
        node_name="custom-node-name"
    )

Imperative Workflows

For scenarios where workflows need to be constructed programmatically (e.g., based on a configuration file), flytekit provides the ImperativeWorkflow class (often aliased as Workflow).

from flytekit import Workflow

wb = Workflow(name="programmatic_workflow")
# Define inputs
wf_input = wb.add_workflow_input("val", int)
# Add entities
node = wb.add_entity(add_one, x=wf_input)
# Define outputs
wb.add_workflow_output("final_result", node.outputs["o0"])

In an ImperativeWorkflow, you manually manage the CompilationState by calling add_entity, add_workflow_input, and add_workflow_output.

Failure Handling

Flyte allows you to define behavior for when a workflow node fails.

Failure Policies

The workflow() decorator accepts a failure_policy (from WorkflowFailurePolicy):

• FAIL_IMMEDIATELY: The workflow stops as soon as any node fails.
• FAIL_AFTER_EXECUTABLE_NODES_COMPLETE: The workflow continues executing nodes that do not depend on the failed node.

On-Failure Handlers

You can specify an on_failure handler—a task or workflow that runs if the workflow fails. The handler's signature must accept the workflow's inputs and can optionally accept an err parameter of type FlyteError.

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

@workflow(on_failure=cleanup)
def wf_with_handler(val: int) -> int:
    return add_one(x=val)

Internally, flytekit validates that the on_failure handler's inputs are a superset of the workflow's inputs and assigns it a special node ID fn0.

Implementation Details and Constraints

• DNS Compliance: Node IDs are "DNSified" via _dnsify to ensure they are valid Kubernetes resource names.
• Promise Restrictions: Promise objects do not support standard Python truthiness (if promise:) or iteration (for x in promise:). Use bitwise operators & and | for logical operations if supported by the underlying type.
• Return Values: Workflows must return Promise objects originating from nodes. Returning local constants or native Python variables will result in an AssertionError during compilation.
• Tuples: Flyte uses tuples to represent multiple outputs. If a task returns a single NamedTuple, flytekit preserves the attribute names, allowing you to access them via node.outputs["attr_name"].