Workflow composition and nodes

Flyte workflows are the primary mechanism for composing tasks into complex execution graphs. In flytekit, a workflow is defined as a Python function that describes how data flows between tasks. Internally, flytekit translates these function calls into a graph of nodes, where each node represents an execution step.

Composing Workflows with Decorators

The most common way to create a workflow is by using the @workflow decorator. When you call a task inside a decorated function, flytekit does not execute the task immediately. Instead, it records the call as a Node in the workflow's execution graph.

from flytekit import task, workflow

@task
def t1(a: int) -> int:
    return a + 1

@task
def t2(a: int) -> int:
    return a * 2

@workflow
def my_workflow(val: int) -> int:
    # Data dependency: t2 depends on the output of t1
    res1 = t1(a=val)
    res2 = t2(a=res1)
    return res2

In this example, flytekit automatically determines the execution order because t2 requires the output of t1. This output is wrapped in a Promise object, which represents a value that will be available at runtime.

Workflow Configuration

The @workflow decorator accepts parameters to control the overall behavior of the execution:

• interruptible: A boolean indicating if the workflow can be scheduled on lower-cost, interruptible instances.
• failure_policy: Uses the WorkflowFailurePolicy enum to decide whether to FAIL_IMMEDIATELY or FAIL_AFTER_EXECUTABLE_NODES_COMPLETE when a node fails.

from flytekit import workflow, WorkflowFailurePolicy

@workflow(interruptible=True, failure_policy=WorkflowFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE)
def robust_workflow(val: int) -> int:
    ...

Managing Execution Nodes

While data dependencies usually define the graph, you sometimes need to enforce an execution order between tasks that do not share data. Flytekit uses the Node class to represent these steps and provides tools to manipulate them.

Explicit Dependencies with create_node

If you need to run task_b after task_a, but task_b doesn't take any inputs from task_a, you can use create_node from flytekit.core.node_creation and the >> operator.

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

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

@task
def compute():
    print("Computing...")

@workflow
def ordered_workflow():
    setup_node = create_node(setup)
    compute_node = create_node(compute)
    
    # Enforce that setup runs before compute
    setup_node >> compute_node

The >> operator is a shorthand for the runs_before method on the Node class. Internally, setup_node >> compute_node calls setup_node.runs_before(compute_node), which appends setup_node to the _upstream_nodes list of compute_node.

Node Overrides

The Node class provides a with_overrides method that allows you to customize execution settings for a specific step without changing the task definition itself. This is useful for adjusting resources or retries for a single instance of a task.

from flytekit import Resources

@workflow
def override_workflow(val: int):
    node = create_node(t1, a=val)
    node.with_overrides(
        node_name="custom-t1-node",
        requests=Resources(cpu="2", mem="4Gi"),
        limits=Resources(cpu="4", mem="8Gi"),
        retries=3,
        timeout=3600  # seconds
    )

The with_overrides method modifies the NodeMetadata and resource specifications. It supports:

• Resources: Setting requests and limits using flytekit.Resources.
• Retries: Configuring a RetryStrategy.
• Timeouts: Accepting int (seconds) or datetime.timedelta.
• Caching: Overriding cache behavior via the cache parameter (accepts bool or a Cache object).

Imperative Workflows

For scenarios where the workflow structure is dynamic or generated at runtime, flytekit provides the ImperativeWorkflow class. This allows you to build a workflow programmatically by adding inputs, tasks, and outputs one by one.

from flytekit import ImperativeWorkflow

# Initialize the workflow
wb = ImperativeWorkflow(name="dynamic_workflow")

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

# Add tasks (entities)
node = wb.add_entity(t1, a=in1)

# Add outputs
wb.add_workflow_output("final_result", node.outputs["o0"])

When you call add_entity, flytekit creates a Node and manages the bindings (connections between inputs and outputs). During local execution, the ImperativeWorkflow.execute method iterates through these nodes in the order they were added, resolving Promise objects using an internal intermediate_node_outputs map.