Workflow composition and nodes

Flyte workflows are defined using the @workflow decorator. While they look like standard Python functions, they are executed at compile time to construct a Directed Acyclic Graph (DAG) of execution units called Nodes.

Defining Workflows and Nodes

When you call a @task or another @workflow inside a function decorated with @workflow, flytekit does not execute the task immediately. Instead, it creates a Node (defined in flytekit.core.node.Node) and returns a Promise object.

A Node represents a single execution step in the DAG. It encapsulates the Flyte entity being run (a task, sub-workflow, or launch plan), its metadata, and its input bindings.

from flytekit import task, workflow

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

@workflow
def my_workflow(val: int) -> int:
    # Calling add_one creates a Node internally.
    # 'result' is a Promise, not the actual integer.
    result = add_one(x=val)
    return result

Connecting Nodes via Data Dependencies

Data dependencies are the primary way to define the structure of your DAG. When you pass the output of one task (a Promise) as an input to another task, flytekit automatically creates an edge between the corresponding nodes.

In the following example from workflow.py, the node for z depends on the output of the node for x:

@workflow
def my_wf_example(a: int) -> typing.Tuple[int, int]:
    # Node 1
    x = add_5(a=a)

    # Node 2: Depends on Node 1 because it uses 'x'
    z = add_5(a=x)

    return x, z

Explicit Execution Order

Sometimes you need to ensure one task runs before another even if there is no data being passed between them (e.g., a task that writes to a database followed by a task that reads from it). You can use the >> operator to define an explicit dependency.

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

@workflow
def ordered_workflow():
    node_a = task_a()
    node_b = task_b()
    
    # node_a must complete before node_b starts
    node_a >> node_b

Customizing Node Execution with Overrides

The Node class provides a with_overrides method that allows you to customize the execution parameters of a specific task instance within a workflow without changing the task definition itself. This is useful for adjusting resources, retries, or timeouts for specific steps.

Common overrides include:

• Resources: Set requests and limits using flytekit.Resources.
• Retries: Set the number of retries.
• Timeout: Set a timeout as an integer (seconds) or datetime.timedelta.
• Node Name: Change the display name of the node in the Flyte UI using node_name.

from flytekit import Resources

@workflow
def resource_heavy_wf(x: int) -> str:
    return my_mappable_task(a=x).with_overrides(
        node_name="heavy-cpu-task",
        requests=Resources(cpu="10M"),
        retries=3,
        timeout=600
    )

Internally, with_overrides modifies the NodeMetadata and _resources attributes of the Node instance. For example, it uses _dnsify to ensure the node_name is a valid Kubernetes DNS subdomain name.

Caching Overrides

When overriding cache settings, flytekit prefers the use of a Cache object over deprecated individual parameters like cache_serialize.

from flytekit.core.node import Cache

# Overriding a node to enable caching with a specific version
node.with_overrides(cache=Cache(version="v2", serialize=True))

Important Constraints

1. Compile-time Execution: The body of a @workflow function runs only once during registration/compilation. You cannot use Python control flow (like if x > 10:) based on task outputs because those outputs are Promise objects, not actual values. Use the conditional construct for logic based on task results.
2. Promises: Variables inside a workflow are Promise objects. If you try to perform Python operations on them (like x + 1 where x is a task output), it will fail unless the operation is supported by the Promise class.
3. Sub-workflows: You can call other workflows within a workflow. Flytekit treats the sub-workflow as a single Node in the parent DAG.