Workflow composition and nodes

Flyte workflows in flytekit are defined by composing tasks and other workflows into a directed acyclic graph (DAG). Internally, flytekit represents each step in this graph as a Node, which encapsulates the execution entity (task, sub-workflow, or launch plan) and its configuration.

Defining Workflows with Decorators

The most common way to compose a workflow is using the @workflow decorator. When you decorate a function with @workflow, flytekit transforms the Python function into a WorkflowBase object. Inside the function body, calling a task does not execute it immediately; instead, it creates a Node and returns Promise objects representing the future outputs of that node.

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 in the workflow graph
    result = add_one(x=val)
    return result

When my_workflow is called, flytekit enters a compilation state (managed by FlyteContext). Each task call is intercepted by flyte_entity_call_handler in flytekit/core/promise.py, which creates a Node and links it to the workflow.

Node Composition and Dependencies

Workflows connect nodes primarily through data dependencies. When you pass a Promise (the return value of one task) as an input to another task, flytekit automatically establishes an execution dependency.

Explicit Execution Order

If two nodes do not share data but must run in a specific order, you can use the >> operator or the .runs_before() method. To do this, you must first create the nodes explicitly using create_node from flytekit.core.node_creation.

from flytekit import task, workflow, create_node

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

@task
def compute(x: int) -> int:
    return x * 2

@workflow
def ordered_workflow(x: int) -> int:
    setup_node = create_node(setup)
    compute_node = create_node(compute, x=x)
    
    # Ensure setup runs before compute
    setup_node >> compute_node
    
    return compute_node.o0

The Node.runs_before method (and the __rshift__ operator) modifies the _upstream_nodes list of the target node, ensuring the Flyte engine respects this order during execution.

Sub-workflows

Workflows can be nested. When a workflow calls another workflow, flytekit treats the sub-workflow as a single node in the parent graph. This allows for modularity and reuse of complex logic.

@workflow
def sub_workflow(a: int) -> int:
    return add_one(x=a)

@workflow
def parent_workflow(a: int) -> int:
    # sub_workflow is treated as a node here
    return sub_workflow(a=a)

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 call without changing the task definition itself. This is useful for adjusting resources, retries, or timeouts for specific instances of a task.

from flytekit import Resources

@workflow
def resource_workflow(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="heavy-computation-node"
    )

Internally, with_overrides updates the Node object's attributes:

• Resources: Handled by convert_resources_to_resource_model and stored in self._resources.
• Metadata: The _override_node_metadata method updates NodeMetadata, including retries, interruptible, and timeout.
• Naming: The node_name argument is passed to _dnsify to ensure the node ID is Kubernetes-compliant.

Imperative Workflows

For scenarios where the workflow structure is dynamic or programmatically generated, flytekit provides the ImperativeWorkflow class. This allows you to build a workflow step-by-step using an API rather than a decorated function.

from flytekit.core.workflow import ImperativeWorkflow

# Create the workflow container
wb = ImperativeWorkflow(name="dynamic_workflow")

# Add inputs
in1 = wb.add_workflow_input("val", int)

# Add tasks (entities)
node = wb.add_entity(add_one, x=in1)

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

The ImperativeWorkflow manually manages its CompilationState and provides methods like add_entity and add_workflow_output to construct the underlying Node list and Binding models that the Flyte engine requires.

Internal Node Representation

The Node class in flytekit/core/node.py is the central data structure for workflow construction. Key properties include:

• id: A unique, DNS-compliant identifier for the node within the workflow.
• bindings: A list of Binding objects that map workflow inputs or upstream node outputs to the inputs of this node's entity.
• flyte_entity: The actual task, sub-workflow, or launch plan this node executes.
• upstream_nodes: A list of nodes that must complete before this node can start.

When a workflow is compiled for registration, these Node objects are translated into the Flyte IDL (Interface Definition Language) models that the Flyte backend understands.