Task authoring and execution

Flyte tasks are the fundamental building blocks of workflows in flytekit. They represent a single unit of execution, characterized by a strong interface (typed inputs and outputs) and declarative configuration.

Declaring Tasks

The most common way to define a task is by using the @task decorator on a Python function. flytekit automatically detects the interface from the function's type hints and docstrings.

from flytekit import task
import typing

@task
def greet(name: str) -> str:
    """
    A simple task that greets a user.
    :param name: The name of the person to greet.
    :return: A greeting string.
    """
    return f"Hello, {name}!"

When you decorate a function with @task, flytekit creates an instance of PythonFunctionTask. This class handles the translation between Flyte's type system and Python native types, as well as the execution logic.

Task Configuration

The @task decorator accepts several parameters to control execution behavior, resource allocation, and metadata:

• Retries: Use retries to specify how many times the task should be retried on failure.
• Caching: Enable caching with cache=True and provide a cache_version. flytekit uses the TaskMetadata class to manage these settings.
• Resources: Request specific compute resources using requests and limits with the Resources class.
• Timeout: Set a maximum execution duration using timeout (either an integer in seconds or a datetime.timedelta).

from datetime import timedelta
from flytekit import task, Resources

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

Task Abstractions

Internally, flytekit uses a hierarchy of classes to represent different task types:

1. Task: The base class in flytekit.core.base_task. it captures information defined in the FlyteIDL TaskTemplate.
2. PythonTask: Inherits from Task and adds support for Python native interfaces. It includes methods like get_input_types() and dispatch_execute().
3. PythonFunctionTask: The primary class for tasks defined via the @task decorator. It wraps a user-defined Python function and manages its execution.

Task Plugins

For tasks that require specialized backends (like Spark, SQL, or specialized Kubernetes Pods), flytekit uses a plugin system. You can pass a configuration object to the task_config parameter of the @task decorator.

# Example using a hypothetical Spark plugin
@task(task_config=Spark(spark_conf={"spark.executor.memory": "2g"}), retries=2)
def spark_task(x: int) -> str:
    ...

The TaskPlugins factory in flytekit.core.task manages these plugins, mapping configuration types to specific PythonFunctionTask implementations.

Data Artifacts

Tasks can produce and consume versioned, partitioned data artifacts. This allows for fine-grained data tracking and triggering. You use Annotated type hints and the Artifact class to declare these dependencies.

from typing import Annotated
from flytekit import task, Artifact
import pandas as pd

Pricing = Artifact(name="pricing", partition_keys=["region"])

@task
def get_pricing(region: str) -> Annotated[pd.DataFrame, Pricing]:
    df = pd.DataFrame({"price": [10, 20]})
    # Bind the partition value at runtime
    return Pricing.create_from(df, region=region)

The create_from method on the Artifact class interacts with the FlyteContext's output_metadata_tracker to attach partition metadata to the task output.

Execution Modes

flytekit supports several specialized execution behaviors beyond standard synchronous tasks:

Dynamic Tasks

Dynamic tasks allow you to generate a sub-workflow at runtime based on task inputs. You use the @dynamic decorator (which sets execution_mode=ExecutionBehavior.DYNAMIC). Internally, PythonFunctionTask.compile_into_workflow generates a DynamicJobSpec that Flyte Propeller executes.

Async and Eager Tasks

• AsyncPythonFunctionTask: Used when the decorated function is a coroutine (async def).
• EagerAsyncPythonFunctionTask: Implements "eager workflows" where Python code acts as the orchestrator, and each task call creates a new execution on the Flyte cluster. This is enabled by setting is_eager=True in TaskMetadata.

Local Execution

When you call a task locally (e.g., in a unit test), flytekit invokes local_execute. This method:

1. Translates Python native inputs into Flyte literals using translate_inputs_to_literals.
2. Checks the LocalTaskCache if caching is enabled.
3. Invokes sandbox_execute, which eventually calls the user's execute method.
4. Wraps the results back into Promise objects or native types.

# Local execution is as simple as calling the function
result = greet(name="Flyte User")
assert result == "Hello, Flyte User!"

Task Resolvers

When a task runs on a remote Flyte cluster, the container needs to know how to find and load the Python task object. This is handled by TaskResolverMixin. The default_task_resolver identifies tasks by their module and function name, allowing pyflyte-execute to rehydrate the task object at runtime.