Task authoring and execution

Flyte tasks are the fundamental building blocks of a Flyte workflow. They represent a discrete unit of work, versioned and strongly typed, that can be executed independently or as part of a larger pipeline. In flytekit, tasks are primarily authored using the @task decorator, which transforms a standard Python function into a PythonFunctionTask.

Declaring Tasks with @task

The most common way to define a task is by decorating a Python function with @task. This decorator automatically handles the conversion of your Python function into a Flyte-compatible entity by inspecting its type hints and docstrings.

from flytekit import task
import typing

@task
def greet(name: str) -> str:
    """
    A simple task that greets a user.
    """
    return f"Hello, {name}!"

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

1. Interface Inference: It uses transform_function_to_interface (from flytekit.core.interface) to extract input and output types from the function's type annotations.
2. Metadata Creation: It initializes a TaskMetadata object to store configuration like retries, timeouts, and caching settings.
3. Task Instantiation: It creates an instance of PythonFunctionTask (or a specialized plugin class if task_config is provided).

Task Configuration

You can customize task behavior by passing arguments to the @task decorator. These arguments populate the TaskMetadata and other internal attributes.

from flytekit import task
from datetime import timedelta

@task(
    retries=3,
    timeout=timedelta(minutes=5),
    cache=True,
    cache_version="1.0",
    environment={"MY_ENV_VAR": "value"}
)
def heavy_computation(x: int) -> int:
    return x * x

Key configuration parameters include:

• retries: Number of times to retry the task on failure. This is stored in TaskMetadata.retries.
• timeout: Maximum duration for a single execution, managed via TaskMetadata.timeout.
• cache and cache_version: Enables caching of results based on input values. TaskMetadata validates that if cache is True, a cache_version must be provided.
• task_config: Used for plugin-specific configurations (e.g., Spark, Pod). TaskPlugins.find_pythontask_plugin uses the type of task_config to determine which specialized PythonFunctionTask subclass to instantiate.

Core Task Abstractions

Flytekit uses a hierarchy of classes to manage task definitions and execution:

• Task: The base class for all tasks in flytekit. It captures the Flyte IDL TaskTemplate information, including the interface, metadata, and task_type.
• PythonTask: A subclass of Task designed for tasks with a Python-native interface. It provides methods like _literal_map_to_python_input and _output_to_literal_map to bridge the gap between Flyte's type system (Literals) and Python objects.
• PythonFunctionTask: The specific implementation used for @task. It holds a reference to the actual Python function (self._task_function) and implements the execute method by calling that function.

Execution Flow

When a task is executed (either locally or on a Flyte cluster), the dispatch_execute method in PythonTask orchestrates the process:

1. Pre-execution: pre_execute is called to set up the environment (e.g., initializing a Spark session).
2. Input Translation: The input LiteralMap is converted to Python native values using TypeEngine.literal_map_to_kwargs.
3. User Code Execution: The execute method is called with the native inputs. For PythonFunctionTask, this simply runs the decorated function.
4. Output Translation: The Python return values are converted back into a Flyte LiteralMap via _output_to_literal_map.

Dynamic Tasks

Dynamic tasks allow you to generate a workflow structure at runtime based on input data. They are declared using the @dynamic decorator, which is a specialized version of @task with execution_mode set to PythonFunctionTask.ExecutionBehavior.DYNAMIC.

from flytekit import task, dynamic
import typing

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

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

Internally, when a dynamic task runs:

1. It executes the function body to gather a set of nodes (task calls).
2. compile_into_workflow is called to produce a DynamicJobSpec.
3. This spec is returned to the Flyte engine, which then schedules the generated subworkflow.

[!WARNING] Dynamic workflows should generally be kept under 50 tasks to avoid excessive overhead during compilation and processing.

Task Resolvers

When a task is executed on a remote cluster, Flyte needs to know how to find and load the Python code. This is handled by TaskResolverMixin. The default resolver, default_task_resolver, identifies tasks by their module and name.

If you need custom loading logic (e.g., loading tasks from a database or a dynamic source), you can implement your own resolver by subclassing TaskResolverMixin and providing it to the @task decorator via the task_resolver parameter.

class MyCustomResolver(TaskResolverMixin):
    def load_task(self, loader_args: List[str]) -> Task:
        # Custom logic to rehydrate the task object
        ...

    def loader_args(self, settings: SerializationSettings, t: Task) -> List[str]:
        # Arguments needed to identify the task at runtime
        ...