Task authoring and execution

Flyte tasks are the fundamental building blocks of a workflow. In flytekit, tasks are primarily authored using the @task decorator, which transforms a standard Python function into a PythonFunctionTask. This abstraction handles the translation between Python types and the Flyte IDL, manages execution metadata, and provides hooks for local and remote execution.

Declaring Tasks

The most common way to define a task is by decorating a typed Python function. flytekit uses the function's type annotations to automatically derive the task's interface (inputs and outputs).

from flytekit import task
import typing

@task
def my_task(x: int, y: typing.Dict[str, str]) -> str:
    return f"{x} - {y}"

Internally, the @task decorator in flytekit/core/task.py performs several steps:

1. Metadata Construction: It creates a TaskMetadata object to store configuration like retries, timeouts, and caching.
2. Interface Discovery: It calls transform_function_to_interface to map Python types to Flyte's type system.
3. Task Instantiation: It selects the appropriate task class (e.g., PythonFunctionTask for standard functions or AsyncPythonFunctionTask for async functions) and instantiates it.
4. Wrapper Update: It uses functools.update_wrapper so the resulting task object retains the original function's docstrings and metadata.

Task Configuration

You can configure task behavior by passing arguments to the @task decorator. These settings are stored in TaskMetadata and PythonTask attributes.

from flytekit import task, Resources, Cache

@task(
    retries=3,
    timeout=3600,
    requests=Resources(cpu="2", mem="500Mi"),
    cache=True,
    cache_version="1.0",
    environment={"MY_ENV_VAR": "value"}
)
def configured_task(x: int) -> int:
    return x + 1

• Caching: When cache=True is set, flytekit requires a cache_version. The Task.local_execute method checks the LocalTaskCache before execution. If a hit is found, it returns the cached LiteralMap instead of running the user code.
• Resources: requests and limits allow you to specify the hardware requirements for the task container on the Flyte cluster.
• Timeouts: The timeout parameter can be an int (seconds) or datetime.timedelta. TaskMetadata.__post_init__ ensures this is normalized to a timedelta.

Core Task Abstractions

The task hierarchy in flytekit provides a layered approach to execution:

1. Task: The base class (flytekit/core/base_task.py) that captures the Flyte IDL TaskTemplate. It defines the core local_execute logic and the dispatch_execute interface.
2. PythonTask: A subclass that adds support for Python-native interfaces. It handles the asynchronous conversion of native outputs back to Flyte literals in _output_to_literal_map.
3. PythonAutoContainerTask: Used for tasks that run within a container. It integrates with TaskResolverMixin to determine how the task should be rehydrated in a remote environment.
4. PythonFunctionTask: The standard implementation for tasks defined by a Python function. It stores the task_function and invokes it during the execute phase.

Task Resolvers

When a task is executed on a Flyte cluster, the container needs to know which task to run. This is handled by a TaskResolverMixin. The DefaultTaskResolver serializes the task's module and name into the container's command-line arguments:

pyflyte-execute --resolver flytekit.core.python_auto_container.default_task_resolver \
-- task-module my_project.tasks task-name my_task

At runtime, the resolver uses importlib.import_module to load the task object back into memory. Because of this, tasks must generally be defined at the module level to be importable.

Execution Modes

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

Default Execution

In DEFAULT mode, the task function is executed normally. For local execution, Task.local_execute translates inputs into literals, calls dispatch_execute, and then converts the results back into Promise objects.

Dynamic Tasks

Dynamic tasks (declared with @dynamic) allow you to generate a workflow structure at runtime based on inputs.

from flytekit import dynamic, task

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

@dynamic
def my_dynamic_task(n: int) -> typing.List[int]:
    return [t1(a=i) for i in range(n)]

When a dynamic task runs on the Flyte platform, PythonFunctionTask.compile_into_workflow is called. It produces a DynamicJobSpec containing the generated nodes and task templates, which Flyte Propeller then executes as a sub-workflow.

Eager Workflows

Eager workflows (@eager) allow for fully imperative execution where every task call is immediately dispatched to the Flyte backend (or run locally).

from flytekit import eager, task

@eager
async def my_eager_workflow(x: int) -> int:
    # This task is executed immediately, and the result is awaited
    first_res = await add_one(x=x)
    return await double(x=first_res)

Eager tasks use an internal Controller to manage the execution queue. If an eager workflow fails, flytekit can use an EagerFailureHandlerTask to clean up any orphaned executions that were started by the parent.

Local Testing and Mocking

Flytekit provides utilities for testing tasks without a full Flyte backend. You can call a task directly like a function, or use task_mock to replace its execution logic during a test.

from flytekit.testing import task_mock

@task
def t1(i: int) -> int:
    return i + 1

def test_my_logic():
    with task_mock(t1) as m:
        m.side_effect = lambda i: i * 10
        assert t1(i=2) == 20

The task_mock utility temporarily replaces the execute method of the PythonTask instance, ensuring that the mock is only active within the context manager.