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 translation between Python's native types and the Flyte IDL, manages execution metadata, and provides hooks for both local and remote execution.

Declaring Tasks

The @task decorator in flytekit.core.task is the primary entry point for authoring tasks. It supports both bare decoration and keyword arguments for configuration.

from flytekit import task
import typing

@task
def simple_task(x: int) -> str:
    return f"Value: {x}"

@task(
    cache=True, 
    cache_version="1.0", 
    retries=3, 
    timeout=3600,
    container_image="my-custom-image:latest"
)
def configured_task(x: int) -> int:
    return x + 1

When you decorate a function, flytekit performs several actions:

1. Interface Detection: It uses transform_function_to_interface to inspect your function's type hints and create a TypedInterface.
2. Metadata Construction: It creates a TaskMetadata object to store settings like retries, cache, and timeout.
3. Plugin Dispatch: It uses TaskPlugins.find_pythontask_plugin to determine which task class to instantiate. For standard functions, this is PythonFunctionTask. For async functions, it uses AsyncPythonFunctionTask.
4. Registration: The task instance is appended to FlyteEntities.entities, making it discoverable during serialization.

Task Configuration and Metadata

The TaskMetadata class in base_task holds the execution parameters for a task. Some parameters have strict validation rules enforced in __post_init__:

• If cache=True, a cache_version must be provided.
• cache_serialize and cache_ignore_input_vars can only be enabled if cache=True.
• timeout can be an int (seconds) or a datetime.timedelta.

Task Class Hierarchy

Flytekit uses a layered class hierarchy to separate concerns between the Flyte IDL, Python interfaces, and containerization.

• Task: The abstract base class in base_task. It captures the core Flyte IDL TaskTemplate information (name, type, interface).
• PythonTask: Adds support for Python-native interfaces. It handles the conversion between Flyte literals and Python objects using the TypeEngine.
• PythonAutoContainerTask: Found in python_auto_container, this class manages the container image and the default command (pyflyte-execute) used to run the task on a cluster.
• PythonFunctionTask: The implementation used for @task-decorated functions. It holds a reference to the user's task_function.

Custom Task Types

While most users use @task, you can create custom task types by subclassing PythonTask or PythonInstanceTask. For example, the Echo task in flytekit.core.task demonstrates a task that simply returns its inputs:

class Echo(PythonTask):
    _TASK_TYPE = "echo"

    def __init__(self, name: str, inputs: Optional[Dict[str, Type]] = None, **kwargs):
        outputs = dict(zip(output_name_generator(len(inputs)), inputs.values()))
        super().__init__(
            task_type=self._TASK_TYPE,
            name=name,
            interface=Interface(inputs=inputs, outputs=outputs),
            **kwargs,
        )

    def execute(self, **kwargs) -> Any:
        values = list(kwargs.values())
        return values[0] if len(values) == 1 else tuple(values)

Execution Lifecycle

Flytekit tasks follow a specific execution flow whether running locally or on a remote cluster.

Local Execution

When you call a task directly in a Python script, Task.__call__ invokes flyte_entity_call_handler. This eventually triggers local_execute, which:

1. Translates native Python inputs into Flyte literals using translate_inputs_to_literals.
2. Checks the LocalTaskCache if caching is enabled.
3. Calls sandbox_execute, which sets up a local execution context.
4. Wraps the results back into Promise objects or a VoidPromise.

Dispatch Execute

The dispatch_execute method in PythonTask is the core execution engine used during runtime (both local and remote). It follows these steps:

1. pre_execute: Modifies the ExecutionParameters (e.g., setting up a Spark session).
2. Input Translation: Converts the input_literal_map into native Python kwargs via _literal_map_to_python_input.
3. execute: Invokes the actual user function (or the overridden execute method).
4. post_execute: Allows for cleanup or output modification. If a task raises IgnoreOutputs, the results are discarded.
5. Output Translation: Converts Python return values back into a LiteralMap via _output_to_literal_map.

Advanced Task Behaviors

Dynamic Tasks

A dynamic task is declared using the @dynamic decorator, which is a shortcut for @task(execution_mode=ExecutionBehavior.DYNAMIC).

When dispatch_execute runs a dynamic task, it calls dynamic_execute. Instead of returning simple values, it invokes compile_into_workflow to produce a DynamicJobSpec. This spec contains a dynamically generated workflow that Flyte Propeller then executes.

Eager Workflows

Eager workflows, declared with @eager, allow for fully dynamic execution where Python constructs like loops and conditionals can control Flyte entity execution.

The EagerAsyncPythonFunctionTask uses a Controller (found in worker_queue) to manage sub-executions. If an eager workflow fails, flytekit uses an EagerFailureHandlerTask to terminate any orphaned sub-executions that might still be running on the cluster.

Task Resolvers

When a task runs on a cluster, the container needs to know how to find and load the Python task object. This is handled by TaskResolverMixin. The default_task_resolver works by:

1. Recording the module and function name during serialization (loader_args).
2. Importing the module and retrieving the task attribute during execution (load_task).

Note that the default resolver cannot handle nested or local functions because they are not importable at the module level. If you define a task inside another function, flytekit will raise a ValueError unless you are in a test module (starting with test_).