Core Task Architecture

Flytekit tasks are the fundamental building blocks of execution. While you typically interact with them via the @task decorator, they are implemented through a structured hierarchy of classes that handle everything from type translation to container command generation.

The Task Hierarchy

At the root of the architecture is the Task class in flytekit.core.base_task. This class is a thin wrapper around the Flyte IDL TaskTemplate. It captures metadata and defines the interface but does not have a native Python execution body.

Most tasks you write inherit from PythonTask, which adds support for a Python native Interface. This allows flytekit to translate between Python types (like int, str, or pandas.DataFrame) and Flyte's internal literal system.

The hierarchy typically follows this path:

1. Task: The base IDL-compatible class.
2. PythonTask: Adds Python type awareness.
3. PythonAutoContainerTask: (Internal) Automatically handles container image resolution and generates the pyflyte-execute command.
4. PythonFunctionTask: The implementation used for @task decorated functions.

Python Function Tasks

When you use the @task decorator, flytekit instantiates a PythonFunctionTask. This class is designed to wrap a standard Python function and make it executable on the Flyte platform.

from flytekit import task

@task(retries=3, cache=True, cache_version="1.0")
def square(n: int) -> int:
    return n * n

Internal Execution Flow

When a task runs, it goes through two distinct paths depending on the environment:

1. Local Execution: When you call square(n=5) in a script or test, flytekit invokes Task.local_execute. This method handles local caching (via LocalTaskCache) and calls sandbox_execute, which eventually triggers the user's function.
2. Remote Execution: On the Flyte cluster, the container starts with a command like pyflyte-execute. This entry point calls dispatch_execute.

Inside dispatch_execute (found in flytekit.core.base_task.PythonTask), flytekit performs the following steps:

• Input Translation: Converts Flyte LiteralMap inputs into Python native values using _literal_map_to_python_input.
• Pre-execution: Calls pre_execute to set up the environment (e.g., initializing a Spark session).
• User Code Execution: Calls the actual execute method (which, in PythonFunctionTask, simply runs the decorated function).
• Output Translation: Converts the Python return values back into a Flyte LiteralMap using _output_to_literal_map.

Raw Container Tasks

If you need to run code that isn't Python, or you have a pre-built binary in a container, you use ContainerTask. Unlike PythonFunctionTask, which generates its own command, ContainerTask requires you to specify the image, command, and arguments manually.

You can use template strings to map Flyte inputs to command-line arguments:

from flytekit import ContainerTask, kwtypes

calculate_area = ContainerTask(
    name="calculate_ellipse_area",
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs=kwtypes(a=float, b=float),
    outputs=kwtypes(area=float),
    image="ghcr.io/flyteorg/rawcontainers-python:v2",
    command=[
        "python",
        "calculate-area.py",
        "{{.inputs.a}}",
        "{{.inputs.b}}",
        "/var/outputs",
    ],
)

In this example, {{.inputs.a}} is a placeholder that Flyte Propeller replaces with the actual input value at runtime. The ContainerTask.execute method (in flytekit.core.container_task) even supports local execution by using docker-py to run the container on your local machine.

Task Metadata and Configuration

The TaskMetadata class in flytekit.core.base_task stores the execution policies for a task. When you pass arguments to the @task decorator, they are encapsulated into this object.

Key attributes include:

• retries: Managed by retry_strategy, determining how many times the Flyte platform will attempt to re-run the task on failure.
• timeout: A datetime.timedelta that limits the maximum duration of a single execution.
• cache and cache_version: Controls whether Flyte should skip execution if the same inputs have been processed before.
• interruptible: Indicates if the task can run on lower-cost, preemptible nodes (like AWS Spot Instances).

Task Resolvers

A critical problem in distributed execution is "rehydration": how does a container know which Python function to run? Flytekit solves this using the TaskResolverMixin.

The DefaultTaskResolver (used by most tasks) works by:

1. Serialization: At registration time, it records the module name and the function name (e.g., my_project.tasks.square).
2. Loading: At runtime, the pyflyte-execute command receives these as loader_args. The resolver uses importlib to import the module and retrieve the task object by name.

If you have a non-standard way of loading tasks (e.g., loading from a database or a dynamic source), you can implement a custom resolver by inheriting from TaskResolverMixin and overriding load_task and loader_args.

Handling Task Outputs

Flytekit provides a specialized exception, IgnoreOutputs, for scenarios where a task completes but its outputs should not be recorded or passed downstream. This is particularly useful in distributed training where only the rank-0 process should emit the final model or metrics. When IgnoreOutputs is raised, the Flyte platform treats the task as successful but ignores any generated data.