Task authoring and execution

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

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 with @task, flytekit performs several internal actions:

1. Interface Detection: It uses transform_function_to_interface to inspect the function's type hints and docstrings, creating a TypedInterface that Flyte understands.
2. Task Instantiation: It creates an instance of PythonFunctionTask (or a specialized subclass if a task_config is provided).
3. Metadata Association: It bundles execution settings like retries and timeouts into a TaskMetadata object.

Task Configuration and Metadata

You can configure task behavior by passing arguments to the @task decorator. These settings are captured in the TaskMetadata class found in flytekit.core.base_task.

from datetime import timedelta
from flytekit import task, Resources

@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"
)
def resource_intensive_task(data: list[int]) -> int:
    return sum(data)

Key metadata attributes include:

• retries: The number of times Flyte should retry the task on failure.
• timeout: The maximum duration for a single execution. Internally, TaskMetadata.__post_init__ ensures this is a datetime.timedelta.
• cache and cache_version: Enables caching of results. If cache=True, a cache_version must be provided, or TaskMetadata will raise a ValueError.
• interruptible: Indicates if the task can run on lower-priority, pre-emptible nodes (e.g., AWS Spot Instances).

Task Execution Flow

Flyte tasks follow different execution paths depending on whether they are running locally or on a remote Flyte cluster.

Local Execution

When you call a task function directly in a Python script, flytekit triggers Task.__call__, which routes to local_execute.

1. Input Translation: local_execute converts Python native inputs into Flyte literals using translate_inputs_to_literals.
2. Cache Check: If caching is enabled and LocalConfig.cache_enabled is true, it checks LocalTaskCache for a hit.
3. Sandbox Execution: If there's a cache miss, it calls sandbox_execute, which invokes the user's code via dispatch_execute.
4. Output Wrapping: The results are wrapped back into Promise objects for use in workflows.

Remote Execution

On a Flyte cluster, the entry point is dispatch_execute. This method is responsible for:

1. Pre-execution: Running pre_execute to set up the environment (e.g., initializing a Spark session).
2. Type Conversion: Converting the LiteralMap provided by the Flyte engine into Python native types using _literal_map_to_python_input.
3. User Code Invocation: Running the actual execute method (which, for PythonFunctionTask, calls the decorated function).
4. Post-execution: Running post_execute for cleanup or output modification.
5. Result Serialization: Converting Python outputs back into a LiteralMap via _output_to_literal_map.

Specialized Task Types

Dynamic Tasks

A dynamic task is declared using the @dynamic decorator. It acts like a hybrid between a task and a workflow, allowing you to generate new execution nodes at runtime based on input data.

from flytekit import task, dynamic

@task
def process_item(item: int) -> int:
    return item * 2

@dynamic
def dynamic_parallel_process(items: list[int]) -> list[int]:
    return [process_item(item=i) for i in items]

Internally, PythonFunctionTask handles this by setting execution_mode to ExecutionBehavior.DYNAMIC. During execution, it calls compile_into_workflow to produce a DynamicJobSpec, which Flyte Propeller then executes as a sub-workflow.

Async and Eager Tasks

Flytekit supports asynchronous tasks using AsyncPythonFunctionTask. If you decorate an async def function with @task, flytekit automatically selects this class.

Eager tasks (or eager workflows) allow for even more dynamic behavior by allowing Python code to control task execution flow directly, similar to a standard Python script but with each call being a Flyte execution. These are implemented via EagerAsyncPythonFunctionTask and require setting is_eager=True in the metadata.

Task Resolvers

When a task is executed on a cluster, Flyte needs to know how to find and load the Python code. This is handled by TaskResolverMixin. The default_task_resolver in flytekit.core.python_auto_container identifies tasks by their module path and function name.

If you need custom loading logic (e.g., loading tasks from a database or a dynamic source), you can implement a custom resolver by overriding load_task and loader_args.

from flytekit.core.base_task import TaskResolverMixin

class MyCustomResolver(TaskResolverMixin):
    def location(self) -> str:
        return "my_package.resolvers.MyCustomResolver"

    def load_task(self, loader_args: list[str]) -> Task:
        # Custom logic to import and return the task
        ...