Skip to main content

Task authoring and execution

Flyte tasks are the fundamental building blocks of Flytekit. They represent a single unit of execution that is versioned, strongly typed, and independently executable. In Flytekit, tasks are typically defined using the @task decorator, which transforms a standard Python function into a PythonFunctionTask.

Declaring Tasks​

The most common way to declare a task is by decorating a Python function with @task. Flytekit uses Python type hints to automatically infer the task's interface (inputs and outputs).

from flytekit import task
import typing

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

@task
def add_numbers(x: int, y: int) -> int:
    return x + y

When you decorate a function, Flytekit creates an instance of flytekit.core.python_function_task.PythonFunctionTask. This class captures the function's signature and metadata, allowing it to be serialized and executed on the Flyte platform.

Task Metadata and Configuration​

You can configure task behavior by passing arguments to the @task decorator. These arguments are stored in a TaskMetadata object (defined in flytekit.core.base_task).

Common configuration options include:

  • Caching: Enable caching to avoid re-running tasks with the same inputs.
@task(cache=True, cache_version="1.0")
def expensive_computation(data: typing.List[int]) -> int:
    return sum(data)
  • Retries: Specify how many times Flyte should retry the task on failure.
@task(retries=3)
def flaky_task(x: int) -> int:
    return x + 1
  • Resources: Request specific CPU, memory, or GPU resources using flytekit.Resources.
from flytekit import Resources

@task(requests=Resources(cpu="2", mem="500Mi"), limits=Resources(cpu="4", mem="1Gi"))
def resource_intensive_task(x: int) -> int:
    return x * 2
  • Timeouts: Set a maximum duration for task execution.
import datetime

@task(timeout=datetime.timedelta(minutes=5))
def time_limited_task(x: int) -> int:
    return x

Core Task Abstractions​

Flytekit provides a hierarchy of classes to handle different task types and execution modes:

  1. Task (flytekit.core.base_task.Task): The base class for all tasks. It captures the Flyte IDL TaskTemplate information but lacks Python-native interfaces.
  2. PythonTask (flytekit.core.base_task.PythonTask): A subclass of Task that adds a Python-native Interface. It handles the translation between Python types and Flyte's internal literal system via the TypeEngine.
  3. PythonFunctionTask (flytekit.core.python_function_task.PythonFunctionTask): The standard class for tasks defined via a Python function. It implements execute() by calling the decorated function.
  4. AsyncPythonFunctionTask: Used when the decorated function is defined with async def.

The Execution Flow​

When a task is executed (either locally or on a cluster), Flytekit follows a structured dispatch process:

  1. local_execute: Used during local development. It translates native Python inputs into Flyte Literal objects, checks the local cache, and calls sandbox_execute.
  2. dispatch_execute: The entry point for both local and remote execution. It handles pre_execute hooks (e.g., setting up Spark sessions), converts LiteralMap inputs back to Python native types, and invokes the user's code.
  3. execute: The method that actually runs the user-defined logic. For PythonFunctionTask, this simply calls the wrapped function.

Task Resolvers​

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

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

  1. Identifying the module and function name during serialization.
  2. Generating a command like pyflyte-execute --resolver flytekit.core.python_auto_container.default_task_resolver -- task-module my_module task-name my_task.
  3. Importing the module and retrieving the task object at runtime using load_task.

Special Task Types​

Dynamic Tasks​

A task can be marked as dynamic by using the @dynamic decorator (which sets execution_mode to DYNAMIC). Dynamic tasks allow you to generate a new workflow structure at runtime based on inputs.

from flytekit import dynamic

@dynamic
def dynamic_task(n: int) -> typing.List[int]:
    return [add_numbers(x=i, y=i) for i in range(n)]

Internally, PythonFunctionTask.dynamic_execute compiles the returned entities into a DynamicJobSpec, which Flyte Propeller then executes as a sub-workflow.

Eager Tasks​

Eager tasks (declared with @eager) allow for more imperative-style execution where tasks are awaited and their results used immediately to decide the next step. These are implemented by EagerAsyncPythonFunctionTask.

Ignore Outputs​

In some scenarios, such as distributed training, a task might not want to return its results to the Flyte engine. You can raise flytekit.core.base_task.IgnoreOutputs to signal that the task's outputs should be ignored.

from flytekit.core.base_task import IgnoreOutputs

@task
def distributed_worker():
    # ... perform work ...
    raise IgnoreOutputs()