Task authoring and execution

Flyte tasks are the fundamental building blocks of workflows in flytekit. They represent a single unit of work, characterized by a strong interface (typed inputs and outputs), versioning, and independent executability.

Declaring Tasks

The most common way to define a task is by using the @task decorator from flytekit.core.task. When you decorate a Python function, flytekit creates an instance of PythonFunctionTask (or AsyncPythonFunctionTask for async def functions).

from flytekit import task
import typing

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

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

Internally, the @task decorator uses transform_function_to_interface to extract the Flyte interface from your Python type hints. This interface is then used for type validation and serialization when the task is executed on a Flyte cluster.

Task Configuration

The @task decorator accepts several parameters to control execution behavior, resource allocation, and metadata. These are encapsulated in the TaskMetadata class in flytekit.core.base_task.

from flytekit import task, Resources
from datetime import timedelta

@task(
    retries=3,
    timeout=timedelta(minutes=5),
    requests=Resources(cpu="1", mem="2Gi"),
    limits=Resources(cpu="2", mem="4Gi"),
    cache=True,
    cache_version="1.0",
    environment={"MY_ENV_VAR": "value"}
)
def configured_task(x: int) -> int:
    return x + 1

Key configuration options include:

• retries: Number of times to retry the task on failure.
• timeout: Maximum duration for a single execution.
• cache and cache_version: Enables caching of results based on input values.
• requests and limits: Defines compute resource requirements using the Resources class.
• container_image: Specifies a custom Docker image for the task.

Core Task Abstractions

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

1. Task (flytekit.core.base_task): The base class for all tasks. it captures information defined in the Flyte IDL TaskTemplate.
2. PythonTask (flytekit.core.base_task): A base class for tasks with a Python-native interface. It handles the translation between Flyte literals and Python types.
3. PythonFunctionTask (flytekit.core.python_function_task): The standard implementation for tasks wrapping a Python function. It manages the execution of the user-defined function body.
4. PythonInstanceTask (flytekit.core.python_function_task): Used for tasks that do not have a user-defined function body but have a platform-defined execute method (common in plugins).

Custom Task Types

If you need to create a new task type (e.g., for a new backend plugin), you typically extend PythonTask or PythonFunctionTask and implement the execute method.

from flytekit.core.base_task import PythonTask
from flytekit.core.interface import Interface

class MyCustomTask(PythonTask):
    def __init__(self, name: str, **kwargs):
        super().__init__(
            task_type="my_custom_task",
            name=name,
            interface=Interface(inputs={"x": int}, outputs={"o0": int}),
            **kwargs
        )

    def execute(self, **kwargs) -> Any:
        # Custom execution logic here
        return kwargs["x"] * 2

Task Execution Flow

When a task is invoked, flytekit follows a specific execution path depending on whether it is running locally or on a Flyte cluster.

Local Execution

When you call a task function directly in Python, flytekit triggers local_execute.

1. Input Translation: Inputs are converted to Flyte literals using translate_inputs_to_literals.
2. Caching: If caching is enabled, flytekit checks the LocalTaskCache.
3. Dispatch: The dispatch_execute method is called, which invokes the user's execute method.
4. Output Translation: Results are wrapped back into Promise objects.

Remote Execution

On a Flyte cluster, the entry point is typically pyflyte-execute, which uses a TaskResolverMixin to locate and load the task.

1. Loading: The default_task_resolver imports the module and retrieves the task object.
2. Pre-execution: pre_execute is called to set up the environment (e.g., initializing a Spark session).
3. Execution: dispatch_execute converts the input LiteralMap to Python native values and runs the task's execute method.
4. Post-execution: post_execute allows for cleanup or output modification.

Special Task Behaviors

Dynamic Tasks

A dynamic task is declared using the @dynamic decorator. It acts like a task but produces a workflow at runtime. Internally, PythonFunctionTask handles this by setting execution_mode to ExecutionBehavior.DYNAMIC.

Eager Tasks

Eager tasks (declared with @eager) allow for more flexible, Pythonic execution where task results can be used to decide which subsequent tasks to run. These are implemented via EagerAsyncPythonFunctionTask, which manages a Controller to coordinate executions with the Flyte backend.

Ignoring Outputs

In distributed training scenarios, you might want to signal that a task's outputs should be ignored. You can raise the IgnoreOutputs exception within your task to achieve this.

from flytekit.core.base_task import IgnoreOutputs

@task
def distributed_task():
    # ... perform work ...
    if not is_rank_zero:
        raise IgnoreOutputs("Only rank 0 returns data")
    return "result"

Caching Gotchas

When using the Cache object for configuration, do not pass the deprecated parameters cache_serialize, cache_version, or cache_ignore_input_vars directly to the @task decorator. Doing so will raise a ValueError. Instead, configure these within the Cache object itself.

from flytekit.core.cache import Cache

# Correct usage
@task(cache=Cache(version="1.0", serialize=True))
def my_cached_task(x: int) -> int:
    return x