Task authoring and execution

Flyte tasks are the fundamental building blocks of Flyte workflows. They represent independent, versioned, and declarative units of execution with strong interfaces. In flytekit, tasks are primarily authored using the @task decorator, which transforms a standard Python function into a PythonFunctionTask.

Declaring Tasks

You define a task by decorating a Python function with @task. The function's type annotations are used to derive the Flyte interface, ensuring type safety across the platform.

from flytekit import task
import typing

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

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

Task Configuration and Metadata

The @task decorator accepts various parameters to configure the task's behavior on the Flyte platform. These settings are encapsulated internally in the TaskMetadata class (found in base_task.py).

• Caching: Enable caching to avoid re-running tasks with identical inputs.
• Retries: Specify the number of times Flyte should retry the task on failure.
• Timeout: Set a maximum duration for the task execution.
• Resources: Request specific CPU, memory, or GPU resources.

from flytekit import task, Resources

@task(
    cache=True, 
    cache_version="1.0", 
    retries=3, 
    timeout=3600,
    requests=Resources(cpu="2", mem="500Mi"),
    limits=Resources(cpu="4", mem="1Gi")
)
def resource_intensive_task(data: typing.List[float]) -> float:
    return sum(data)

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

Core Task Abstractions

Flytekit uses a layered class hierarchy to manage task behavior:

1. Task (base_task.py): The base class for all tasks. it captures the Flyte IDL specification (name, type, interface, metadata) and handles the core lifecycle hooks like local_execute and dispatch_execute.
2. PythonTask (base_task.py): Extends Task to support Python-native interfaces. It handles the translation between Flyte literals and Python native types using the TypeEngine.
3. PythonAutoContainerTask (python_auto_container.py): A base class for tasks that run within a container. It automatically generates the pyflyte-execute command used by the Flyte platform to run the task.
4. PythonFunctionTask (python_function_task.py): The most common task type, used for tasks defined by a Python function. It auto-detects the interface and handles function execution.

Task Execution Flow

When you call a task, its behavior depends on the context:

• Local Execution: When called in a local script, Task.local_execute is invoked. It translates native Python inputs into Flyte literals, checks the LocalTaskCache if caching is enabled, and then calls execute.
• Workflow Compilation: When called inside a @workflow, the task's compile method is called. Instead of executing the function, it creates a Node and links it into the workflow graph using Promise objects to represent future outputs.
• Remote Dispatch: On the Flyte platform, the dispatch_execute method is called. It handles the full lifecycle: pre_execute (setup), converting input literals to Python types, running the user's execute method, and finally post_execute (cleanup) and converting outputs back to literals.

Task Serialization and Resolvers

For a task to run on a remote Flyte cluster, the platform needs to know how to "rehydrate" the Python task object inside the container. This is handled by TaskResolverMixin.

The DefaultTaskResolver (in python_auto_container.py) works by capturing the module and name of the task. At execution time, it uses importlib.import_module to load the task.

# Internal command generated for a task
pyflyte-execute --inputs s3://... --resolver flytekit.core.python_auto_container.default_task_resolver -- task-module my_module task-name my_task

If you define tasks in a notebook, flytekit uses the DefaultNotebookTaskResolver, which loads tasks from a pickled file (cloudpickle) to preserve the state of the notebook environment.

Dynamic Tasks

Dynamic tasks allow you to generate a workflow at runtime based on the task's inputs. You declare them using the @dynamic decorator (which is a specialized PythonFunctionTask with ExecutionBehavior.DYNAMIC).

from flytekit import dynamic, task

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

@dynamic
def my_dynamic_task(n: int) -> typing.List[int]:
    return [process_item(item=i) for i in range(n)]

Inside a dynamic task, you can use native Python control flow (like loops and conditionals) to call other tasks. Flytekit compiles the result of the dynamic task into a DynamicJobSpec, which the Flyte backend then executes as a sub-workflow.

Eager Workflows

Eager workflows (@eager) allow for even more flexibility by executing Flyte entities immediately using an asynchronous Controller. Unlike dynamic tasks, they are not compiled into a static spec.

from flytekit import task, eager
import asyncio

@task
def add(x: int, y: int) -> int:
    return x + y

@eager
async def eager_workflow(x: int) -> int:
    # Tasks are awaited and executed immediately on the backend
    first = await add(x=x, y=5)
    if first > 10:
        return await add(x=first, y=10)
    return first

# Local execution
if __name__ == "__main__":
    print(asyncio.run(eager_workflow(x=2)))

Eager workflows require a FlyteRemote connection to interact with the backend. If an eager workflow fails, the EagerFailureHandlerTask is triggered to terminate any orphaned executions that were kicked off during the run.

Gotchas and Limitations

• Nested Functions: The DefaultTaskResolver cannot handle nested or local functions because they cannot be imported by name. Tasks must be defined at the module level.
• Caching Invariants: Caching requires a cache_version. If you enable cache=True without a version, TaskMetadata will raise a ValueError.
• Map Tasks: map_task only supports standard PythonFunctionTask or PythonInstanceTask. It does not support dynamic or eager tasks as the sub-task.
• Decks: The disable_deck parameter is deprecated in favor of enable_deck. If both are provided to a PythonTask, it raises a ValueError.