Workflow composition and nodes

Flytekit workflows are the primary mechanism for composing tasks and other entities into a directed acyclic graph (DAG). When you define a workflow using the @workflow decorator, flytekit compiles the function body into a set of Nodes, each representing an execution step, and Bindings, which define how data flows between these steps.

Workflow Composition with Decorators

The most common way to compose workflows is by using the @workflow decorator. Within the decorated function, you call tasks as if they were regular Python functions. However, instead of returning actual values, these calls return Promise objects.

from flytekit import task, workflow

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

@task
def multiply(x: int, y: int) -> int:
    return x * y

@workflow
def math_workflow(val: int) -> int:
    # Each task call creates a Node internally
    step1 = add_one(x=val)
    step2 = multiply(x=step1, y=2)
    return step2

When math_workflow is called during compilation, flytekit tracks the creation of each Node. The step1 object is a Promise that represents the output of the first node. Passing step1 into multiply creates a dependency (a binding) between the two nodes.

Internal Node Creation

Internally, every time a task or sub-workflow is invoked within a workflow context, flytekit uses the flyte_entity_call_handler (found in flytekit/core/promise.py) to wrap the entity in a Node.

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

• ID: A unique identifier within the workflow (often derived from the task name).
• Flyte Entity: The actual task, launch plan, or sub-workflow being executed.
• Bindings: A list of Binding models that map inputs to the outputs of upstream nodes or workflow-level inputs.
• Metadata: Execution settings like timeouts and retries.

Customizing Node Execution

Sometimes a specific step in your workflow requires more resources or a different retry policy than the task's default. You can use the with_overrides method on the object returned by a task call to modify the underlying Node configuration.

from flytekit import Resources

@workflow
def resource_workflow(val: int) -> int:
    # Override resources and retries for this specific node
    return add_one(x=val).with_overrides(
        requests=Resources(cpu="2", mem="200Mi"),
        retries=3,
        node_name="heavy-math-step"
    )

The Node.with_overrides method allows you to customize:

• node_name: Changes the ID of the node in the Flyte DAG.
• requests and limits: Sets Resources (CPU, Memory, GPU, Ephemeral Storage).
• timeout: A datetime.timedelta or integer seconds.
• retries: Number of times to retry on failure.
• interruptible: Boolean indicating if the node can run on spot/preemptible instances.

Internally, with_overrides calls _override_node_metadata to update the NodeMetadata object and convert_resources_to_resource_model to handle hardware specifications.

Explicit Dependencies

While data flow usually defines the execution order, you may sometimes need to enforce that a task runs after another even if it doesn't consume its output (e.g., a task that performs a side effect). You can use the >> operator or the runs_before method.

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

@task
def compute() -> int:
    return 42

@workflow
def dependency_workflow() -> int:
    s = setup()
    c = compute()
    
    # Force 'setup' to run before 'compute'
    s >> c
    return c

The Node.__rshift__ operator is a convenience alias for Node.runs_before. It appends the current node to the _upstream_nodes list of the target node, ensuring the Flyte propeller executes them in the correct sequence.

Imperative Workflows

For scenarios where the workflow structure is dynamic or generated at runtime, flytekit provides the ImperativeWorkflow class. This allows you to build a workflow programmatically without the @workflow decorator.

from flytekit.core.workflow import ImperativeWorkflow

# Initialize the workflow
wb = ImperativeWorkflow(name="dynamic_math")

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

# Add entities (tasks) and connect them
node1 = wb.add_entity(add_one, x=val_input)
node2 = wb.add_entity(multiply, x=node1.outputs["o0"], y=2)

# Define workflow outputs
wb.add_workflow_output("result", node2.outputs["o0"])

The ImperativeWorkflow class (in flytekit/core/workflow.py) maintains a CompilationState that tracks nodes as they are added via add_entity. This is particularly useful for building "workflow builders" or UI-driven workflow generators.

Data Flow and Promises

Data is passed between nodes using Promise objects. A Promise is a placeholder for a value that will exist at execution time.

• When you access node.outputs["output_name"], you get a NodeOutput wrapped in a Promise.
• When this Promise is passed as an input to another task, flytekit creates a Binding.
• During local execution, the WorkflowBase.local_execute method resolves these promises by executing nodes in topological order and passing the resulting literals to downstream tasks.

Binding Mechanism

The Binding model (from flytekit.models.literals) connects a task's input variable to a source. This source can be:

1. A Scalar: A constant value passed directly.
2. A Promise: An output from a previous Node.
3. A Workflow Input: A value provided when the workflow starts.

This binding system ensures that Flyte can reconstruct the DAG and manage data movement between distributed workers during remote execution.