Skip to main content

Task authoring and execution

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

Defining Tasks with the @task Decorator​

The primary way to author a task in flytekit is by using the @task decorator from flytekit.core.task. This decorator transforms a standard Python function into a PythonFunctionTask.

from flytekit import task
import typing

@task
def greet(name: str) -> str:
    return f"Hello, {name}!"

When you decorate a function with @task, flytekit automatically detects the interface (inputs and outputs) using Python type hints. This information is used to create a TypedInterface that the Flyte backend uses for data validation and orchestration.

Task Configuration​

The @task decorator accepts several parameters to control the execution behavior and resource allocation of the task:

  • retries: The number of times to retry the task on failure.
  • timeout: A datetime.timedelta or integer (seconds) representing the maximum duration for a single execution.
  • requests / limits: Define compute resource requirements (CPU, memory, storage) using the Resources class.
  • container_image: Specify a custom Docker image for this specific task, overriding the default image.
  • environment: A dictionary of environment variables to be set during execution.
  • cache: Enables caching of results. You can pass a boolean or a Cache object for advanced configuration.

Example of a configured task:

from datetime import timedelta
from flytekit import task, Resources

@task(
    retries=3,
    timeout=timedelta(minutes=5),
    requests=Resources(cpu="2", mem="500Mi"),
    limits=Resources(cpu="4", mem="1Gi"),
    cache=True,
    cache_version="1.0"
)
def heavy_computation(data: list[int]) -> int:
    return sum(data)

Core Task Abstractions​

Internally, flytekit uses a hierarchy of classes to manage task behavior.

The Task Base Class​

The base_task.Task class is the root of all tasks. It captures the metadata and interface required by the Flyte IDL (Interface Definition Language). It defines the dispatch_execute method, which is the entry point for execution both locally and on the Flyte platform.

PythonTask​

base_task.PythonTask extends Task to handle tasks with a Python-native interface. It manages the translation between Flyte's internal Literal types and Python types using the TypeEngine.

PythonFunctionTask​

python_function_task.PythonFunctionTask is the implementation used by the @task decorator. It wraps the user's Python function and handles its execution.

Task Execution Flow​

When a task is executed, it goes through several stages managed by the Task and PythonTask classes:

  1. pre_execute: Invoked before the task body. This is used to set up execution parameters or modify the context (e.g., initializing a Spark session).
  2. Input Translation: The dispatch_execute method calls _literal_map_to_python_input to convert Flyte LiteralMap inputs into Python-native keyword arguments.
  3. execute: The actual user-defined function is called with the translated inputs.
  4. post_execute: Invoked after the task body. It can be used for cleanup or to alter the return values.
  5. Output Translation: The _output_to_literal_map method converts the Python return values back into a Flyte LiteralMap.

Local Execution​

When you call a task function directly in a Python script, flytekit triggers local_execute. This bypasses the Flyte backend and runs the code locally, while still performing type validation and (optionally) local caching.

# This triggers local_execute
result = greet(name="Flyte")

Advanced Task Types​

Dynamic Tasks​

Dynamic tasks are defined using the @dynamic decorator. They allow you to generate a sub-workflow at runtime based on the task's inputs. Internally, this sets the execution_mode of a PythonFunctionTask to ExecutionBehavior.DYNAMIC.

from flytekit import dynamic

@dynamic
def my_dynamic_task(n: int) -> list[str]:
    return [greet(name=f"User {i}") for i in range(n)]

Eager Tasks​

Eager tasks (using @eager) allow for more flexible, Pythonic execution where task results can be used to decide subsequent task calls within the same execution context. These are implemented via EagerAsyncPythonFunctionTask and require an asynchronous environment.

Reference Tasks​

ReferenceTask allows you to point to a task that is already registered on a Flyte cluster without providing the implementation. This is useful for composing workflows across different projects or teams.

from flytekit import reference_task

@reference_task(
    project="flytesnacks",
    domain="development",
    name="core.greet",
    version="v1"
)
def remote_greet(name: str) -> str:
    ...

Task Plugins and Customization​

Flytekit supports extensibility through TaskPlugins. By providing a task_config to the @task decorator, you can trigger specialized behavior handled by specific plugin classes (e.g., Spark, SQL, Pod).

The TaskPlugins factory (in flytekit.core.task) maps configuration types to their corresponding PythonFunctionTask implementation. When a task is defined with a specific config, flytekit looks up the registered plugin to handle serialization and execution.