Task authoring and execution

Flyte tasks are the fundamental building blocks of a workflow. In flytekit, tasks are defined by decorating a Python function with the @task decorator, which transforms the function into a PythonFunctionTask object. This object captures the function's interface (inputs and outputs), metadata (like retries and timeouts), and execution logic.

Declaring Tasks

The primary way to author a task is using the @task decorator. flytekit uses Python type hints to automatically derive the TypedInterface for the task, ensuring that data types are consistent across the Flyte platform.

from flytekit import task

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

When you apply @task to a function, flytekit performs the following:

1. Interface Extraction: It calls transform_function_to_interface to inspect the function signature and create a Flyte-compatible interface.
2. Task Instantiation: It creates an instance of PythonFunctionTask (defined in flytekit/core/python_function_task.py).
3. Plugin Selection: It checks TaskPlugins to see if a specialized plugin should be used based on the task_config provided.

Task Metadata and Configuration

You can configure task behavior by passing arguments to the @task decorator. These arguments are stored in a TaskMetadata object (defined in flytekit/core/base_task.py).

from datetime import timedelta
from flytekit import task, Resources

@task(
    retries=3,
    timeout=timedelta(minutes=5),
    requests=Resources(cpu="1", mem="2Gi"),
    limits=Resources(cpu="2", mem="4Gi"),
    cache=True,
    cache_version="1.0"
)
def resource_intensive_task(data: list[int]) -> int:
    return sum(data)

Key configuration options include:

• retries: The number of times Flyte should retry the task on failure.
• timeout: The maximum duration a single execution of the task is allowed to run.
• cache and cache_version: Enables caching of task outputs. If cache=True, a cache_version must be provided.
• requests and limits: Specifies the compute resources (CPU, memory, GPU) required for the task using the Resources class.

Core Task Abstractions

flytekit uses a class hierarchy to manage different types of tasks:

• Task: The base class in flytekit/core/base_task.py. it captures the Flyte IDL TaskTemplate information but lacks a Python-native interface.
• PythonTask: A subclass of Task that adds a python_interface. This is the base for tasks that have a Python-native typed interface.
• PythonFunctionTask: The most common task type, which wraps a standard Python function. It handles the translation between Flyte literals and Python native types during execution.

Custom Task Types

For tasks that don't fit the standard Python function model (like SQL queries or pre-built container logic), you can inherit from PythonTask. For example, a SQL task might look like this:

from flytekit.core.base_task import PythonTask
from flytekit.core.interface import Interface

class MySqlTask(PythonTask):
    def __init__(self, name: str, query: str, **kwargs):
        super().__init__(
            task_type="sql",
            name=name,
            interface=Interface(inputs={}, outputs={"result": int}),
            **kwargs
        )
        self._query = query

    def execute(self, **kwargs):
        # Implementation for executing the SQL query
        return 42

Task Execution Flow

When a task is executed, flytekit follows a specific lifecycle managed by the dispatch_execute method in PythonTask:

1. pre_execute: Prepares the execution environment (e.g., setting up a Spark session).
2. Input Translation: Converts Flyte LiteralMap inputs into Python native values using _literal_map_to_python_input.
3. execute: Invokes the actual user code (the decorated function or the overridden execute method).
4. post_execute: Performs cleanup or output modification.
5. Output Translation: Converts Python native return values back into a Flyte LiteralMap using _output_to_literal_map.

Local vs. Remote Execution

flytekit supports running tasks locally for testing. When you call a task function directly in Python, it triggers local_execute.

# Local execution
result = add_one(x=10)
assert result == 11

In local_execute, flytekit mimics the Flyte platform's behavior by translating inputs to literals and back, ensuring that your type transformers and custom logic work as expected before deploying to a cluster.

Specialized Task Behaviors

Async Tasks

If you decorate an async function with @task, flytekit instantiates an AsyncPythonFunctionTask. This allows you to use await within your task body.

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

Eager Workflows

Eager workflows (declared with the @eager decorator) allow for dynamic execution where Python code acts as the orchestrator. Unlike standard tasks, eager workflows can await other tasks and make decisions based on their results at runtime.

from flytekit import eager

@eager
async def my_eager_workflow(x: int) -> int:
    out = await add_one(x=x)
    return await add_one(x=out)

Internally, @eager creates an EagerAsyncPythonFunctionTask which manages a Controller to communicate with the Flyte backend during execution.

Reference Tasks

A reference_task is a pointer to a task that is already registered on a Flyte cluster. It allows you to use tasks defined in other projects or languages without having the source code available locally.

from flytekit import reference_task

@reference_task(
    project="flytesnacks",
    domain="development",
    name="core.control_flow.merge_sort.merge",
    version="v1"
)
def merge(sorted_list1: list[int], sorted_list2: list[int]) -> list[int]:
    ...

The function body of a reference task is ignored; only the signature is used for workflow compilation.