Task authoring and execution

Flyte tasks are the fundamental building blocks of flytekit. They represent a single unit of work with a strongly typed interface, allowing for independent execution, versioning, and unit testing.

Declaring Tasks with Decorators

The most common way to define a task in flytekit is using the @task decorator from flytekit.core.task. This decorator wraps a standard Python function and converts it into a PythonFunctionTask.

from flytekit import task

@task(cache=True, cache_version="1.0", retries=3)
def square(n: int) -> int:
    return n * n

When you use @task, flytekit performs several actions:

1. Interface Detection: It uses transform_function_to_interface to inspect the function's type hints and docstrings to build a TypedInterface.
2. Metadata Configuration: It packages parameters like cache, retries, and timeout into a TaskMetadata object.
3. Plugin Selection: It identifies the appropriate task plugin. By default, it uses PythonFunctionTask, but if the function is async, it uses AsyncPythonFunctionTask.

Task Configuration and Metadata

The TaskMetadata class in flytekit.core.base_task stores execution-related settings. Key attributes include:

• cache: Enables caching of results based on input values.
• cache_version: A string that, when changed, forces a cache miss even if inputs remain the same.
• retries: The number of times the Flyte platform should retry the task on failure.
• timeout: A datetime.timedelta or integer seconds specifying the maximum execution duration.

Core Task Abstractions

Flytekit uses a hierarchical class structure to manage different types of tasks:

The Base Task

The Task class in flytekit.core.base_task is the root abstraction. It captures the information required by the Flyte IDL TaskTemplate. It does not have a Python-native interface and is primarily used as a foundation for more specific task types.

PythonTask

PythonTask extends Task by adding a python_interface. This class is intended for tasks that have a Python-native signature but might not execute a standard Python function body (for example, a task that generates a SQL query).

PythonFunctionTask

PythonFunctionTask (found in flytekit.core.python_function_task) is the workhorse for most flytekit users. It handles:

• Local Execution: When you call a task function directly in Python, local_execute is triggered. It translates Python native inputs into Flyte literals, checks the LocalTaskCache, and then calls execute.
• Remote Dispatch: During runtime on a Flyte cluster, dispatch_execute is called. It handles the translation of Flyte literals back into Python types before running the user's code.

Specialized Execution Modes

PythonFunctionTask supports different ExecutionBehavior modes that change how the task body is interpreted.

Dynamic Tasks

A dynamic task is declared using the @dynamic decorator (a partial of @task with execution_mode=DYNAMIC).

from flytekit import dynamic

@dynamic
def my_dynamic_task(s: str) -> list[str]:
    # This code runs at execution time to produce a workflow
    return [s, s[::-1]]

Internally, dynamic_execute mimics a workflow's compilation. It runs the function body to produce a DynamicJobSpec, which contains a set of nodes and tasks that Flyte Propeller then executes as a sub-workflow.

Eager Tasks

Eager tasks (using ExecutionBehavior.EAGER) allow for more flexible, Pythonic control flow where the execution of tasks can depend on the results of previous tasks within the same function, similar to a standard Python script but executed on the Flyte backend.

The EagerAsyncPythonFunctionTask manages this by using a Controller to communicate with the Flyte backend, submitting tasks to a worker queue and awaiting their results.

Authoring Custom Task Types

To create a new task type (e.g., for a specific database or external service), you should inherit from PythonTask or PythonFunctionTask.

The SQLTask in flytekit.models.sql (and its implementations) is a prime example. It overrides get_sql to provide the SQL definition to the Flyte backend while typically leaving execute as a NotImplementedError for local execution, requiring users to mock it.

class MyCustomTask(PythonTask):
    def __init__(self, name, custom_config, **kwargs):
        super().__init__(task_type="my_custom_type", name=name, **kwargs)
        self.custom_config = custom_config

    def execute(self, **kwargs) -> Any:
        # Define local execution behavior here
        pass

Task Resolvers

When a task is executed on a remote cluster, the container needs to know how to find and load the specific task object. This is handled by TaskResolverMixin.

The default_task_resolver in flytekit.core.python_auto_container is the standard implementation. It:

1. Serializes: Captures the module and function name in loader_args.
2. Loads: Uses importlib to import the module and retrieve the task object during container startup via load_task.

Testing Tasks

Because tasks often interact with the Flyte platform or external systems, flytekit provides utilities in flytekit.testing to mock their behavior during unit tests.

The task_mock context manager allows you to replace a task's execution logic with a mock:

from flytekit.testing import task_mock

@task
def t1(i: int) -> int:
    ...

with task_mock(t1) as m:
    m.side_effect = lambda x: x + 1
    assert t1(i=1) == 2

This is particularly useful for PythonTask subclasses that do not have a native execute implementation.