Workflow composition and nodes

Flyte workflows are the primary mechanism for composing tasks and other workflows into a Directed Acyclic Graph (DAG). In flytekit, this composition happens at compile time when a function decorated with @workflow is evaluated.

Composing the DAG

When you decorate a Python function with @workflow, flytekit treats the function body as a DSL for building a graph. Instead of executing the logic immediately, flytekit records the calls to tasks and sub-workflows to construct the DAG.

Implicit Dependencies via Data Flow

The most common way to define dependencies is by passing the output of one task as an input to another. flytekit uses Promise objects to track these data dependencies.

from flytekit import task, workflow
import typing

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

@workflow
def my_workflow(a: int) -> typing.Tuple[int, str]:
    # x and y are Promises representing the future outputs of t1
    x, y = t1(a=a)
    
    # Passing x to another call creates an implicit dependency
    res_val, res_msg = t1(a=x)
    
    return res_val, y

In this example, the second call to t1 implicitly depends on the first because it consumes x. Internally, flytekit.core.promise.Promise objects are passed between tasks during compilation to build the flytekit.models.literals.Binding list for each node.

Explicit Dependencies with the Shift Operator

If two tasks do not share data but must run in a specific order (e.g., a setup task followed by a processing task), use the >> (right-shift) operator.

from flytekit import task, workflow, create_node

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

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

@workflow
def ordered_workflow(a: int) -> int:
    s = create_node(setup)
    c = compute(a=a)
    
    # setup must complete before compute starts
    s >> c
    
    return c

The Node.__rshift__ method in flytekit/core/node.py implements this by calling runs_before, which appends the upstream node to the _upstream_nodes list of the downstream node.

Node-Level Customization

Every time a task or workflow is called within a workflow, flytekit creates a Node. You can customize the execution behavior of these nodes using the .with_overrides() method.

Resource Requests and Limits

You can specify the hardware requirements for a specific node using the Resources class.

from flytekit import task, workflow, Resources

@task
def heavy_task(data: list) -> int:
    return len(data)

@workflow
def resource_wf(data: list) -> int:
    return heavy_task(data=data).with_overrides(
        requests=Resources(cpu="2", mem="1Gi"),
        limits=Resources(cpu="4", mem="2Gi")
    )

The Node.with_overrides method in flytekit/core/node.py validates these resources and converts them into a ResourceSpec model.

Retries and Timeouts

For tasks that might be flaky or have strict execution windows, you can set retries and timeouts at the call site.

import datetime

@workflow
def resilient_wf(a: int) -> int:
    # Set a 30-second timeout and 3 retry attempts
    return t1(a=a).with_overrides(
        timeout=datetime.timedelta(seconds=30),
        retries=3
    )

Internally, _override_node_metadata updates the NodeMetadata object associated with the node, which is eventually serialized into the workflow definition.

Caching

You can enable or override caching behavior for a specific node using the cache and cache_version parameters.

from flytekit import Cache

@workflow
def cached_wf(a: int) -> int:
    return t1(a=a).with_overrides(
        cache=True,
        cache_version="v1"
    )

Advanced Composition

Workflows can be composed of other workflows, allowing for modular design.

Sub-workflows

Calling a workflow from within another workflow creates a sub-workflow node.

@workflow
def sub_wf(a: int) -> int:
    return t1(a=a)[0]

@workflow
def parent_wf(a: int) -> int:
    # sub_wf is treated as a single node in parent_wf's DAG
    return sub_wf(a=a)

Imperative Workflows

While the @workflow decorator is the standard way to define DAGs, ImperativeWorkflow (found in flytekit/core/workflow.py) allows for programmatic construction of workflows, which is useful for dynamic graph generation.

from flytekit import ImperativeWorkflow

wb = ImperativeWorkflow(name="programmatic_wf")
wb.add_workflow_input("in1", int)
node = wb.add_entity(t1, a=wb.inputs["in1"])
wb.add_workflow_output("out1", node.outputs["val"])

Implementation Details

• Compile-time Evaluation: The body of a @workflow function is executed only once during compilation. Standard Python control flow (like if or for) inside the function will be evaluated at compile time, not at runtime on the cluster. For runtime branching, use flytekit.conditional.
• Node IDs: flytekit automatically generates unique IDs for nodes. If you use with_overrides(node_name="..."), the _dnsify function ensures the name is compliant with Kubernetes naming standards.
• Global Input Node: All workflow inputs are managed by a special GLOBAL_START_NODE (defined in flytekit/core/workflow.py) with the ID start-node. Promises for workflow inputs point back to this node.