Skip to main content

Task authoring and execution

Flyte tasks are the fundamental building blocks of a workflow. In flytekit, tasks are defined as versioned, independently executable units of logic with strong interfaces. While the underlying Flyte engine treats tasks as containerized execution units, flytekit provides high-level abstractions to define these tasks directly in Python.

Declaring Tasks with the @task Decorator​

The most common way to define a task is by using the @task decorator from flytekit.core.task. 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, flytekit automatically:

  1. Detects the Interface: It uses Python type hints to define the task's inputs and outputs.
  2. Captures Metadata: It records the function name and module for serialization.
  3. Enables Local Execution: The resulting object remains callable like a normal function for unit testing.

Requirements for Task Functions​

Task functions must be accessible at the module level so that the Flyte container can re-import and execute them. The PythonFunctionTask constructor (in flytekit/core/python_function_task.py) enforces this:

  • It cannot be a nested or local function (unless in a test module starting with test_).
  • If you use custom decorators, you must use functools.wraps to preserve the function's identity.

Task Configuration and Metadata​

The @task decorator accepts several parameters to control how the task behaves on the Flyte platform. These settings are encapsulated in the TaskMetadata class found in flytekit/core/base_task.py.

Retries and Timeouts​

You can protect against transient failures or runaway processes by specifying retries and timeouts:

from datetime import timedelta

@task(retries=3, timeout=timedelta(minutes=5))
def flaky_task(x: int) -> int:
    ...

Caching​

Caching allows Flyte to skip execution if a task is called with the same inputs and cache version.

from flytekit.core.cache import Cache

@task(cache=True, cache_version="1.0")
def expensive_computation(data: list[int]) -> int:
    ...

Internally, TaskMetadata validates that if cache=True is set, a cache_version must also be provided. You can also use cache_serialize=True to ensure that concurrent executions with identical inputs run serially to avoid redundant work.

Resource Requests​

Tasks can request specific hardware resources like CPU, memory, or GPUs using the Resources class:

from flytekit import Resources

@task(requests=Resources(cpu="2", mem="500Mi"), limits=Resources(cpu="4", mem="1Gi"))
def resource_intensive_task():
    ...

Task Execution Modes​

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

Default Execution​

In DEFAULT mode, the task function is executed exactly as written. This is the standard behavior for most tasks.

Dynamic Tasks​

Dynamic tasks allow you to generate a workflow structure at runtime based on the task's inputs. You declare these using the @dynamic decorator (which is a specialized version of @task with execution_mode=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)]

When a dynamic task runs:

  1. It executes the user code to produce a DynamicJobSpec.
  2. Flyte Propeller then executes the nodes defined in that spec as a sub-workflow.

Eager Tasks​

Eager tasks (or "eager workflows") allow for more flexible, Pythonic control flow that feels like local execution but runs on the Flyte backend. They are implemented via EagerAsyncPythonFunctionTask. In this mode, every task call inside the eager function creates a separate execution on the cluster, and the eager task awaits the result.

Core Abstractions​

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

  • Task: The base class capturing the Flyte IDL TaskTemplate. It handles the translation between Flyte's type system (Literals) and Python types.
  • PythonTask: A subclass that adds Python-native interfaces, environment variables, and Flyte Decks support.
  • PythonFunctionTask: The implementation for tasks backed by a Python function. It handles the dispatch_execute logic, which is the entry point when the task runs in a container.

The Execution Flow​

When a task is executed on the Flyte platform, the container entry point calls dispatch_execute. The internal flow is:

  1. pre_execute: Sets up the execution environment (e.g., initializing Spark sessions).
  2. Input Translation: Converts LiteralMap inputs from the Flyte engine into Python native types using the TypeEngine.
  3. execute: Invokes the actual Python function.
  4. post_execute: Performs cleanup or output modification.
  5. Output Translation: Converts Python return values back into a LiteralMap to be sent back to the Flyte engine.

Task Resolvers​

When a task runs in a container, Flyte needs to know how to find and load the task object. This is handled by TaskResolverMixin. The default_task_resolver (in flytekit/core/python_auto_container.py) serializes the task's module and name into the container's command-line arguments.

For example, a task t1 in my_module.py might be invoked with:

pyflyte-execute --resolver flytekit.core.python_auto_container.default_task_resolver -- task-module my_module task-name t1

The resolver uses importlib to load the module and retrieve the task instance.