Workflow composition and nodes

Workflows in flytekit are the primary mechanism for composing tasks into complex, multi-step pipelines. They are defined as declarative Python functions that describe how data flows between tasks. Internally, flytekit translates these functions into a Directed Acyclic Graph (DAG) where each step is represented by a Node.

Defining Workflows

You create a workflow by decorating a Python function with the @workflow decorator. Unlike tasks, the body of a workflow function is executed at serialization time (compile-time) to build the execution graph.

from flytekit import task, workflow

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

@workflow
def my_pipeline(val: int) -> int:
    result = add_one(x=val)
    return result

When you call add_one(x=val) inside the workflow, flytekit does not execute the task immediately. Instead, it creates a Node and returns a Promise object. This promise represents a future value that will be available when the task completes during actual execution.

Workflow Metadata and Policies

The @workflow decorator accepts several parameters to control the behavior of the entire pipeline:

• failure_policy: Determines what happens when a node fails. WorkflowFailurePolicy.FAIL_IMMEDIATELY (default) stops the workflow, while FAIL_AFTER_EXECUTABLE_NODES_COMPLETE allows other independent nodes to finish.
• interruptible: A boolean indicating if the workflow's tasks should be run on interruptible (spot) instances by default.
• on_failure: A task or workflow to execute if this workflow fails.

from flytekit import workflow, WorkflowFailurePolicy

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

Understanding Nodes

A Node (defined in flytekit.core.node.Node) is the fundamental unit of a workflow graph. It encapsulates a Flyte entity (like a task or a sub-workflow) along with its metadata, inputs, and upstream dependencies.

Data Dependencies

The most common way to connect nodes is through data flow. When you pass the output of one task as an input to another, flytekit automatically creates a dependency between the corresponding nodes.

@workflow
def data_flow_wf(a: int) -> int:
    # Node 1 is created for t1
    res1 = t1(a=a)
    # Node 2 is created for t2, with an upstream dependency on Node 1
    res2 = t2(b=res1)
    return res2

Explicit Execution Order

Sometimes you need to ensure a task runs after another even if there is no data being passed between them (e.g., a task that performs a side effect like cleaning up a database). You can use the >> operator to define an explicit execution order.

@workflow
def explicit_order_wf(name: str):
    c = create_cluster(name=name)
    t = run_analysis(name=name)
    d = delete_cluster(name=name)
    
    # Ensure analysis runs after cluster creation and deletion runs after analysis
    c >> t >> d

Internally, the __rshift__ method in the Node class calls runs_before(other), which appends the current node to the _upstream_nodes list of the target node.

Node Overrides

You can customize the execution environment of a specific node without changing the underlying task definition using the .with_overrides() method. This is useful for adjusting resources or retries for a specific step in a workflow.

from flytekit import Resources

@workflow
def override_wf(val: int) -> int:
    # Task t1 is configured with specific resource limits and retries for this node
    node = t1(a=val).with_overrides(
        requests=Resources(cpu="2", mem="4Gi"),
        retries=3,
        node_name="heavy-lifting-node"
    )
    return node.outputs["o0"]

The Node.with_overrides method allows you to modify:

• requests and limits: CPU, memory, and GPU requirements using flytekit.Resources.
• timeout: A datetime.timedelta or integer seconds.
• retries: Number of times to retry the node on failure.
• container_image: Use a specific Docker image for this node.
• node_name: Provide a custom ID for the node in the Flyte UI.

Handling Failures

Workflows can define an on_failure handler to perform cleanup or notification tasks. If the workflow fails, Flyte will execute the specified entity.

@task
def clean_up(err: str):
    print(f"Workflow failed with error: {err}")

@workflow(on_failure=clean_up)
def failure_handling_wf(name: str):
    t1(name=name)

When an on_failure task is triggered, flytekit can automatically pass the error message if the handler task accepts an input named err of type str or FlyteError. This is managed in WorkflowBase.__call__, which catches exceptions and invokes the on_failure entity with the appropriate error context.