Task authoring and execution

Flyte tasks are the fundamental building blocks of workflows in flytekit. They represent a single unit of execution, defined by a typed interface (inputs and outputs) and a specific execution behavior.

Declaring Tasks

The most common way to author a task is using the @task decorator. This decorator transforms a standard Python function into a PythonFunctionTask instance.

from flytekit import task
import typing

@task(cache=True, cache_version="1.0", retries=3)
def my_task(x: int, y: typing.Dict[str, str]) -> str:
    return f"{x} - {y}"

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

1. Metadata Construction: It creates a TaskMetadata object (from flytekit.core.base_task) to store execution policies like retries, timeouts, and caching.
2. Interface Extraction: It uses Python type hints to derive a TypedInterface. This interface is used by the Flyte backend to validate data flow between tasks.
3. Plugin Selection: It identifies the correct task class. For standard functions, this is PythonFunctionTask. For async def functions, it automatically selects AsyncPythonFunctionTask.
4. Global Registration: The task instance is appended to FlyteEntities.entities, allowing flytekit to track all entities defined in a module.

Task Metadata and Configuration

The TaskMetadata class defines how a task should behave on the Flyte platform. Key attributes include:

• Caching: Controlled by cache, cache_version, and cache_serialize. Caching requires a cache_version string. If cache_serialize is enabled, Flyte ensures that concurrent executions with the same inputs are run serially to avoid redundant work.
• Retries: The retries integer determines how many times the platform will attempt to re-run the task upon failure.
• Interruptible: A boolean indicating if the task can be scheduled on lower-priority, pre-emptible nodes (e.g., AWS Spot Instances).
• Timeout: Can be an int (seconds) or datetime.timedelta. TaskMetadata.__post_init__ automatically converts integers to timedeltas.

Execution Models

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

Default Execution

In the default mode, the task runs as a single unit of work. When executed locally, flytekit calls the underlying function directly. When executed on a cluster, the dispatch_execute method handles the lifecycle:

1. Pre-execute: Calls pre_execute to set up environment-specific parameters (like Spark sessions).
2. Input Conversion: Translates Flyte literals from the backend into native Python types using the TypeEngine.
3. User Code: Invokes the actual Python function.
4. Post-execute: Runs post_execute for cleanup or output modification.
5. Output Conversion: Translates Python return values back into Flyte literals.

Dynamic Tasks

Dynamic tasks are declared using the @dynamic decorator (which is a partial application of @task with execution_mode=DYNAMIC).

from flytekit import dynamic

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

Internally, a dynamic task does not execute its body as a single unit. Instead, it compiles the body into a DynamicJobSpec at runtime. This spec contains a new workflow graph generated based on the task's inputs. This allows for workflows where the graph structure depends on runtime data (e.g., processing a variable number of files).

Eager Workflows

Eager workflows (using the @eager decorator) allow for a more "Pythonic" async execution style. Unlike dynamic tasks, they are not compiled into a static spec.

from flytekit import eager, task

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

@eager
async def eager_workflow(x: int) -> int:
    # This looks like standard Python async code
    out = await add_one(x=x)
    return await add_one(x=out)

When running on a backend, EagerAsyncPythonFunctionTask uses a Controller (from flytekit.core.worker_queue) to manage a queue of executions. Every time a Flyte entity is called within an eager function, it triggers a remote execution on the Flyte cluster.

Task Rehydration and Resolvers

When a task runs on a hosted Flyte cluster, it executes inside a container. The container needs to know which specific task object to instantiate and run. This is handled by the TaskResolverMixin.

The default command for a hosted task looks like this:

pyflyte-execute --resolver flytekit.core.python_auto_container.default_task_resolver \
-- task-module my_module task-name my_task_function

The default_task_resolver uses the module path and function name to import and re-create the task instance. Note: Because the default resolver relies on module-level imports, task functions cannot be nested or local functions (unless they are in a test module starting with test_). If you need custom rehydration logic, you must implement a class that inherits from TaskResolverMixin and provides loader_args and load_task methods.

Advanced Authoring Tools

kwtypes

The kwtypes helper (in flytekit.core.base_task) is used to define ordered type mappings, which is critical for maintaining interface consistency in Flyte.

from flytekit import kwtypes
from typing import Annotated
import pandas as pd

# Used for structured data column definitions
stats_cols = kwtypes(mean=float, std=float)

@task
def compute_stats() -> Annotated[pd.DataFrame, stats_cols]:
    ...

PythonInstanceTask

For tasks that do not wrap a user-defined function but instead have platform-defined logic, use PythonInstanceTask. This is the base for many specialized plugins.

class MyCustomTask(PythonInstanceTask):
    def __init__(self, name, **kwargs):
        super().__init__(name=name, task_config=None, task_type="my-type", **kwargs)

    def execute(self, **kwargs):
        # Implementation logic here
        ...

IgnoreOutputs

If a task produces data that should be discarded (common in distributed training where only the rank-0 process might output a model), the task can raise the IgnoreOutputs exception. The PythonTask.dispatch_execute logic is designed to catch this and signal to the platform that the execution finished successfully without producing the declared outputs.