Task authoring and execution

Flyte tasks are the fundamental building blocks of flytekit. They represent a single unit of execution, characterized by a strong interface (typed inputs and outputs), versioning, and declarative configuration. In flytekit, tasks are typically defined using the @task decorator, which transforms a standard Python function into a PythonFunctionTask.

Declaring Tasks

The most common way to author a task is by decorating a Python function with @task. flytekit uses Python type hints to automatically derive the task's interface, which is then used by the Flyte engine to ensure type safety across the workflow.

from flytekit import task
import typing

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

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

Internally, the @task decorator in flytekit/core/task.py creates an instance of PythonFunctionTask. This class captures the function's metadata and provides the logic for both local and remote execution.

Task Metadata and Configuration

You can configure task behavior—such as retries, timeouts, and caching—by passing arguments to the @task decorator. These settings are stored in the TaskMetadata class (found in flytekit/core/base_task.py).

from datetime import timedelta
from flytekit import task

@task(
    retries=3,
    timeout=timedelta(minutes=60),
    cache=True,
    cache_version="1.0",
    interruptible=True
)
def heavy_computation(data: list[float]) -> float:
    ...

The TaskMetadata class validates these parameters. For instance, if cache=True is set, cache_version must also be provided, or flytekit will raise a ValueError.

Task Execution Flow

flytekit handles task execution differently depending on whether it is running locally or on a remote Flyte cluster.

Local Execution

When you call a task function directly in a Python script, flytekit invokes Task.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. Calls sandbox_execute, which eventually triggers the user's execute method.
4. Wraps the results back into Promise objects or native Python types.

Remote Execution

On a Flyte cluster, the execution is managed by dispatch_execute in PythonTask. The process involves:

1. Pre-execution: pre_execute is called to set up the environment (e.g., initializing a Spark session).
2. Input Translation: _literal_map_to_python_input converts the LiteralMap received from the Flyte engine into Python native types.
3. User Code Execution: The execute method runs the actual Python function.
4. Output Translation: _output_to_literal_map converts the function's return values back into a LiteralMap for the Flyte engine.

Advanced Task Types

While @task covers most use cases, flytekit provides base classes for specialized execution patterns.

PythonFunctionTask Behaviors

The PythonFunctionTask supports different ExecutionBehavior modes:

• DEFAULT: Standard task execution.
• DYNAMIC: Used for tasks that generate a new workflow at runtime based on inputs. This is triggered by the @dynamic decorator.
• EAGER: Used for "eager workflows" where Python code acts as the orchestrator, allowing for more flexible control flow (e.g., if/else based on task outputs) while still executing tasks on the cluster.

Custom Task Plugins

If you need to integrate with external systems (like Spark, SQL, or specialized hardware), you can use task plugins. These are registered via TaskPlugins.register_pythontask_plugin.

# Example of using a Spark task (requires flytekitplugins-spark)
from flytekitplugins.spark import Spark
from flytekit import task

@task(
    task_config=Spark(
        spark_conf={"spark.driver.memory": "2g"},
        hadoop_conf={"fs.s3a.access.key": "key"}
    )
)
def my_spark_task(df: typing.Any) -> typing.Any:
    ...

When a task_config is provided, flytekit looks up the corresponding plugin class. If found, it instantiates that specific plugin class instead of the base PythonFunctionTask.

Container Image Configuration

By default, tasks run in the default Flyte container image. You can override this for specific tasks using the container_image parameter. This supports Go-style templating to reference images defined in your Flyte configuration.

@task(container_image="{{.images.my_custom_image.fqn}}:{{.images.default.tag}}")
def specialized_task():
    ...

This ensures that tasks with heavy or conflicting dependencies can run in isolated environments without bloating the primary project image.