Workflow composition and nodes

Flyte workflows are defined as Directed Acyclic Graphs (DAGs) where each vertex is a Node and each edge represents a data or execution dependency. In flytekit, this structure is typically built using the @workflow decorator, which captures the relationships between tasks at compile time.

Defining Workflows

When you decorate a function with @workflow, flytekit treats the function body as a DSL for defining a DAG. Instead of executing the logic immediately, flytekit runs the function in a "compilation" mode to record how tasks are called and how their outputs flow into subsequent tasks.

from flytekit import task, workflow

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

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

@workflow
def my_workflow(a: int) -> int:
    res1 = t1(a=a)
    res2 = t2(a=a, b=res1)
    return res2

Internally, the WorkflowBase class (found in workflow.py) manages this process. When a task like t1 is called inside a workflow, the flyte_entity_call_handler in promise.py detects that the system is in CompilationMode. Instead of running the task, it calls create_and_link_node, which:

1. Creates a Node object representing that specific task execution.
2. Generates Promise objects for the task's outputs.
3. Binds the task's inputs to either workflow inputs or Promise objects from upstream nodes.

Understanding Nodes

A Node is the fundamental unit of execution in a Flyte DAG. Every time you call a task, sub-workflow, or launch plan within a workflow, flytekit creates a Node instance.

The Node class (defined in flytekit/core/node.py) encapsulates:

• ID: A unique identifier within the workflow (e.g., n0, n1).
• Flyte Entity: The actual task, workflow, or launch plan the node executes.
• Bindings: A list of Binding objects that map inputs to their sources.
• Upstream Nodes: A list of nodes that must complete before this node can run.

Explicit Node Creation

While most nodes are created implicitly by calling tasks, you can create them explicitly using create_node from flytekit/core/node_creation.py. This is useful when you need to define execution order without a direct data dependency.

from flytekit import task, workflow, create_node

@task
def t1():
    print("Task 1")

@task
def t2():
    print("Task 2")

@workflow
def manual_dependency_wf():
    # Create nodes without calling them as functions
    t1_node = create_node(t1)
    t2_node = create_node(t2)

    # Force t1 to run before t2 using the shift operator
    t1_node >> t2_node

The >> operator is a shorthand for the runs_before method on the Node class. It adds the left-hand node to the _upstream_nodes list of the right-hand node, ensuring the Flyte engine respects this order during execution.

Connecting Inputs and Outputs

Data flow in flytekit is managed through Promise objects. When a node is created, it returns Promise objects (or a tuple of them) representing its future outputs.

When you pass a Promise as an argument to another task:

1. create_and_link_node inspects the input.
2. It identifies the Node that produced the Promise.
3. It adds that producing node to the current node's upstream_nodes.
4. It creates a Binding that tells the Flyte engine to pull the specific output variable from the upstream node.

Node Overrides

You can customize the execution behavior of individual nodes using the with_overrides method. This allows you to set resource requirements, timeouts, and retries for a specific instance of a task without changing the task definition itself.

@workflow
def override_wf(a: int) -> int:
    return t1(a=a).with_overrides(
        node_name="custom-t1-node",
        requests=Resources(cpu="2", mem="500Mi"),
        retries=3,
        timeout=datetime.timedelta(minutes=5)
    )

The with_overrides method in node.py updates the NodeMetadata and _resources attributes of the Node instance. It supports:

• Resources: CPU, memory, and GPU limits/requests via Resources objects.
• Metadata: timeout, retries, and interruptible flags.
• Caching: Overriding cache settings using the Cache object.
• Container Image: Specifying a different image for a specific node.

Imperative Workflows

For scenarios where the DAG structure is determined dynamically (e.g., based on a configuration file), flytekit provides the ImperativeWorkflow class. Unlike the @workflow decorator, which uses Python's function execution to build the DAG, imperative workflows allow you to add nodes and inputs programmatically.

from flytekit import Workflow

wb = Workflow(name="imperative_workflow")
wb.add_workflow_input("in1", int)
node = wb.add_entity(t1, a=wb.inputs["in1"])
wb.add_workflow_output("out1", node.outputs["o0"])

In an imperative workflow, you use add_entity to create nodes and add_workflow_output to define the final results. This bypasses the standard Python function body while still producing a valid Flyte DAG.