Task authoring and execution

Flyte tasks are the fundamental building blocks of a workflow. They represent a discrete unit of work with a strongly typed interface, allowing flytekit to handle data movement, retries, and execution environment management.

Declaring Tasks

The primary way to define a task in flytekit is by using the @task decorator on a Python function. When you decorate a function, flytekit automatically inspects the function signature to determine its inputs and outputs, creating a PythonFunctionTask.

from flytekit import task

@task
def greet(name: str) -> str:
    return f"Hello, {name}!"

In this example, flytekit uses the type hints name: str and -> str to define the task's interface. Internally, the PythonFunctionTask class handles the conversion between Flyte's internal type system (IDL) and Python native types.

Task Configuration

You can customize task behavior by passing arguments to the @task decorator. These configurations are stored in the TaskMetadata class and influence how the Flyte engine executes the task.

Common configuration options include:

• Retries: Automatically retry the task on failure.
• Caching: Cache results based on input values and a version string.
• Timeout: Limit the maximum duration of a single execution.
• Resources: Specify CPU, memory, and GPU requirements.

import datetime
from flytekit import task, Resources

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

If you enable caching (cache=True), you must also provide a cache_version. The TaskMetadata.__post_init__ method enforces this requirement, raising a ValueError if the version is missing.

Core Task Abstractions

Flytekit uses a hierarchy of classes to represent different types of tasks:

1. Task: The base class located in flytekit.core.base_task. it captures the low-level Flyte IDL specification, including the task_type, name, and interface.
2. PythonTask: A subclass of Task that adds support for Python-native interfaces. It manages the Interface object which maps Python types to Flyte types.
3. PythonFunctionTask: The most common task type, which wraps a user-defined Python function. It implements the execute method by calling the underlying _task_function.

The Execution Lifecycle

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

1. pre_execute: Prepares the execution environment. This is often used by plugins to set up specific contexts (e.g., a Spark session).
2. Input Translation: Converts Flyte LiteralMap inputs into Python-native values using the TypeEngine.
3. execute: Invokes the actual user code.
4. post_execute: Performs cleanup or output modification.
5. Output Translation: Converts the Python return values back into a Flyte LiteralMap.

Advanced Execution Behaviors

The PythonFunctionTask supports different execution modes via the ExecutionBehavior enum:

Dynamic Tasks

A task can be marked as dynamic by setting execution_mode=PythonFunctionTask.ExecutionBehavior.DYNAMIC (or using the @dynamic decorator). Dynamic tasks allow you to generate a new workflow structure at runtime based on the task's inputs.

Internally, dynamic_execute compiles the code within the task into a DynamicJobSpec, which Flyte Propeller then executes as a sub-workflow.

Eager Tasks

Eager tasks (declared via EagerAsyncPythonFunctionTask) allow for more flexible, imperative-style execution where Python code acts as the orchestrator. Unlike standard tasks that are compiled into a static graph, eager tasks can make decisions and trigger other tasks dynamically during their execution.

from flytekit import eager

@eager
async def my_eager_workflow(x: int) -> int:
    # This code runs and can await other tasks
    out = await some_task(x=x)
    if out > 10:
        return await task_a(val=out)
    return await task_b(val=out)

Task Resolvers

When a task runs on a remote Flyte cluster, the container needs to know how to find and load the specific Python function. This is handled by the TaskResolverMixin.

The default_task_resolver in flytekit works by:

1. Capturing the module and function name during serialization (loader_args).
2. Importing the module and retrieving the function attribute during execution (load_task).

If you have custom loading requirements (e.g., loading tasks from a database or a dynamic source), you can implement a custom resolver by subclassing TaskResolverMixin and passing it to the @task decorator.