Task authoring and execution

Flyte tasks are the fundamental building blocks of flytekit. They represent a single unit of work, characterized by a versioned, strongly-typed interface and independent executability. In flytekit, tasks are primarily authored 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 decorating a Python function with @task. flytekit uses Python type hints to automatically infer the task's input and output types, which are then mapped to the Flyte IDL.

import typing
from flytekit import task

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

@task
def add_numbers(a: int, b: int) -> int:
    return a + b

When you call greet(name="Flyte") in a local script, flytekit executes the function directly. However, when used within a @workflow, the call returns a Promise object, which represents a future value that will be computed on the Flyte platform.

Task Metadata and Configuration

The @task decorator accepts several parameters to control the task's behavior, such as retries, caching, and resource requirements. These are encapsulated in the TaskMetadata class internally.

from datetime import timedelta
from flytekit import task, Resources

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

Key configuration options include:

• cache: Enables caching of results. If True, cache_version must also be provided.
• retries: The number of times to retry the task on failure.
• timeout: The maximum duration for a single execution.
• requests / limits: Define the Resources (CPU, memory, storage) required for the task.

Core Task Abstractions

The task system is built on a hierarchy of classes in flytekit.core.base_task:

1. Task: The base class for all tasks. it captures the TaskTemplate information required by the Flyte backend, including the task_type, name, and interface.
2. PythonTask: A subclass of Task that adds a python_interface. It handles the translation between Flyte literals and Python native types using the TypeEngine.
3. PythonFunctionTask: The implementation used for functions decorated with @task. It stores the original task_function and manages its execution.

Execution Flow

When a task is executed (either locally or on the cluster), it follows a specific lifecycle managed by dispatch_execute in PythonTask:

1. pre_execute: Prepares the execution environment (e.g., setting up Spark sessions).
2. Input Translation: Converts the LiteralMap of inputs from the Flyte engine into Python native objects using _literal_map_to_python_input.
3. execute: Invokes the actual user-defined function with the translated inputs.
4. post_execute: Allows for cleanup or output modification.
5. Output Translation: Converts the Python return values back into a LiteralMap via _output_to_literal_map.

Dynamic Tasks

Dynamic tasks allow you to define the execution graph at runtime based on input data. They are declared using the @dynamic decorator, which is a specialized PythonFunctionTask with ExecutionBehavior.DYNAMIC.

from flytekit import task, dynamic

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

@dynamic
def dynamic_workflow(count: int) -> typing.List[int]:
    return [process_item(item=i) for i in range(count)]

Internally, when a dynamic task runs on the Flyte platform, it executes the function body to produce a DynamicJobSpec. This spec contains a set of nodes (tasks or sub-workflows) that the Flyte engine then schedules and executes.

Map Tasks

Map tasks allow you to run a single task across a list of inputs in parallel. This is more efficient than a dynamic workflow for simple "map" operations because it avoids the overhead of creating a full dynamic job spec.

from flytekit import task, map_task

@task
def square(x: int) -> int:
    return x * x

@workflow
def my_map_workflow(inputs: typing.List[int]) -> typing.List[int]:
    return map_task(square)(x=inputs)

Flytekit provides map_task which, depending on the backend configuration, uses either the legacy MapPythonTask or the newer ArrayNodeMapTask.

Local Execution and Testing

One of flytekit's strengths is the ability to run tasks locally. When you call a task function directly, Task.__call__ triggers the flyte_entity_call_handler. If it detects a local execution context, it invokes local_execute, which performs the same input/output translation as the remote path but runs the code in your local process.

if __name__ == "__main__":
    # This runs locally without needing a Flyte cluster
    result = greet(name="Developer")
    print(result)

For tasks that require complex environments, you can use sandbox_execute to simulate the execution environment more closely, including local caching if enabled via LocalTaskCache.