Task authoring and execution

Flyte tasks are the fundamental building blocks of Flyte workflows. They represent a single unit of execution, characterized by a versioned, strongly-typed interface and declarative configuration. In flytekit, tasks are primarily authored using the @task decorator, which transforms a standard Python function into a PythonFunctionTask.

Declaring Tasks

The most common way to define a task in flytekit is by decorating a Python function with the @task decorator. This decorator automatically infers the task's interface (inputs and outputs) from the function's type hints.

from flytekit import task
import typing

@task
def add_two(a: int) -> typing.NamedTuple("Outputs", [("result", int), ("msg", str)]):
    res = a + 2
    return res, f"Result is {res}"

When you apply the @task decorator, flytekit performs several internal steps:

1. Interface Extraction: It uses transform_function_to_interface (from flytekit.core.interface) to convert Python type hints into a TypedInterface that Flyte understands.
2. Metadata Construction: It creates a TaskMetadata object (defined in flytekit.core.base_task) to store configuration like retries, timeouts, and caching settings.
3. Task Instantiation: It instantiates a PythonFunctionTask (or a specialized plugin subclass). This object captures the function body, the interface, and the metadata.

Constraints on Task Functions

Task functions must be accessible at the module level so that the Flyte container can import and execute them. The PythonFunctionTask constructor in flytekit/core/python_function_task.py enforces this:

if (
    not istestfunction(func=task_function)
    and isnested(func=task_function)
    and not is_functools_wrapped_module_level(task_function)
):
    raise ValueError(
        "TaskFunction cannot be a nested/inner or local function. "
        "It should be accessible at a module level for Flyte to execute it."
    )

Task Configuration

The @task decorator accepts various parameters to control how the task behaves on the Flyte platform.

Caching

Caching allows Flyte to skip execution if a task is called with the same inputs and a matching cache_version. You can use the Cache object for fine-grained control.

from flytekit import task, Cache

@task(cache=True, cache_version="1.0", cache_serialize=True)
def cached_task(x: int) -> int:
    return x * x

# Or using the modern Cache object
@task(cache=Cache(version="1.0", serialize=True, ignored_inputs=("debug_flag",)))
def advanced_cached_task(x: int, debug_flag: bool) -> int:
    return x * x

Internally, TaskMetadata validates these settings. For instance, if cache=True is set, a cache_version must be provided, or a ValueError is raised during __post_init__.

Resources and Retries

You can specify resource requirements (CPU, memory, GPU) and retry logic directly in the decorator.

from flytekit import task, Resources

@task(
    requests=Resources(cpu="1", mem="2Gi"),
    limits=Resources(cpu="2", mem="4Gi"),
    retries=3,
    timeout=3600  # seconds
)
def resource_intensive_task(data: list) -> int:
    return len(data)

Core Task Abstractions

Flytekit uses a hierarchy of classes to manage task behavior:

• Task: The base class in flytekit.core.base_task. It maps closely to the Flyte IDL TaskTemplate and handles the core logic for local_execute and dispatch_execute.
• PythonTask: A subclass of Task that introduces a Python-native interface. It manages the conversion between Flyte's internal Literal types and Python types using the TypeEngine.
• PythonFunctionTask: The class used for tasks defined via @task. It holds a reference to the actual Python function (self._task_function) and implements the execute method by calling that function.

Execution Flow

When a task is executed, the following methods are involved:

1. local_execute: Used during local runs. It translates native Python inputs to Flyte Literals, checks the local cache, calls sandbox_execute, and then translates the results back to Python objects.
2. dispatch_execute: The entry point for both local sandbox and remote runtime execution. It handles pre_execute hooks (e.g., setting up Spark sessions), invokes the user's execute method, and processes post_execute logic.
3. execute: In PythonFunctionTask, this method simply calls the wrapped Python function with the provided keyword arguments.

Specialized Task Types

Dynamic Tasks

A dynamic task is a task that can generate a new workflow at runtime based on its inputs. This is achieved by setting the execution_mode to DYNAMIC.

from flytekit import task, dynamic

@dynamic
def my_dynamic_task(n: int) -> list[int]:
    return [add_two(a=i) for i in range(n)]

When execution_mode == ExecutionBehavior.DYNAMIC, the PythonFunctionTask.execute method calls dynamic_execute, which compiles the returned entities into a DynamicJobSpec.

Eager Tasks

Eager tasks (or eager workflows) allow for more flexible, Pythonic execution where each task call is immediately dispatched to the Flyte backend. These are implemented via EagerAsyncPythonFunctionTask.

from flytekit import eager

@eager
async def my_eager_task(x: int) -> int:
    out = await add_two(a=x)
    if out > 10:
        return await add_two(a=out)
    return out

Eager tasks use a Controller to manage the worker queue and communicate with the Flyte backend during execution. They also support failure handling via EagerFailureHandlerTask, which cleans up running sub-tasks if the parent eager task fails.