Task authoring and execution

Flyte tasks are the fundamental building blocks of execution in flytekit. They represent a single unit of work, versioned and strongly typed, that can be executed independently or as part of a workflow.

Declaring Tasks

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

from flytekit import task

@task
def add_one(x: int) -> int:
    return x + 1

When you decorate a function with @task, flytekit automatically infers the task's interface (inputs and outputs) from the function's type hints using the transform_function_to_interface utility. This ensures that the task is compatible with Flyte's type system.

Task Metadata and Configuration

You can configure task behavior by passing arguments to the @task decorator. These arguments are captured in the TaskMetadata class, which controls execution properties like retries, timeouts, and caching.

from flytekit import task
from datetime import timedelta

@task(
    retries=3,
    timeout=timedelta(minutes=5),
    cache=True,
    cache_version="1.0",
    interruptible=True
)
def heavy_computation(data: list[int]) -> int:
    return sum(data)

Key configuration options include:

• Caching: Set cache=True and provide a cache_version. Flyte uses these to avoid re-running tasks with identical inputs. Internally, TaskMetadata validates that cache_version is present if cache is enabled.
• Retries: The retries parameter (an integer) is converted into a RetryStrategy model for the Flyte backend.
• Interruptible: Indicates the task can run on lower-cost, pre-emptible nodes.
• Resources: Use requests and limits to specify CPU, memory, and GPU requirements.

Core Task Abstractions

Flytekit uses a class hierarchy to manage different types of tasks:

1. Task: The base class in flytekit.core.base_task. It captures the Flyte IDL TaskTemplate information, such as task_type, name, and interface.
2. PythonTask: Inherits from Task and adds support for Python native interfaces. It handles the translation between Flyte literals and Python types via the TypeEngine.
3. PythonFunctionTask: The most common implementation, used for tasks defined as Python functions. It stores the actual task_function and handles its execution.

Custom Task Types

For specialized execution environments (like SQL queries or Spark jobs), flytekit provides specialized subclasses. You can register these using the TaskPlugins factory:

# Internal mechanism for finding the right task class based on config
task_plugin = TaskPlugins.find_pythontask_plugin(type(task_config))

Task Execution Flow

Task execution in flytekit follows a structured lifecycle managed by the dispatch_execute method in PythonTask.

1. Pre-execution

Before the user code runs, pre_execute is called. This method can mutate the ExecutionParameters to set up the environment, such as initializing a Spark session or configuring secrets.

2. Input Translation

Flytekit receives inputs as a LiteralMap (Flyte's internal data format). The _literal_map_to_python_input method uses the TypeEngine to convert these into Python-native objects that your function can process.

3. User Code Execution

The execute method invokes the underlying Python function. If the task is running locally, flytekit catches exceptions and provides context-aware error messages.

4. Output Translation

After the function returns, _output_to_literal_map converts the Python return values back into a LiteralMap. This involves:

• Validating the output against the expected interface.
• Using TypeEngine.async_to_literal to handle data serialization (including offloading large files to remote storage).

Dynamic Workflows

Dynamic workflows, declared with the @dynamic decorator, are a special type of task where the execution_mode is set to ExecutionBehavior.DYNAMIC.

from flytekit import dynamic

@dynamic
def my_dynamic_task(n: int) -> list[int]:
    return [add_one(x=i) for i in range(n)]

Unlike a standard task, a dynamic task's execute method calls dynamic_execute. Instead of returning simple values, it compiles a new workflow at runtime based on the inputs provided. This generated workflow is returned as a DynamicJobSpec, which the Flyte engine then executes as a subworkflow.

Local Execution

Tasks can be executed locally just like regular Python functions. When you call a task object (e.g., my_task(x=10)), the flyte_entity_call_handler determines the execution context. If it's a local call, it triggers local_execute, which simulates the Flyte runtime by:

1. Translating native Python inputs to literals.
2. Checking the LocalTaskCache if caching is enabled.
3. Running sandbox_execute to invoke the task logic.
4. Wrapping the results in Promise objects.