Task authoring and execution

Flyte tasks are the fundamental building blocks of flytekit, representing independent, versioned, and unit-testable units of logic. In flytekit, tasks are primarily authored using the @task decorator, which transforms a standard Python function into a PythonFunctionTask.

Declaring Tasks with @task

When you decorate a Python function with @task, flytekit inspects the function's signature and type hints to automatically generate a TypedInterface. This interface defines the inputs and outputs that the Flyte backend uses to orchestrate data flow.

from flytekit import task

@task
def add_numbers(x: int, y: int) -> int:
    return x + y

Internally, the @task decorator (defined in flytekit/core/task.py) wraps your function in a PythonFunctionTask. It uses transform_function_to_interface to map Python types to Flyte's IDL types.

Configuration and Metadata

You can customize task behavior by passing arguments to the @task decorator. These arguments are stored in a TaskMetadata object (defined in flytekit/core/base_task.py), which controls execution parameters like retries, timeouts, and caching.

from datetime import timedelta
from flytekit import task

@task(
    retries=3,
    timeout=timedelta(minutes=5),
    cache=True,
    cache_version="1.0"
)
def fetch_data(url: str) -> str:
    ...

Key configuration options include:

• retries: The number of times Flyte should attempt to re-run the task if it fails.
• timeout: A datetime.timedelta or integer (seconds) specifying the maximum duration for a single execution.
• cache and cache_version: Enables memoization. If cache=True, you must provide a cache_version. Bumping this version forces a cache miss even if inputs remain the same.
• interruptible: Indicates the task can run on lower-priority (spot) instances.

Core Task Abstractions

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

The Task Base Class

The base_task.Task class is the root of all tasks. it captures the information required by the Flyte IDL TaskTemplate. It defines the core lifecycle methods:

• execute(**kwargs): The actual logic to run.
• pre_execute(user_params): Invoked before execution to set up the environment (e.g., initializing a Spark session).
• dispatch_execute(ctx, input_literal_map): The entry point used by the Flyte runtime to translate Flyte literals into Python native values before calling execute.

PythonTask and PythonAutoContainerTask

PythonTask adds a Python-native interface to the base Task. PythonAutoContainerTask (in flytekit/core/python_auto_container.py) further extends this by automatically capturing container information, such as the Docker image and environment variables, required to run the task on a remote cluster.

PythonFunctionTask

This is the most common task type. It holds a reference to the original Python function in self._task_function. When dispatch_execute is called, it handles the conversion of inputs via the TypeEngine, executes the function, and then converts the results back into Flyte literals.

Execution Modes

Tasks in flytekit can behave differently depending on their execution_mode, defined by the PythonFunctionTask.ExecutionBehavior enum.

Default Execution

In the DEFAULT mode, the task runs as a single unit of work. When called locally, it executes the function directly. When run on a cluster, the Flyte engine invokes the container and calls dispatch_execute.

Dynamic Tasks

A task decorated with @dynamic (which sets execution_mode to DYNAMIC) acts as a workflow generator.

from flytekit import dynamic, task

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

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

When a dynamic task executes, it doesn't just return a value; it returns a DynamicJobSpec. Internally, PythonFunctionTask.compile_into_workflow is called to generate a new workflow graph based on the task's execution, which Flyte then schedules.

Eager Tasks

Eager tasks (using EagerAsyncPythonFunctionTask) allow for more flexible, Pythonic execution where every task invocation creates a stack frame on the Flyte cluster. This is useful for complex conditional logic that is difficult to express in a static workflow graph.

Task Resolvers

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

The default_task_resolver (in flytekit/core/python_auto_container.py) works by:

1. Serialization: Capturing the module name and the task's name within that module.
2. Loading: Using importlib.import_module to re-import the module in the container and retrieve the task object by name.

If you have custom loading requirements (e.g., loading tasks from a database or a dynamic source), you can implement your own resolver by overriding load_task and loader_args.

Local Execution and Testing

One of flytekit's strengths is the ability to run tasks locally. When you call a task object like a function, Task.__call__ triggers local_execute.

# This runs locally without a Flyte cluster
result = add_numbers(x=10, y=20)

During local execution:

1. Inputs are translated to Flyte literals to simulate the remote environment.
2. local_execute checks the LocalTaskCache if caching is enabled.
3. sandbox_execute is called, which eventually runs the user's execute method.
4. Outputs are wrapped back into Promise objects or native Python types.

This process ensures that your tasks behave consistently between your local development environment and the production Flyte cluster.