Workflow composition and nodes

Workflows in flytekit are the primary mechanism for composing tasks and other workflows into a directed acyclic graph (DAG). This composition is achieved by defining a function decorated with @workflow, where each call to a Flyte entity (like a task) creates a Node in the graph.

Defining Workflows

You define a workflow by decorating a Python function with the @workflow decorator. The function's parameters represent the workflow's inputs, and its return values define the workflow's outputs.

from flytekit import task, workflow

@task
def say_hello(name: str) -> str:
    return f"Hello, {name}!"

@workflow
def my_wf(name: str) -> str:
    greeting = say_hello(name=name)
    return greeting

Internally, the @workflow decorator transforms the function into a WorkflowBase object (found in flytekit/core/workflow.py). When this function is called during compilation, flytekit tracks the execution flow to build a static representation of the DAG.

Nodes and Data Dependencies

A Node (implemented in flytekit/core/node.py) represents a single execution step within a workflow. In most cases, you do not create nodes manually; they are created implicitly when you call a task or sub-workflow inside a workflow function.

Implicit Node Creation

When you call say_hello(name=name) inside my_wf, flytekit creates a Node for that task. The relationship between nodes is typically established through data flow: if the output of Node A is passed as an input to Node B, Node B implicitly depends on Node A.

Explicit Node Creation and Ordering

If you have tasks that do not share data but must run in a specific order, you can use the create_node function from flytekit/core/node_creation.py and the >> operator.

from flytekit import task, workflow, create_node

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

@task
def run_test():
    print("Running test...")

@workflow
def ordered_wf():
    setup_node = create_node(setup)
    test_node = create_node(run_test)
    
    # Enforce that setup runs before test
    setup_node >> test_node

The Node class implements __rshift__, which calls runs_before(other). This method appends the current node to the _upstream_nodes list of the target node, ensuring the Flyte engine respects this execution order.

Customizing Node Execution

You can override the default execution settings for a specific node using the .with_overrides() method. This is useful for setting resource requirements, retries, or timeouts for a single step without changing the task definition itself.

from flytekit import task, workflow, Resources

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

@workflow
def resource_wf(data: list) -> int:
    # Override CPU and Memory for this specific call
    return compute_heavy_task(data=data).with_overrides(
        requests=Resources(cpu="2", mem="500Mi"),
        limits=Resources(cpu="4", mem="1Gi"),
        retries=3
    )

The Node.with_overrides method in flytekit/core/node.py handles these configurations. It updates the NodeMetadata and internal _resources attributes, which are then used by the Flyte backend during execution.

Workflow-Level Configuration

The @workflow decorator accepts parameters that control the behavior of the entire workflow, such as how to handle failures or whether the workflow is interruptible.

from flytekit import workflow, WorkflowFailurePolicy

@workflow(
    failure_policy=WorkflowFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE,
    interruptible=True
)
def robust_wf(x: int) -> int:
    # ... workflow logic ...
    return x

• failure_policy: Determines if the workflow should stop immediately on any node failure (FAIL_IMMEDIATELY) or continue running other independent nodes (FAIL_AFTER_EXECUTABLE_NODES_COMPLETE).
• interruptible: When set to True, it indicates that the workflow's nodes can be scheduled on lower-priority, interruptible instances (like AWS Spot Instances).

Imperative Workflows

While the @workflow decorator provides a declarative way to define graphs, the ImperativeWorkflow class (in flytekit/core/workflow.py) allows you to build workflows programmatically. This is useful for dynamic scenarios where the graph structure depends on runtime logic.

from flytekit import ImperativeWorkflow, task

@task
def t1(a: int) -> int:
    return a + 1

# Create workflow programmatically
wb = ImperativeWorkflow(name="my_imperative_wf")
wb.add_workflow_input("in1", int)
node = wb.add_entity(t1, a=wb.inputs["in1"])
wb.add_workflow_output("out1", node.outputs["o0"])

The ImperativeWorkflow class maintains its own CompilationState and provides methods like add_entity to manually link tasks and inputs into the workflow structure.