Workflow composition and nodes

Flyte workflows are directed acyclic graphs (DAGs) where each node represents an executable unit, such as a task or a sub-workflow. In flytekit, you define these graphs by composing functions and entities, which the framework then translates into a structured execution plan.

Defining Workflows with Decorators

The most common way to compose a workflow is using the @workflow decorator. When you decorate a Python function with @workflow, flytekit executes the function body at compile time to build the DAG.

from flytekit import task, workflow

@task
def add(a: int, b: int) -> int:
    return a + b

@task
def square(c: int) -> int:
    return c * c

@workflow
def math_workflow(x: int, y: int) -> int:
    sum_result = add(a=x, b=y)
    return square(c=sum_result)

In this example, flytekit tracks the flow of data from add to square. Because square requires the output of add, flytekit automatically creates a dependency between the two nodes.

Workflow Metadata and Policies

You can configure how the workflow behaves during execution by passing arguments to the decorator. For instance, the failure_policy determines if the workflow should stop immediately upon a node failure or continue running other independent nodes.

from flytekit import workflow, WorkflowFailurePolicy

@workflow(
    interruptible=True, 
    failure_policy=WorkflowFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE
)
def robust_workflow(a: int) -> int:
    # ... workflow body ...
    pass

The WorkflowFailurePolicy enum in flytekit/core/workflow.py defines two modes:

• FAIL_IMMEDIATELY: The default behavior; the workflow fails as soon as any node fails.
• FAIL_AFTER_EXECUTABLE_NODES_COMPLETE: The workflow continues executing nodes that do not depend on the failed node.

Understanding Nodes

A Node is the fundamental building block of a Flyte DAG. While the @workflow decorator handles node creation implicitly, the flytekit.core.node.Node class manages the underlying metadata, bindings, and dependencies.

Every time you call a task inside a workflow, flytekit creates a node. This node encapsulates:

• Inputs: Bindings that map workflow inputs or upstream node outputs to the task's arguments.
• Metadata: Configuration like timeouts, retries, and interruptibility.
• Upstream Dependencies: A list of nodes that must complete before this node can start.

Explicit Dependencies with the Shift Operator

Sometimes you need to enforce an execution order even when there is no data dependency between tasks (e.g., a task that performs a side effect like cleaning up a database). You can use the >> operator or the runs_before method to define these explicit dependencies.

from flytekit import task, workflow, create_node

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

@task
def work():
    print("Doing work...")

@workflow
def explicit_dependency_wf():
    setup_node = create_node(setup)
    work_node = create_node(work)
    
    # Ensure setup runs before work
    setup_node >> work_node

Internally, the Node.__rshift__ method in flytekit/core/node.py calls self.runs_before(other), which appends the current node to the _upstream_nodes list of the target node.

Customizing Node Execution

You can override the default configuration of a specific node using the with_overrides method. This is useful for tailoring resource requirements or execution constraints for a single step in a workflow without modifying the underlying task definition.

from flytekit import task, workflow, Resources

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

@workflow
def resource_wf(data: list) -> int:
    return memory_intensive_task(data=data).with_overrides(
        requests=Resources(cpu="2", mem="500Mi"),
        limits=Resources(cpu="4", mem="1Gi"),
        retries=3,
        timeout=600 # seconds
    )

The Node.with_overrides method allows you to modify:

• Resources: CPU, memory, and GPU requests/limits via the Resources class.
• Retries: The number of times Flyte should attempt to re-run the node on failure.
• Timeouts: A datetime.timedelta or integer seconds after which the node is terminated.
• Container Image: Override the default image for just this node.

Imperative Workflows

While the @workflow decorator is standard, ImperativeWorkflow provides a programmatic way to construct workflows. This is particularly useful when the workflow structure is dynamic or generated based on external configuration.

from flytekit.core.workflow import ImperativeWorkflow
from my_tasks import t1, t2

# 1. Initialize the workflow
wb = ImperativeWorkflow(name="dynamic_workflow")

# 2. Add inputs
wb.add_workflow_input("in1", str)

# 3. Add entities (tasks/workflows) and capture nodes
node_1 = wb.add_entity(t1, a=wb.inputs["in1"])
node_2 = wb.add_entity(t2)

# 4. Define outputs
wb.add_workflow_output("final_result", node_1.outputs["o0"])

Unlike the functional approach, ImperativeWorkflow requires you to explicitly add inputs and outputs. The add_entity method returns a Node object, which you then use to link data or define execution order. Internally, ImperativeWorkflow maintains a CompilationState to track the nodes as they are added, ensuring they are processed in a topological order during local execution.