Workflow composition and nodes

Flyte workflows are declarative entities that define a directed acyclic graph (DAG) of tasks and sub-workflows. In flytekit, you define these graphs using the @workflow decorator on a Python function.

When you call a task inside a workflow, flytekit does not execute the task immediately. Instead, it creates a Node in the workflow's execution graph and returns a Promise object. These promises represent future values and are used to connect the outputs of one node to the inputs of another, establishing data dependencies.

Defining Workflows

To create a workflow, apply the @workflow decorator to a Python function. The function's signature defines the workflow's interface (inputs and outputs).

from flytekit import workflow, task

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

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

Internally, the workflow decorator (found in flytekit/core/workflow.py) evaluates the function body at serialization time to construct the DAG. Unlike tasks, the body of a workflow function is typically run only once during compilation/registration, not during every execution on the Flyte platform.

Workflow Metadata and Policies

You can configure how the workflow behaves during failures using parameters in the @workflow decorator:

• failure_policy: Determines what happens when a node fails. WorkflowFailurePolicy.FAIL_IMMEDIATELY (default) stops the workflow as soon as any node fails. WorkflowFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE allows other independent nodes to finish before failing the workflow.
• interruptible: A boolean indicating if the tasks within the workflow can be scheduled on interruptible (spot) instances.
• on_failure: A task or workflow to execute if the workflow fails. This is useful for cleanup operations.

from flytekit.core.workflow import WorkflowFailurePolicy

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

Nodes and Execution Order

Every time you call a task or sub-workflow within a workflow function, flytekit creates a Node (defined in flytekit/core/node.py). A Node encapsulates the execution entity (task or workflow), its inputs (bindings), and metadata.

Data Dependencies

Data dependencies are created naturally by passing the output Promise of one task as an input to another.

@workflow
def chained_workflow(val: int) -> int:
    a = add_one(x=val)
    b = add_one(x=a)  # 'b' depends on 'a'
    return b

Control Flow Dependencies

Sometimes you need to ensure a specific execution order even when there is no data dependency (e.g., a setup task must run before a processing task). You can use the >> operator to define these dependencies between nodes.

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

@workflow
def ordered_workflow(val: int) -> int:
    s = setup()
    a = add_one(x=val)
    s >> a  # Ensures setup runs before add_one
    return a

The Node.__rshift__ method implements this behavior by adding the left-hand node to the _upstream_nodes list of the right-hand node.

Customizing Nodes with Overrides

You can customize the execution parameters of a specific node using the with_overrides method. This is useful for setting resource requirements, retries, or timeouts for a single task invocation without changing the task definition itself.

from flytekit import Resources

@workflow
def override_workflow(val: int) -> int:
    # Request specific CPU and memory for this specific call
    a = add_one(x=val).with_overrides(
        requests=Resources(cpu="2", mem="4Gi"),
        retries=3,
        node_name="custom-add-one-node"
    )
    return a

The Node.with_overrides method (in flytekit/core/node.py) allows you to modify:

• requests and limits: Resource specifications using flytekit.Resources.
• timeout: Execution timeout as an int (seconds) or datetime.timedelta.
• retries: Number of times to retry the node on failure.
• interruptible: Override the workflow-level interruptible setting for this node.
• container_image: Use a specific Docker image for this node.

Promises and Local Execution

In flytekit, the return value of a task call inside a workflow is a Promise. A Promise (defined in flytekit/core/promise.py) acts as a proxy for the actual value.

When running a workflow locally, flytekit manages the transition between these states:

1. Compilation State: Task calls return Promise objects containing a NodeOutput reference.
2. Local Execution State: The WorkflowBase.local_execute method translates inputs into literals, executes tasks in topological order, and resolves Promise objects into actual values using translate_inputs_to_literals.

You can also access specific attributes or indices of a Promise if the underlying task returns a complex type like a NamedTuple, dict, or list:

@task
def get_map() -> dict:
    return {"a": 1, "b": 2}

@workflow
def map_workflow() -> int:
    m = get_map()
    return m["a"]  # Indexing into a Promise

This indexing creates a new Promise with an attr_path that flytekit resolves during execution.