Task authoring and execution

Flyte tasks are the fundamental building blocks of execution in flytekit. They represent a single unit of work, characterized by a strong interface (typed inputs and outputs) and declarative metadata.

Declaring Tasks with the @task Decorator

The primary way to define a task is by using the @task decorator on a Python function. This creates a PythonFunctionTask instance that automatically infers its interface from the function's type hints and docstring.

from flytekit import task

@task
def square(n: int) -> int:
    """
    A simple task that squares an integer.
    """
    return n * n

When you apply @task, flytekit uses transform_function_to_interface (from flytekit.core.interface) to parse the function signature. This ensures that the task's inputs and outputs are strictly typed according to Flyte's type system.

Task Metadata and Configuration

The @task decorator accepts several parameters to control execution behavior, which are encapsulated in the TaskMetadata class (found in flytekit.core.base_task).

• Caching: Enable caching by passing cache=True and a cache_version. Caching is managed by the Cache class in flytekit.core.cache.
• Retries: Specify the number of retries on failure using the retries parameter.
• Timeouts: Set a maximum execution duration using timeout (either an integer in seconds or a datetime.timedelta).
• Resources: Request specific compute resources like CPU, memory, and GPU using the Resources class.

from flytekit import task, Resources
from datetime import timedelta

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

Internally, TaskMetadata validates these parameters. For instance, it raises a ValueError if cache=True is set without a cache_version.

Core Task Abstractions

Flytekit organizes task behavior through a hierarchy of classes in flytekit.core.base_task and flytekit.core.python_function_task.

The Base Task Class

The Task class is the root abstraction. It captures the Flyte IDL TaskTemplate and defines the basic lifecycle methods:

• pre_execute: Prepares the execution context (e.g., setting up Spark sessions).
• execute: The actual logic of the task.
• dispatch_execute: Translates Flyte literals to Python native types, calls execute, and translates results back to literals.

PythonTask and PythonFunctionTask

PythonTask extends Task to provide a Python-native Interface. PythonFunctionTask further specializes this by wrapping a user-defined Python function.

When a PythonFunctionTask is executed, it follows this flow in dispatch_execute:

1. Input Translation: Converts the LiteralMap from the Flyte engine into Python native values using _literal_map_to_python_input.
2. Execution: Invokes the decorated function with the native inputs.
3. Output Translation: Converts the function's return values back into a LiteralMap via _output_to_literal_map.

Task Execution Modes

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

Local Execution

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

# Local execution
result = square(n=5)
print(result) # Outputs 25

Dynamic Tasks

A task can be marked as dynamic by using the @dynamic decorator (which sets execution_mode to DYNAMIC). Dynamic tasks allow you to generate new tasks or workflows at runtime based on input data.

from flytekit import dynamic

@dynamic
def dynamic_subworkflow(n: int) -> list[int]:
    return [square(n=i) for i in range(n)]

In dynamic_execute, flytekit compiles the generated entities into a DynamicJobSpec, which the Flyte propeller then executes as a sub-workflow.

Eager Tasks

Eager tasks (defined via EagerAsyncPythonFunctionTask) allow for more flexible, imperative-style execution where Python code acts as the orchestrator. Unlike standard tasks, eager tasks can await the results of other tasks directly within the function body.

Task Resolvers

For a task to run on a remote Flyte cluster, the container needs to know how to locate and load the task code. This is handled by the TaskResolverMixin. The default_task_resolver (in flytekit.core.python_auto_container) identifies tasks by their module path and function name.

When a task is serialized, flytekit generates a command like:

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

The resolver's load_task method uses these arguments to import the module and retrieve the task object at runtime.