The Core Task Abstraction

Tasks in flytekit are the fundamental building blocks of workflows, representing atomic units of execution. While users typically interact with tasks via the @task decorator, flytekit implements a multi-layered abstraction that separates the task's interface, its infrastructure requirements, and its execution logic.

The Task Hierarchy

Flytekit uses a class hierarchy to build up the functionality required to run Python code on a distributed platform:

1. base_task.Task: The root abstract class. it is a direct representation of the Flyte IDL TaskTemplate. It captures metadata and the TypedInterface but does not have a Python-native interface.
2. base_task.PythonTask: Adds a python_interface to the base task. It uses the TypeEngine to translate between Python native types and Flyte's Literal types.
3. python_auto_container.PythonAutoContainerTask: Handles the infrastructure details. It manages container images, resource requests/limits, environment variables, and the command-line arguments needed to run the task in a remote container.
4. python_function_task.PythonFunctionTask: The most common implementation. It wraps a standard Python function, automatically detecting its interface and handling the execution of that function.

Creating Tasks with @task

The @task decorator is the primary entry point for defining tasks. When you decorate a function, flytekit creates an instance of PythonFunctionTask.

from flytekit import task

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

Task Configuration and Metadata

Configuration passed to the decorator is stored in the base_task.TaskMetadata class. This class ensures that parameters like caching and timeouts are valid before the task is even registered.

• Caching: If cache=True is set, cache_version must also be provided. You can also use cache_serialize=True to ensure that identical instances of a task (same inputs) are executed serially to avoid redundant work.
• Timeouts: The timeout parameter can be an int (seconds) or a datetime.timedelta.
• Retries: The retries parameter is converted into a _literal_models.RetryStrategy during serialization.

# Internal validation in TaskMetadata
if self.cache and not self.cache_version:
    raise ValueError("Caching is enabled ``cache=True`` but ``cache_version`` is not set.")

Infrastructure and Containerization

The PythonAutoContainerTask class is responsible for defining how the task runs in a container. It provides methods to generate the container definition used by the Flyte backend.

Resource Management

You can specify resource requirements directly in the task decorator. These are captured in a ResourceSpec within the task instance.

from flytekit import Resources

@task(
    requests=Resources(cpu="1", mem="2Gi"),
    limits=Resources(cpu="2", mem="4Gi"),
)
def resource_intensive_task(data: list) -> str:
    ...

Custom Pod Templates

For advanced Kubernetes configurations, you can provide a PodTemplate. When a PodTemplate is present, PythonAutoContainerTask.get_container returns None, and the logic shifts to get_k8s_pod, which merges the container definition into the provided pod specification.

from flytekit.core.pod_template import PodTemplate
from kubernetes.client.models import V1Container, V1PodSpec

@task(
    pod_template=PodTemplate(
        primary_container_name="primary",
        pod_spec=V1PodSpec(
            containers=[V1Container(name="primary")]
        ),
    )
)
def task_with_custom_pod(i: str):
    ...

Task Resolution and Execution

When a task runs on a remote Flyte cluster, the container starts by running pyflyte-execute. The container needs to know which specific Python task to load and execute. This is handled by the TaskResolverMixin.

The Default Task Resolver

The python_auto_container.DefaultTaskResolver is the standard mechanism for locating tasks. It works by:

1. Serialization: Capturing the module name and the task name (the function name or variable name).
2. Loading: Using importlib.import_module to re-import the module in the container and retrieve the task object.

The command generated by PythonAutoContainerTask.get_default_command looks like this:

pyflyte-execute --inputs {{.input}} --output-prefix {{.outputPrefix}} \
  --resolver flytekit.core.python_auto_container.default_task_resolver \
  -- task-module my_module task-name my_task_function

Execution Flow

The execution logic is split between local and remote paths:

• local_execute: Used when you call a task directly in a Python script or test. It handles input translation, checks the LocalTaskCache, and calls dispatch_execute.
• dispatch_execute: The entry point for both local and remote execution. It prepares the ExecutionParameters (user space params), translates input LiteralMap to Python native values, and finally calls the user's execute method.

# Simplified dispatch_execute flow in PythonFunctionTask
def dispatch_execute(self, ctx, input_literal_map):
    # 1. Pre-execute (setup context)
    new_user_params = self.pre_execute(ctx.user_space_params)
    
    # 2. Translate inputs
    native_inputs = self._literal_map_to_python_input(input_literal_map, exec_ctx)
    
    # 3. Execute user code
    native_outputs = self.execute(**native_inputs)
    
    # 4. Translate outputs back to Literals
    return run_sync(self._output_to_literal_map, native_outputs, exec_ctx)

Specialized Execution Modes

Dynamic Tasks

A dynamic task is a task that generates a workflow at runtime. This is achieved by setting the execution_mode to ExecutionBehavior.DYNAMIC. When dispatch_execute runs a dynamic task, it doesn't just return values; it returns a DynamicJobSpec containing the generated nodes and tasks.

from flytekit import dynamic

@dynamic
def my_dynamic_task(n: int):
    for i in range(n):
        square(n=i)

Eager Tasks

Eager tasks (EagerAsyncPythonFunctionTask) allow for more imperative, asynchronous execution where Python code acts as the orchestrator. Unlike standard tasks, eager tasks can await the results of other tasks within a live execution environment. They use a Controller to manage the worker queue and communicate with the Flyte backend during runtime.