Task authoring and execution

Flyte tasks are the fundamental building blocks of Flytekit, representing versioned, independently executable units of code with strong interfaces. In Flytekit, tasks are primarily declared using the @task decorator, which transforms a standard Python function into a PythonFunctionTask.

Declaring Tasks

The most common way to define a task is by applying the @task decorator to a Python function. Flytekit uses the function's type hints to automatically infer the task's input and output interfaces.

from flytekit import task
import typing

@task
def process_data(x: int, y: typing.Dict[str, str]) -> str:
    """
    A simple task that takes an integer and a dictionary and returns a string.
    """
    return f"Processed {x} with {y}"

When you call this task within a workflow, it does not execute immediately. Instead, it returns a Promise object, which represents a future value that will be produced when the task runs on the Flyte platform.

Task Configuration

The @task decorator accepts numerous parameters to control execution behavior, resource allocation, and metadata.

Execution Control

You can configure how Flyte handles task failures and execution limits:

• retries: Number of times to retry the task on failure (default is 0).
• timeout: A datetime.timedelta or integer (seconds) specifying the maximum duration for a single execution.
• interruptible: A boolean indicating if the task can be scheduled on lower-priority, pre-emptible nodes.

from datetime import timedelta

@task(retries=3, timeout=timedelta(minutes=5), interruptible=True)
def robust_task(a: int) -> int:
    return a + 1

Caching

Caching allows Flyte to skip task execution if the same inputs have been processed before. While legacy parameters like cache_version exist, the recommended approach is using the Cache object from flytekit.core.cache.

from flytekit import task
from flytekit.core.cache import Cache

@task(cache=Cache(version="v1", serialize=True, ignored_inputs=("secret_key",)))
def cached_task(data: str, secret_key: str) -> str:
    return f"Result for {data}"

Internally, TaskMetadata (found in flytekit/core/base_task.py) stores these settings. The Cache object's serialize=True setting ensures that concurrent executions with identical inputs run serially to avoid redundant work.

Resource Management

Tasks can request specific hardware resources using the Resources object:

from flytekit import task, Resources

@task(
    requests=Resources(cpu="2", mem="500Mi"),
    limits=Resources(cpu="4", mem="1Gi"),
    accelerator=None # Can be used for GPU allocation
)
def resource_intensive_task(x: int) -> int:
    return x * x

Task Abstractions and Hierarchy

Flytekit implements a class hierarchy to manage different task behaviors while maintaining a consistent interface.

The Task Base Class

The Task class in flytekit/core/base_task.py is the root of all tasks. It captures the Flyte IDL TaskTemplate specification. It is responsible for:

• Defining the TypedInterface (inputs and outputs).
• Managing TaskMetadata.
• Handling local_execute, which manages the transition from Python native values to Flyte literals during local testing.

PythonTask and PythonFunctionTask

PythonTask extends the base Task to support Python-native interfaces. The PythonFunctionTask (in flytekit/core/python_function_task.py) is the specific implementation used for tasks defined via the @task decorator.

When a PythonFunctionTask is executed, it follows this internal flow:

1. pre_execute: Sets up the execution environment (e.g., initializing ExecutionParameters).
2. dispatch_execute: Translates Flyte LiteralMap inputs into Python native types using the TypeEngine.
3. execute: Invokes the actual user-defined Python function.
4. post_execute: Handles output translation back to Flyte literals and generates Decks if enabled.

Specialized Task Types

Async and Eager Tasks

Flytekit supports asynchronous execution through AsyncPythonFunctionTask. If you decorate an async def function with @task, Flytekit automatically selects this task type.

For more complex dynamic logic, the @eager decorator (implemented by EagerAsyncPythonFunctionTask) allows you to run tasks and workflows using Python's async/await syntax, where each call to a Flyte entity triggers a remote execution on the Flyte cluster.

Reference Tasks

If you need to call a task that is already registered on a Flyte cluster without having its source code available, use @reference_task:

from flytekit import reference_task

@reference_task(
    project="flytesnacks",
    domain="development",
    name="core.recipes.simple.add",
    version="v1"
)
def remote_add(a: int, b: int) -> int:
    ...

Local Execution and Testing

Tasks can be executed locally just like regular Python functions. When you call my_task(a=1), Flytekit's flyte_entity_call_handler detects the execution context. If it's a local execution, it triggers local_execute, which performs type validation and optionally checks the LocalTaskCache before running the function.

If a task needs to signal that its outputs should be ignored (e.g., in distributed training scenarios), it can raise the IgnoreOutputs exception defined in flytekit/core/base_task.py.