Task authoring and execution

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

Declaring Tasks with the @task Decorator

The most common way to define a task in flytekit 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 greet(name: str) -> str:
    return f"Hello, {name}!"

When you apply @task, flytekit uses transform_function_to_interface to automatically infer the task's input and output types from the function's type hints. This ensures that the task adheres to Flyte's type system.

Task Configuration and Metadata

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

• Retries: Use retries to specify how many times the task should be re-run on failure.
• Timeouts: Use timeout (as an int seconds or datetime.timedelta) to limit the execution duration.
• Caching: Enable caching with cache=True and provide a cache_version. Flyte will skip execution if the same inputs and version are encountered again.
• Resources: Specify requests and limits using the Resources class to define CPU, memory, and GPU requirements.

from datetime import timedelta
from flytekit import task, Resources

@task(
    retries=3,
    timeout=timedelta(minutes=5),
    cache=True,
    cache_version="1.0",
    requests=Resources(cpu="2", mem="500Mi"),
    limits=Resources(cpu="4", mem="1Gi")
)
def resource_intensive_task(data: list[int]) -> int:
    return sum(data)

Core Task Abstractions

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

Task and PythonTask

The Task class in flytekit.core.base_task is the base for all tasks. It maps to the Flyte IDL TaskTemplate. PythonTask extends this to support Python-native interfaces, handling the translation between Python types and Flyte literals via the TypeEngine.

PythonFunctionTask

Located in flytekit.core.python_function_task, this class is the engine behind the @task decorator. It captures the user's function and manages its execution.

• execute: This method is called during local execution or by the Flyte executor. It simply invokes the underlying _task_function.
• dispatch_execute: This is the entry point used by the Flyte platform at runtime. It handles input translation from Flyte literals to Python types, calls execute, and then translates the results back to literals.

PythonInstanceTask

For tasks that do not have a user-defined function body but instead have a platform-defined execute method, you should inherit from PythonInstanceTask. This is useful for building reusable task plugins where the logic is encapsulated within the class itself.

Task Execution Flow

When a task is executed, flytekit manages the transition between the Flyte type system and Python native code.

1. Input Translation: In dispatch_execute, the input_literal_map (Flyte's internal representation) is converted to Python native kwargs using _literal_map_to_python_input.
2. Pre-execution: The pre_execute method is called, allowing for context setup (e.g., initializing a Spark session).
3. User Code Execution: The execute method runs the actual logic.
4. Post-execution: The post_execute method allows for cleanup or output modification.
5. Output Translation: The Python return values are converted back into a LiteralMap via _output_to_literal_map.

Local Execution

You can run tasks locally just like regular Python functions. Flytekit's local_execute method handles this by mocking the Flyte environment. If caching is enabled locally, LocalTaskCache (found in flytekit.core.base_task) will attempt to retrieve results from a local cache before executing the task.

Task Resolvers

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

The default_task_resolver in flytekit.core.python_auto_container is the standard implementation. It works by:

1. loader_args: Identifying the module and function name of the task during serialization.
2. load_task: Importing the module and retrieving the task object at runtime using importlib.

If you need custom loading logic (e.g., loading tasks from a database or a dynamic source), you can implement a custom resolver by inheriting from TaskResolverMixin and passing it to the @task decorator.

Special Task Types

Async Tasks

If you define a task as an async def, flytekit automatically uses AsyncPythonFunctionTask. These tasks are executed within an event loop managed by flytekit's loop_manager.

Eager Tasks

Eager tasks (declared via EagerAsyncPythonFunctionTask) allow for more dynamic execution patterns where Python code acts as the orchestrator, spawning other Flyte executions and awaiting their results. This is indicated by setting is_eager=True in the TaskMetadata.

Reference Tasks

A ReferenceTask is a pointer to a task that already exists on a Flyte cluster. It allows you to call tasks defined in other projects or domains without having the source code available locally.

from flytekit import reference_task

@reference_task(
    project="flytesnacks",
    domain="development",
    name="core.control_flow.merge_sort.merge",
    version="v1"
)
def merge(sorted_list1: list[int], sorted_list2: list[int]) -> list[int]:
    ...