Task authoring and execution

Flyte tasks are the fundamental building blocks of a workflow. They represent a single unit of execution, characterized by a versioned, strongly-typed interface. In flytekit, tasks are typically declared using the @task decorator, which transforms a standard Python function into a PythonFunctionTask.

Declaring Tasks

The most common way to define a task is by decorating a Python function with @task. This allows flytekit to automatically detect the task's interface (inputs and outputs) based on Python type hints.

from flytekit import task

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

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

Task Configuration

The @task decorator accepts several parameters to control execution behavior, resource allocation, and metadata. These configurations are encapsulated in the TaskMetadata class internally.

• Retries: Use retries to handle transient failures.
• Caching: Enable cache and provide a cache_version to avoid redundant computations.
• Resources: Specify requests and limits for CPU, memory, and GPU using the Resources class.
• Timeout: Set a timeout (as an int seconds or datetime.timedelta) to terminate long-running executions.

import datetime
from flytekit import task, Resources

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

Core Task Abstractions

Flytekit uses a hierarchy of classes to manage different task types and execution logic.

The Task Base Class

The base_task.Task class is the lowest-level abstraction. it captures information defined in the Flyte IDL (Interface Definition Language) and does not have a native Python interface. It is generally not used directly by developers.

PythonTask

The base_task.PythonTask class extends the base Task to include a Python-native Interface. It handles the translation between Flyte's internal type system and Python types.

PythonFunctionTask

The python_function_task.PythonFunctionTask is the primary class used for tasks defined with the @task decorator. It wraps the user-defined task_function and manages its execution via the execute method.

Execution Modes

Flytekit supports different execution behaviors through the ExecutionBehavior enum in PythonFunctionTask.

Default Execution

In the DEFAULT mode, the task runs as a single unit of work. The dispatch_execute method handles the conversion of input LiteralMaps to Python native values, calls the user's execute function, and converts the results back to LiteralMaps.

Dynamic Tasks

Dynamic tasks (declared with @dynamic) allow you to generate a workflow graph at runtime based on inputs. Internally, this sets the execution_mode to DYNAMIC. When executed, the task function returns a DynamicJobSpec instead of standard outputs, which Flyte then uses to schedule further tasks.

Eager Tasks

Eager tasks (declared with @eager) allow for more flexible, Pythonic execution where you can use standard Python control flow (like if statements or loops) that depend on task outputs.

The EagerAsyncPythonFunctionTask class implements this by treating Python as the orchestrator. Every task invocation within an eager task creates a new execution on the Flyte cluster (or locally) rather than just adding a node to a static graph.

from flytekit import eager, task

@task
def get_count() -> int:
    return 5

@eager
async def eager_workflow():
    count = await get_count()
    if count > 0:
        # Standard Python logic allowed here
        print(f"Count is {count}")

Task Plugins

For specialized execution environments (like Spark, Ray, or SQLAlchemy), flytekit uses a plugin system. You can pass a task_config to the @task decorator to trigger a specific plugin.

from flytekit import task
from flytekitplugins.spark import Spark

@task(
    task_config=Spark(
        spark_conf={"spark.driver.memory": "1000M"},
        hadoop_conf={"fs.s3a.access.key": "key"},
    )
)
def spark_task(x: int) -> int:
    # This function will be executed in a Spark context
    return x + 1

Internally, TaskPlugins.find_pythontask_plugin looks up the appropriate PythonFunctionTask subclass registered for the given task_config type.

Local Execution and Testing

Flytekit is designed to be testable. When you call a task directly in a script or test, it triggers local_execute. This method bypasses the Flyte backend and runs the code locally, while still performing type validation and local caching if configured.

If a task fails during local execution, flytekit raises the original Python exception to simplify debugging. During remote execution, these are wrapped in FlyteUserRuntimeException or FlyteNonRecoverableSystemException to provide better context for the Flyte platform.