Task authoring and execution

Flyte tasks are the fundamental building blocks of a Flyte workflow. They represent a single unit of execution, characterized by a strong interface (typed inputs and outputs), versioning, and independent executability. In flytekit, tasks are primarily authored using the @task decorator, which transforms a standard Python function into a PythonFunctionTask.

Declaring Tasks

To define a task in flytekit, you decorate a Python function with the @task decorator. The decorator automatically inspects the function's type hints to generate a TypedInterface for the Flyte backend.

from flytekit import task

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

When you call this function, flytekit handles the translation between Python native types and Flyte's internal type system. Internally, the @task decorator instantiates a flytekit.core.python_function_task.PythonFunctionTask. This class captures the function body, its module location, and its interface.

Task Metadata and Configuration

The @task decorator accepts numerous parameters to control how the task is executed on the Flyte platform. These settings are encapsulated in the flytekit.core.base_task.TaskMetadata class.

Common configuration options include:

• Retries: Specify how many times the task should be retried on failure.
• Timeout: Set a maximum duration for the task execution using datetime.timedelta or an integer (seconds).
• Caching: Enable caching to avoid re-running tasks with the same inputs.
• Resources: Request specific CPU, memory, or GPU resources using the requests and limits parameters.

from datetime import timedelta
from flytekit import task, Resources

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

In flytekit/core/base_task.py, the TaskMetadata class ensures that these parameters are valid. For example, it raises a ValueError if cache=True is set without a cache_version.

Task Execution Flow

Flyte tasks follow different execution paths depending on whether they are running locally or on a remote Flyte cluster.

Local Execution

When you run a task locally (e.g., in a script or unit test), flytekit invokes the local_execute method of the Task class. This method:

1. Translates Python native inputs into Flyte Literal objects.
2. Checks the LocalTaskCache if caching is enabled.
3. Calls dispatch_execute, which eventually runs your original Python function.
4. Wraps the results back into Promise objects or native Python types.

Remote Execution

On a Flyte cluster, the execution is handled by pyflyte-execute. The platform uses a TaskResolverMixin to locate and load the task code. The default_task_resolver in flytekit identifies the task by its module path and function name.

The entry point on the container calls dispatch_execute, which:

1. Prepares the execution environment via pre_execute.
2. Converts the input LiteralMap from the Flyte engine into Python native variables.
3. Executes the user-defined function.
4. Converts the return values back into a LiteralMap to be sent back to the Flyte engine.

Specialized Task Types

While PythonFunctionTask is the most common, flytekit provides specialized abstractions for different execution behaviors.

Dynamic Tasks

A dynamic task is declared using the @dynamic decorator (which is a PythonFunctionTask with ExecutionBehavior.DYNAMIC). These tasks can generate new tasks or workflows at runtime based on their inputs.

from flytekit import task, dynamic

@task
def process_item(item: int) -> int:
    return item * 2

@dynamic
def dynamic_subworkflow(items: list[int]) -> list[int]:
    return [process_item(item=i) for i in items]

Internally, PythonFunctionTask.compile_into_workflow is called during execution to produce a DynamicJobSpec, which the Flyte engine then executes as a sub-workflow.

Async and Eager Tasks

Flytekit supports asynchronous tasks via AsyncPythonFunctionTask. If you decorate an async def function with @task, flytekit automatically selects this class.

For more complex interactive patterns, EagerAsyncPythonFunctionTask (used via the @eager decorator) allows Python code to act as a coordinator, triggering executions on the Flyte cluster as if they were standard function calls, while maintaining a stack frame on the cluster.

Core Abstractions

The task system is built on a hierarchy of classes in flytekit/core/base_task.py:

• Task: The base class that captures the Flyte IDL TaskTemplate. It defines the interface and metadata but has no Python-specific execution logic.
• PythonTask: Adds a python_interface property, allowing flytekit to map between Python types and Flyte types using the TypeEngine.
• PythonFunctionTask: The implementation for tasks defined by a Python function. It handles the logic for extracting the interface from function signatures and executing the function body.

If you need to ignore specific outputs (for example, in distributed training where only the rank-0 process returns a result), you can raise the IgnoreOutputs exception defined in base_task.py.