Skip to main content

Task authoring and execution

Flyte tasks are the fundamental building blocks of a workflow. In flytekit, tasks are primarily declared using the @task decorator, which transforms a standard Python function into a PythonFunctionTask. This abstraction handles the complexities of type conversion, metadata management, and execution on the Flyte platform.

Declaring Tasks with @task​

When you want to run a piece of logic as a discrete step in a workflow, you decorate a Python function with @task. This decorator captures the function's signature and uses it to define the task's interface (inputs and outputs).

from flytekit import task

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

Internally, the task decorator in flytekit/core/task.py performs several steps:

  1. Interface Extraction: It calls transform_function_to_interface to inspect the function's type hints and docstrings.
  2. Metadata Construction: It creates a TaskMetadata object (defined in flytekit/core/base_task.py) to store configuration like retries, timeouts, and caching settings.
  3. Plugin Resolution: It checks TaskPlugins to see if a specific task_config was provided that requires a specialized task class (e.g., for Spark or Ray).
  4. Instantiation: It instantiates a PythonFunctionTask (or a subclass like AsyncPythonFunctionTask for async functions) and wraps the original function.

Task Configuration​

The @task decorator accepts numerous parameters to control how the task behaves in a production environment:

  • Caching: Use cache=True and cache_version="1.0" to avoid re-running tasks with the same inputs.
  • Resources: Specify requests and limits using the Resources class to request specific CPU, memory, or GPU allocations.
  • Retries: Set retries=n to automatically retry the task on failure.
  • Container Image: Use container_image to specify a custom Docker image for this specific task, overriding the default image used by the workflow.
from flytekit import task, Resources

@task(
    cache=True, 
    cache_version="1.0", 
    requests=Resources(cpu="2", mem="500Mi"),
    retries=3
)
def heavy_computation(data: list[float]) -> float:
    ...

Core Task Abstractions​

Flytekit uses a hierarchy of classes to represent different types of tasks:

Task (Base Class)​

The Task class in flytekit/core/base_task.py is the lowest-level abstraction. it maps directly to the Flyte IDL TaskTemplate. It does not assume a Python-native interface and is primarily used as a base for tasks that might be implemented in other languages or specialized backends.

PythonTask​

PythonTask extends Task by adding a python_interface. This is the base class for tasks that have a Python-native signature but might not be defined by a single Python function body. For example, SQLTask (found in flytekit/core/base_sql_task.py) inherits from PythonTask because it executes a SQL query rather than a Python function, yet it still needs to define its inputs and outputs using Python types.

PythonFunctionTask​

PythonFunctionTask (in flytekit/core/python_function_task.py) is the most common class users interact with. It is designed to wrap a standard Python Callable.

When a PythonFunctionTask is executed, it follows a specific lifecycle:

  1. pre_execute: Prepares the execution environment (e.g., setting up a Spark session).
  2. execute: Invokes the underlying Python function with the provided inputs.
  3. post_execute: Performs cleanup or output modification.

Task Execution Modes​

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

Default Execution​

In DEFAULT mode, the task runs as a single unit of work. The dispatch_execute method handles the translation of Flyte literals (the platform's data format) into Python-native types before calling the user's function.

Dynamic Execution​

When a task is decorated with @dynamic (which sets the behavior to DYNAMIC), the task function is expected to return a list of other task calls or sub-workflows.

  • Compilation: During execution, the task runs the user code to "discover" the workflow structure.
  • DynamicJobSpec: Instead of returning data, it returns a DynamicJobSpec containing the generated workflow template, which Flyte Propeller then executes.

Eager Execution​

Eager tasks (using EagerAsyncPythonFunctionTask) allow for more flexible, Pythonic control flow where task results can be inspected immediately to decide the next step, similar to a standard Python program but with each call being offloaded to the Flyte cluster.

Task Resolvers​

When a task is executed on a remote Flyte cluster, the container needs to know how to find and instantiate the specific task object. This is handled by TaskResolverMixin.

The default_task_resolver (in flytekit/core/python_auto_container.py) works by:

  1. loader_args: Capturing the module name and the task's variable name during serialization.
  2. load_task: Importing the module and retrieving the task object by name when the container starts up.

You can implement a custom resolver by inheriting from TaskResolverMixin and overriding load_task and loader_args if you need to load tasks from a database or a dynamic source.