Skip to main content

Workflow composition and nodes

Flyte workflows are defined as Python functions decorated with @workflow. When you call a task or another workflow inside this function, flytekit does not execute the task immediately. Instead, it records the call as a Node in a directed acyclic graph (DAG) and returns a Promise representing the future output of that node.

Nodes and Execution Steps​

A Node (defined in flytekit.core.node.Node) represents a single execution step in a workflow. Every time you invoke a task within a workflow, flytekit creates a Node instance to encapsulate that execution.

Implicit Node Creation​

In most cases, you don't interact with Node objects directly. They are created implicitly when you call a task:

from flytekit import task, workflow

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

@workflow
def my_workflow(val: int) -> int:
    # This call creates a Node internally
    result = add_one(x=val)
    return result

Internally, the Node stores:

  • An id: A unique identifier within the workflow (often derived from the task name).
  • flyte_entity: The task or sub-workflow to be executed.
  • bindings: How the inputs of this node are connected to outputs of previous nodes or workflow inputs.
  • upstream_nodes: A list of nodes that must complete before this node can start.

Explicit Ordering with >>​

Sometimes you need to ensure one task runs before another even if there is no data dependency between them. You can use the >> operator (implemented via Node.__rshift__) to enforce this order.

@workflow
def ordered_workflow():
    n1 = task_a()
    n2 = task_b()
    
    # Ensure task_a completes before task_b starts
    n1 >> n2

The >> operator calls Node.runs_before(other), which appends the current node to the _upstream_nodes list of the target node.

Promises and Data Flow​

When a task is called inside a workflow, it returns a Promise (defined in flytekit.core.promise.Promise). A Promise is a placeholder for a value that will exist only at runtime.

Connecting Inputs and Outputs​

Data flow is established by passing a Promise from one task as an argument to another. Flytekit uses these connections to build the edges of the DAG.

@workflow
def data_flow_wf(a: int) -> int:
    # res1 is a Promise
    res1 = task_1(a=a)
    # Passing res1 to task_2 creates a dependency
    res2 = task_2(b=res1)
    return res2

Accessing Collection Elements​

If a task returns a collection (like a list or dict), you can access its elements using standard Python indexing. This returns a new Promise with an updated attr_path.

@workflow
def collection_wf() -> str:
    # Assume t1 returns a dict: {"key": ["val1", "val2"]}
    o = t1()
    # Accessing by key and index creates a specialized Promise
    return o["key"][0]

The Promise.__getitem__ method appends the key or index to the _attr_path. During local execution, resolve_attr_path_in_promise uses this path to extract the specific value from the LiteralMap or LiteralCollection.

Customizing Node Execution​

You can override the default execution parameters for a specific node using the .with_overrides() method. This is available on both the Promise returned by a task call and the Node itself.

from flytekit import Resources

@workflow
def resource_wf(val: int) -> int:
    return add_one(x=val).with_overrides(
        requests=Resources(cpu="2", mem="500Mi"),
        limits=Resources(cpu="4", mem="1Gi"),
        retries=3,
        timeout=3600, # seconds
        node_name="custom-node-id"
    )

The Node.with_overrides method modifies the node's metadata and resource requirements:

  • Resources: Sets _resources using convert_resources_to_resource_model.
  • Metadata: Updates _metadata for timeouts, retries, and interruptibility via _override_node_metadata.
  • Naming: Updates self._id using _dnsify(node_name) to ensure Kubernetes compatibility.

Workflow Failure Handling​

Workflows can define a cleanup or recovery strategy using the on_failure parameter in the @workflow decorator.

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

@workflow(on_failure=clean_up)
def failure_handling_wf(val: int) -> int:
    return task_that_might_fail(val=val)

When a workflow is compiled, the on_failure entity is stored in WorkflowBase._on_failure. If the workflow execution fails, Flyte will trigger this entity. If the failure task accepts an err input, flytekit automatically passes the error message to it.

Imperative Workflow Composition​

While the @workflow decorator is the standard way to define workflows, flytekit also provides ImperativeWorkflow for cases where the DAG needs to be constructed programmatically (e.g., based on a configuration file).

from flytekit import ImperativeWorkflow

wb = ImperativeWorkflow(name="my_imperative_wf")
# Add inputs
in1 = wb.add_workflow_input("in1", int)
# Add nodes
node = wb.add_entity(add_one, x=in1)
# Add outputs
wb.add_workflow_output("out1", node.outputs["o0"])

In an ImperativeWorkflow, you explicitly manage the CompilationState and manually add entities and bindings, whereas the @workflow decorator handles this by tracing the function's execution.