Task authoring and execution

Flytekit tasks are the fundamental building blocks of a Flyte workflow. They represent a unit of work with a strongly typed interface, allowing for independent execution, versioning, and unit testing. In flytekit, tasks are primarily declared using the @task decorator, which transforms a standard Python function into a PythonFunctionTask.

Declaring Tasks

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

import typing
from flytekit import task

@task
def process_data(x: int, y: typing.Dict[str, str]) -> str:
    return f"Processed {x} with {y}"

When you call this function, flytekit intercepts the call via the flyte_entity_call_handler. In a local execution context, it invokes local_execute, which handles the translation of Python native values into Flyte literals using the TypeEngine.

Task Hierarchy and Abstractions

Flytekit uses a layered class hierarchy to manage task definitions and execution:

1. base_task.Task: The language-agnostic base class. It captures the TaskTemplate information required by the Flyte IDL, such as the task_type, name, and interface.
2. base_task.PythonTask: A subclass of Task that introduces a python_interface (an instance of flytekit.core.interface.Interface). It handles the conversion between Python types and Flyte's internal LiteralMap.
3. python_function_task.PythonFunctionTask: The standard implementation for tasks wrapping a Python function. It uses transform_function_to_interface to automatically generate the task interface from the function signature.

Task Configuration

The @task decorator accepts several parameters to control execution behavior, resource allocation, and metadata. These are encapsulated internally in the TaskMetadata class.

Caching and Retries

You can enable caching to avoid redundant computations when inputs haven't changed. The cache_version must be updated manually if the logic inside the task changes but the signature remains the same.

from flytekit import task, Cache

@task(cache=True, cache_version="1.0", retries=3)
def compute_expensive_value(a: int) -> int:
    return a * 42

Internally, TaskMetadata validates these settings. For instance, if cache=True is set, cache_version becomes mandatory, as enforced in TaskMetadata.__post_init__.

Resource Management

Tasks can request specific hardware resources like CPU, memory, or GPUs.

from flytekit import task, Resources

@task(
    requests=Resources(cpu="2", mem="500Mi"),
    limits=Resources(cpu="4", mem="1Gi"),
    interruptible=True
)
def resource_intensive_task(data: list) -> int:
    return len(data)

Execution Modes

Flytekit supports different execution behaviors through the PythonFunctionTask.ExecutionBehavior enum.

Default Execution

In the DEFAULT mode, the task runs as a single unit of work. When dispatch_execute is called (either locally or on the cluster), it simply executes the wrapped _task_function.

Dynamic Execution

The @dynamic decorator (a partial application of @task with execution_mode=ExecutionBehavior.DYNAMIC) allows a task to generate a workflow at runtime based on its inputs.

from flytekit import dynamic

@dynamic
def my_dynamic_subwf(a: int) -> typing.List[str]:
    s = []
    for i in range(a):
        # You can use native Python logic like loops here
        s.append(process_data(x=i, y={"key": "val"}))
    return s

When a dynamic task runs:

1. PythonFunctionTask.execute calls dynamic_execute.
2. If running on the Flyte backend, compile_into_workflow is invoked.
3. This produces a DynamicJobSpec containing the generated nodes and task templates, which Flyte Propeller then executes as a subworkflow.

Eager Execution

Eager tasks (declared via EagerAsyncPythonFunctionTask) allow for asynchronous, imperative-style execution where Python code acts as the orchestrator. Unlike dynamic tasks, eager tasks can await results and make decisions based on them in real-time, with each task call creating a new execution on the Flyte cluster.

Task Resolvers

When a task is executed on a remote cluster, Flyte needs to know how to "rehydrate" the Python object from the container. This is handled by TaskResolverMixin.

The default_task_resolver (used by PythonAutoContainerTask) identifies tasks by their module and function name. At serialization time, loader_args generates the command-line arguments (e.g., task-module path.to.module task-name my_task) that pyflyte-execute uses to find and load the task via load_task.

Local Execution and Testing

Flytekit tasks are designed to be unit-testable. When you call a task function directly in a script or test, Task.local_execute is triggered. It performs the following steps:

1. Translates native Python inputs to Flyte Literals.
2. Checks the LocalTaskCache if caching is enabled.
3. Invokes sandbox_execute, which calls dispatch_execute.
4. Translates the resulting LiteralMap back into Python native types (or Promise objects if within a workflow context).