Task authoring and execution

Flyte tasks are the fundamental building blocks of flytekit. They represent a single unit of execution with a strongly typed interface, allowing for independent execution, versioning, and unit testing.

Declaring Tasks

The primary way to define a task in flytekit is by using the @task decorator on a Python function. This decorator transforms a standard Python function into a PythonFunctionTask object.

from flytekit import task

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

When you decorate a function with @task, flytekit automatically:

1. Infers the Interface: It uses Python type hints to determine the input and output types, mapping them to Flyte's internal type system.
2. Captures Metadata: It records the function name and module to allow the Flyte platform to re-instantiate and execute the task in a remote container.
3. Enables Local Execution: You can call the decorated function directly in your local environment for testing, and flytekit will handle the execution.

Task Configuration

The @task decorator accepts several parameters to configure the task's behavior on the Flyte platform.

from datetime import timedelta
from flytekit import task, Resources

@task(
    retries=3,
    cache=True,
    cache_version="1.0",
    timeout=timedelta(minutes=5),
    requests=Resources(cpu="1", mem="1Gi"),
    limits=Resources(cpu="2", mem="2Gi"),
    interruptible=True
)
def resource_intensive_task(x: int) -> int:
    return x * 2

Key configuration options include:

• retries: The number of times to retry the task on failure.
• cache and cache_version: Enables caching of results. If a task is called with the same inputs and the same cache_version, Flyte will reuse the previous result instead of re-running the task.
• timeout: The maximum duration for a single execution of the task.
• requests and limits: Specify the compute resources (CPU, memory, storage) required for the task using the Resources class.
• interruptible: Indicates if the task can run on lower-priority, pre-emptible nodes (like AWS Spot Instances).

Task Abstractions

Flytekit uses a class hierarchy to manage different types of tasks, starting from a language-agnostic base and moving to Python-specific implementations.

The Task Base Class

The base_task.Task class is the root of all tasks. it captures information required by the Flyte IDL (Interface Definition Language), such as the task_type, name, and interface. It defines the dispatch_execute method, which is the entry point for execution both locally and on the Flyte platform.

PythonTask

The base_task.PythonTask class extends Task for tasks that have a Python-native interface. It introduces task_config, a generic parameter used by plugins to provide specialized configuration (e.g., for Spark or SQL tasks).

PythonFunctionTask

The python_function_task.PythonFunctionTask class is the most common implementation. It specifically handles tasks defined by a Python function. Internally, it manages:

• Interface Transformation: Uses transform_function_to_interface to convert Python signatures into Flyte TypedInterface models.
• Execution Logic: Implements the execute method by simply calling the underlying Python function with the provided keyword arguments.

Task Execution Flow

When a task is executed, flytekit manages the transition between Flyte's internal data representation (Literals) and Python-native types.

Local Execution

When you call a task locally, flytekit invokes Task.local_execute. This method:

1. Translates Python-native inputs into Flyte Literal objects.
2. Checks the local cache if caching is enabled.
3. Calls dispatch_execute, which eventually runs the user's Python code.
4. Translates the results back into Promise objects (or native values) for use in workflows.

Remote Execution

On the Flyte platform, the container starts by running pyflyte-execute. This command uses a Task Resolver to find and load the task.

The TaskResolverMixin defines how tasks are serialized and reloaded:

• loader_args: Generates the command-line arguments (like module name and function name) needed to identify the task.
• load_task: Uses those arguments to import the module and retrieve the task object at runtime.

The DefaultTaskResolver in python_auto_container.py handles standard Python functions by recording the module and function name:

# Example of arguments generated by DefaultTaskResolver
# --resolver flytekit.core.python_auto_container.default_task_resolver 
# -- task-module my_module task-name my_task

Specialized Task Types

Flytekit supports several specialized execution behaviors through subclasses of PythonFunctionTask:

• Dynamic Tasks: Created using the @dynamic decorator. These tasks return a DynamicJobSpec, allowing them to generate a new sub-workflow at runtime based on input data.
• Eager Tasks: Implemented by EagerAsyncPythonFunctionTask. These allow for "eager" execution where Python code acts as the orchestrator, invoking other Flyte entities and awaiting their results in a style similar to standard async/await code.
• Reference Tasks: Defined using @reference_task. These are pointers to tasks that already exist on a Flyte deployment, allowing you to call them without having the source code locally.

from flytekit import reference_task

@reference_task(
    project="flytesnacks",
    domain="development",
    name="core.recipes.add_one",
    version="v1"
)
def remote_add_one(x: int) -> int:
    ...