Workflow composition and nodes

Flyte workflows are directed acyclic graphs (DAGs) where each step is represented as a node. While you typically define these graphs implicitly by passing outputs from one task to the inputs of another, flytekit provides fine-grained control over how these nodes are composed, ordered, and configured.

Workflow Composition with Decorators

The most common way to compose a workflow in flytekit is using the @workflow decorator. When you call a task inside a function decorated with @workflow, flytekit does not execute the task immediately. Instead, it captures the call as a Node and creates a Promise for each output.

from flytekit import task, workflow
import typing

@task
def t1(a: int) -> typing.NamedTuple("Outputs", [("val", int), ("msg", str)]):
    return a + 2, "result-" + str(a)

@workflow
def my_workflow(a: int) -> str:
    # Calling t1 creates a Node in the workflow graph.
    # x and y are Promises, not actual values.
    x, y = t1(a=a)
    
    # Passing a Promise to another task call creates a data dependency.
    _, v = t1(a=x)
    return v

Internally, the @workflow decorator transforms the Python function into a PythonFunctionWorkflow. During compilation, flytekit tracks the order of calls and the flow of Promise objects to build the underlying WorkflowTemplate.

Understanding Nodes

A Node (defined in flytekit.core.node.Node) is the fundamental unit of execution in a workflow. Every time you call a task, sub-workflow, or launch plan within a workflow, flytekit creates a node to represent that execution step.

Each Node tracks:

• Flyte Entity: The task or workflow to be executed (flyte_entity).
• Inputs: Bindings that map workflow inputs or previous node outputs to this node's inputs.
• Upstream Dependencies: A list of nodes that must complete before this node starts.
• Metadata: Configuration like retries, timeouts, and interruptibility.

Customizing Node Execution

You can override the default configuration of a specific task execution within a workflow using the with_overrides method. This is useful when a specific step in your workflow requires more memory, a longer timeout, or should be retried differently than the task's default definition.

from flytekit import task, workflow, Resources

@task
def heavy_task(x: int) -> int:
    return x * 2

@workflow
def resource_workflow(x: int) -> int:
    # Override resources and retries for this specific node
    return heavy_task(x=x).with_overrides(
        requests=Resources(cpu="2", mem="500Mi"),
        limits=Resources(cpu="4", mem="1Gi"),
        retries=3,
        node_name="my-custom-node-id"
    )

The with_overrides method (found in flytekit/core/node.py) modifies the node's metadata and resource requirements. Note that these overrides must be static values; flytekit does not support using Promise objects (outputs from other tasks) for override parameters like retries or timeout.

Explicit Dependencies and the Shift Operator

When two tasks do not share data but must run in a specific order, you cannot rely on implicit data-flow dependencies. In these cases, use create_node from flytekit.core.node_creation and the bitwise right-shift operator (>>).

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

@task
def setup():
    print("Setting up...")

@task
def work():
    print("Working...")

@workflow
def ordered_workflow():
    # Manually create nodes to gain access to dependency methods
    setup_node = create_node(setup)
    work_node = create_node(work)

    # Enforce that setup runs before work
    setup_node >> work_node

The >> operator is a shorthand for the runs_before method. Calling node_a >> node_b appends node_a to the _upstream_nodes list of node_b.

Programmatic Workflow Construction

For scenarios where the workflow structure is determined at runtime (e.g., based on a configuration file), you can use the Workflow class (often referred to as Imperative Workflows) instead of the @workflow decorator.

from flytekit import task, Workflow

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

# Construct the workflow programmatically
wb = Workflow(name="imperative_workflow")
wb.add_workflow_input("in1", str)

# add_entity returns a Node
node = wb.add_entity(t1, a=wb.inputs["in1"])

# Define the workflow output by referencing a node's output
wb.add_workflow_output("final_result", node.outputs["o0"])

In this mode, you explicitly manage the addition of inputs, entities (tasks/workflows), and outputs to the Workflow object, which internally manages the Node creation and binding process.