Task authoring and execution

Flyte tasks are the fundamental building blocks of a workflow. In flytekit, tasks are typically declared using the @task decorator, which transforms a standard Python function into a PythonFunctionTask. This abstraction handles the translation between Python's native types and the Flyte IDL, manages execution metadata, and provides hooks for both local and remote execution.

Declaring Tasks

The most common way to author a task is by decorating a type-annotated Python function. flytekit uses these annotations to automatically derive the task's interface.

from flytekit import task
import typing

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

When you call this function, flytekit's flyte_entity_call_handler (found in promise.py) intercepts the call. Depending on the context, it either executes the function locally or creates a node in a workflow graph during compilation.

Task Configuration and Metadata

The @task decorator accepts several parameters that configure how the task behaves on the Flyte platform. These are encapsulated in the TaskMetadata class in base_task.py.

• Caching: Controlled by cache, cache_version, and cache_serialize. Caching requires a cache_version. If cache_serialize=True is set, Flyte ensures that multiple instances of the task with the same inputs are executed serially to avoid redundant work.
• Retries: The retries parameter defines how many times Flyte should attempt to re-run the task upon failure.
• Timeouts: The timeout parameter (either an int representing seconds or a datetime.timedelta) limits the maximum duration of a single execution.
• Resources: You can specify requests and limits for CPU, memory, and GPU resources.

from flytekit import task, Resources

@task(
    cache=True, 
    cache_version="1.0", 
    retries=3, 
    timeout=3600,
    requests=Resources(cpu="2", mem="500Mi"),
    limits=Resources(cpu="4", mem="1Gi")
)
def resource_heavy_task(x: int) -> int:
    return x * x

Core Task Abstractions

The task hierarchy in flytekit provides a flexible foundation for different execution behaviors:

1. Task: The base class for all tasks. It captures the Flyte IDL specification and handles the core local_execute logic, including input/output translation via the TypeEngine.
2. PythonTask: A specialization for tasks with a Python native interface. It implements dispatch_execute, which manages the lifecycle: pre_execute -> input conversion -> execute -> post_execute -> output conversion.
3. PythonAutoContainerTask: Handles the serialization of the task into a container command. It uses a TaskResolverMixin to determine how the task should be rehydrated inside the container at runtime.
4. PythonFunctionTask: The standard class used for functions decorated with @task. It wraps the user's callable and detects its interface.

Plugin-Specific Tasks

For tasks that don't have a user-defined function body but instead rely on a backend plugin (like SQL or Spark), flytekit uses PythonInstanceTask or direct PythonTask subclasses. For example, SQLTask (in base_sql_task.py) defines a typed interface but raises an error in execute because the actual execution is offloaded to a Flyte backend plugin.

Execution Modes

Flytekit supports several specialized execution behaviors through the ExecutionBehavior enum in PythonFunctionTask.

Dynamic Tasks

Declared with the @dynamic decorator, these tasks are executed at runtime to produce a new workflow spec based on the inputs. They are useful when the graph structure depends on data (e.g., processing a variable number of files).

from flytekit import dynamic, task

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

@dynamic
def my_dynamic_task(a: int) -> typing.List[str]:
    s = []
    for i in range(a):
        s.append(t1(a=i))
    return s

Async and Eager Tasks

If you decorate an async def function with @task, flytekit automatically selects AsyncPythonFunctionTask.

Eager Workflows (using the @eager decorator) allow you to use standard Python control flow (like if statements and loops) while still executing tasks on the Flyte cluster. Internally, EagerAsyncPythonFunctionTask uses a Controller and a worker queue to manage remote executions and fetch results as they complete.

from flytekit import task, eager
import asyncio

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

@eager
async def eager_workflow(x: int) -> int:
    # Standard Python logic works here
    if x > 0:
        return await add_one(x=x)
    return 0

# Local execution
if __name__ == "__main__":
    result = asyncio.run(eager_workflow(x=5))

Task Serialization and Resolution

When a task is serialized for the Flyte platform, it must be "rehydratable" in the execution container. This is handled by the TaskResolverMixin.

The default_task_resolver serializes the task's module and name. At runtime, the pyflyte-execute entrypoint uses this resolver to import the module and locate the task object. A typical execution command looks like this:

pyflyte-execute --inputs s3://path/inputs.pb \
                --output-prefix s3://outputs/location \
                --resolver flytekit.core.python_auto_container.default_task_resolver \
                -- task-module my_project.tasks task-name my_task

Constraints and Gotchas

• Nested Functions: The default resolver cannot handle nested or local functions because they cannot be imported by name. Tasks must be defined at the module level.
• Output Rules: Tasks that return nothing should return None or a VoidPromise. If a task returns a single value wrapped in a NamedTuple of length one, flytekit applies special handling to extract the value correctly.
• IgnoreOutputs: The IgnoreOutputs exception can be raised within a task to indicate that its outputs should be discarded, which is particularly useful in distributed training scenarios.