Task authoring and execution

Flyte tasks are the fundamental building blocks of a workflow. In flytekit, tasks are typically defined by decorating a Python function with the @task decorator. This process transforms a standard Python function into a flytekit.core.python_function_task.PythonFunctionTask object, which captures the function's interface, metadata, and execution logic for the Flyte platform.

Declaring Tasks

The most common way to author a task is using the @task decorator. flytekit uses Python type hints to automatically derive the flytekit.core.interface.Interface of the task, which defines the expected inputs and outputs.

from flytekit import task

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

When you call this function locally, it behaves like a normal Python function. However, when used within a @workflow, flytekit intercepts the call to build an execution graph.

Core Task Abstractions

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

• Task (flytekit.core.base_task.Task): The base class for all tasks. it maps closely to the Flyte IDL TaskTemplate. It does not have a Python-native interface by default.
• PythonTask (flytekit.core.base_task.PythonTask): A subclass of Task that introduces a python_interface. This is used for tasks that have Python-native types but might not have a user-defined function body (e.g., plugin-based tasks).
• PythonFunctionTask (flytekit.core.python_function_task.PythonFunctionTask): The primary class used for tasks defined with @task. It wraps a Python Callable and handles the translation between Flyte literals and Python native values during execution.

Configuring Tasks

The @task decorator accepts several arguments to control how the task is executed on the Flyte cluster. These settings are stored in flytekit.core.base_task.TaskMetadata.

Resources and Environment

You can specify resource requirements and environment variables for the task's container:

from flytekit import task, Resources

@task(
    requests=Resources(cpu="1", mem="2Gi"),
    limits=Resources(cpu="2", mem="4Gi"),
    environment={"MY_ENV_VAR": "value"}
)
def resource_intensive_task(x: int) -> int:
    return x * 2

Caching and Retries

Flyte supports memoization to avoid redundant computations. Caching requires a cache_version string; changing this version will invalidate previous cache entries.

@task(cache=True, cache_version="1.0", retries=3)
def cached_task(n: int) -> int:
    # This result will be cached based on the input 'n'
    return n + 1

Internally, TaskMetadata validates these settings. For example, if cache=True is set, cache_version must also be provided, or a ValueError is raised during initialization.

Custom Container Images

If a task requires specific dependencies not present in the default image, you can specify a custom image:

@task(container_image="ghcr.io/my-org/my-custom-image:latest")
def custom_image_task():
    # Task logic requiring specialized dependencies
    pass

Task Execution Flow

When a task is executed (either locally or on the cluster), it follows a specific lifecycle managed by the dispatch_execute method in PythonFunctionTask.

1. pre_execute: Invoked before the task body runs. This is used to set up the execution context, such as initializing a Spark session or configuring secrets.
2. Input Translation: The input_literal_map (Flyte's internal format) is converted to Python native values using _literal_map_to_python_input.
3. execute: The actual user-defined Python function is called with the native inputs.
4. post_execute: Invoked after the function returns. It can be used for cleanup or to modify the return value.
5. Output Translation: The Python return values are converted back into a LiteralMap via _output_to_literal_map to be passed to the next node in the workflow.

Local Execution and Mocking

During local execution, flytekit bypasses the Flyte backend. The local_execute method in base_task.Task handles this by translating inputs, checking a local cache (if enabled via LocalTaskCache), and calling sandbox_execute.

If a task fails during local execution, flytekit wraps the exception to provide context:

# Internal logic in base_task.py
# except TypeTransformerFailedError as exc:
#     exc.args = (f"Failed to convert inputs of task '{self.name}':\n  {exc.args[0]}",)
#     raise

Specialized Task Types

Async Tasks

Tasks defined with async def are automatically instantiated as flytekit.core.python_function_task.AsyncPythonFunctionTask. These tasks are executed within an event loop managed by flytekit's loop_manager.

Dynamic Tasks

A task can be marked as dynamic by using the @dynamic decorator (which sets ExecutionBehavior.DYNAMIC). Dynamic tasks can generate new tasks or workflows at runtime based on their inputs.

from flytekit import dynamic, task

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

@dynamic
def my_dynamic_task(n: int) -> list[str]:
    return [greet(name=str(i)) for i in range(n)]

When a dynamic task runs on the cluster, compile_into_workflow is called to produce a DynamicJobSpec, which Flyte Propeller then uses to schedule the newly generated nodes.

Eager Tasks

Eager tasks (EagerAsyncPythonFunctionTask) allow for more flexible, Pythonic execution where the Python code itself acts as the orchestrator, creating stack frames on the Flyte cluster for each task invocation. This is indicated by is_eager=True in the TaskMetadata.