Skip to main content

Task authoring and execution

Flyte tasks are the fundamental building blocks of a workflow. In flytekit, tasks are typically defined by decorating a Python function, which allows the framework to automatically handle type conversion, metadata tracking, and execution environment setup.

Declaring Tasks​

The most common way to create a task is using the @task decorator from flytekit.core.task. When you apply this decorator to a function, flytekit inspects the function's type hints and docstrings to construct a PythonFunctionTask.

from flytekit import task

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

Internally, the @task decorator acts as a factory. It creates a TaskMetadata object to store configuration and then instantiates a subclass of PythonFunctionTask. If the decorated function is asynchronous (async def), flytekit automatically uses AsyncPythonFunctionTask instead.

Interface Inference​

Flytekit relies on Python type hints to define the task's interface. The transform_function_to_interface function in flytekit.core.interface parses the function signature to create a TypedInterface, which Flyte uses to ensure type safety across different tasks in a workflow.

Task Configuration​

You can configure task behavior by passing arguments to the @task decorator. These settings are captured in the TaskMetadata class within flytekit.core.base_task.

Caching and Retries​

Caching allows Flyte to skip execution if a task has already been run with the same inputs.

from flytekit import task

@task(cache=True, cache_version="1.0", retries=3)
def fetch_data(query: str) -> str:
    ...
  • cache: Enables caching when set to True.
  • cache_version: A string that, when changed, invalidates previous cache entries. This is required if cache=True.
  • retries: The number of times Flyte should attempt to re-run the task upon failure.

Resource Requests and Limits​

You can specify the compute resources required for a task using the Resources object.

from flytekit import task, Resources

@task(
    requests=Resources(cpu="2", mem="500Mi"),
    limits=Resources(cpu="4", mem="1Gi")
)
def memory_intensive_task(data: list) -> int:
    ...

Core Task Abstractions​

Flytekit uses a hierarchical class structure to manage different types of tasks:

  1. Task: The base class in flytekit.core.base_task. it captures the language-agnostic information required by the Flyte IDL, such as the 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 primary class for tasks defined via functions. It implements execute() by calling the decorated Python function.

Custom Task Plugins​

For specialized execution environments (like Spark or SQL), flytekit provides a plugin system. You can register custom task types using TaskPlugins.register_pythontask_plugin. This allows a specific task_config object to trigger a specialized PythonFunctionTask subclass.

Execution Flow​

When a task is executed, flytekit follows a structured lifecycle managed by the dispatch_execute method in PythonTask:

  1. pre_execute: Prepares the execution environment (e.g., setting up a Spark session).
  2. Input Translation: Converts Flyte LiteralMap inputs into Python-native values using _literal_map_to_python_input.
  3. execute: Invokes the actual user code or plugin logic.
  4. post_execute: Performs cleanup or output modification.
  5. Output Translation: Converts Python return values back into a Flyte LiteralMap via _output_to_literal_map.

Local Execution​

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

# Local execution for testing
result = square(n=5)
assert result == 25

Specialized Task Types​

Dynamic Tasks​

Dynamic tasks allow you to generate a new workflow structure at runtime based on input data. You define these by setting execution_mode=PythonFunctionTask.ExecutionBehavior.DYNAMIC (often via the @dynamic decorator).

Eager Tasks​

Eager tasks (EagerAsyncPythonFunctionTask) allow for more flexible, imperative-style execution where Python code acts as the orchestrator, spawning sub-tasks on the Flyte cluster as it runs. This is indicated by is_eager=True in the TaskMetadata.