Skip to main content

Task authoring and execution

Flyte tasks are the fundamental building blocks of a workflow. In flytekit, tasks are declared using decorators that transform Python functions into specialized task objects. These objects manage the lifecycle of a task, from local execution and testing to serialization and remote execution on a Flyte cluster.

Task Declaration and Configuration​

The primary way to author a task is using the @task decorator. This decorator wraps a Python function and instantiates a PythonFunctionTask (or AsyncPythonFunctionTask for coroutines).

from flytekit import task
import datetime

@task(
    cache=True,
    cache_version="1.0",
    retries=3,
    timeout=datetime.timedelta(minutes=60),
    environment={"MY_ENV_VAR": "value"},
)
def my_task(x: int) -> str:
    return str(x)

Task Metadata​

Configuration for a task is captured in the TaskMetadata class. This class validates and stores execution parameters:

  • Caching: Enabled via cache=True. You must provide a cache_version. You can also use cache_serialize=True to ensure identical inputs are processed serially and cache_ignore_input_vars to exclude specific inputs from the cache key.
  • Retries: Specified as an integer. Flyte will retry the task at least n times on failure.
  • Timeout: Can be a datetime.timedelta or an integer representing seconds.
  • Interruptible: Indicates if the task can be scheduled on lower-priority, pre-emptible nodes.

Internally, TaskMetadata.__post_init__ enforces strict rules:

  • If cache is True, cache_version must be set.
  • cache_serialize and cache_ignore_input_vars require cache=True.

Decorator Variants​

Flytekit provides specialized decorators for different execution behaviors:

  • @dynamic: Defined in dynamic_workflow_task.py as a partial of @task with ExecutionBehavior.DYNAMIC. It allows the task body to act like a workflow, generating new tasks and subworkflows at runtime based on input data.
  • @eager: Constructs an EagerAsyncPythonFunctionTask. Eager workflows allow you to use standard Python control flow (like if and for) with Flyte entities by executing them immediately against a backend.

Core Task Abstractions​

The task system is built on a hierarchy of classes that separate Flyte IDL concerns from Python-native execution.

The Task Hierarchy​

  1. Task: The root base class in base_task.py. it captures the FlyteIDL specification, including the task_type, name, and interface (as a TypedInterface).
  2. PythonTask: Adds a Python-native Interface to the base Task. It handles the translation between Flyte literals and Python types using the TypeEngine.
  3. PythonFunctionTask: The most common task type. It wraps a user-defined Python function and auto-detects its interface.

Task Resolution​

When a task runs on a remote cluster, the container needs to know how to rehydrate the Python task object. This is handled by the TaskResolverMixin. The default resolver identifies the task by its module and function name.

At serialization time, flytekit generates command-line arguments for pyflyte-execute:

pyflyte-execute --resolver flytekit.core.python_auto_container.default_task_resolver \
-- task-module my_module task-name my_task_function

Execution Lifecycle​

The dispatch_execute method in PythonTask (and its overrides) defines the lifecycle of a task execution, whether running locally or on a cluster.

1. Pre-execution​

The pre_execute method is called before any inputs are converted. This is used to set up environment-specific contexts, such as a SparkSession.

2. Input Translation​

Flytekit translates the LiteralMap (Flyte's internal data format) into Python-native keyword arguments using _literal_map_to_python_input. This relies on the TypeEngine to map IDL types to Python types.

3. Execution​

The core logic is invoked in the execute method.

  • For PythonFunctionTask, this simply calls the wrapped _task_function(**kwargs).
  • For AsyncPythonFunctionTask, it awaits the coroutine.
  • For EagerAsyncPythonFunctionTask, it manages a Controller and worker queue to coordinate remote executions.

4. Post-execution and Output Translation​

After execute completes, post_execute allows for final cleanup or output modification. Finally, _output_to_literal_map converts the Python return values back into a LiteralMap for Flyte.

5. Decks​

If enabled via enable_deck=True, the task generates Flyte Decks. PythonFunctionTask automatically includes SOURCE_CODE and DEPENDENCIES decks by default.

Local Execution and Caching​

When you call a task directly in a Python script, it triggers local_execute.

# This triggers local_execute
result = my_task(x=10)

local_execute performs the following:

  1. Translates native inputs to literals.
  2. Checks the LocalTaskCache if metadata.cache is enabled and LocalConfig.auto().cache_enabled is true.
  3. If a cache miss occurs, it calls sandbox_execute, which eventually runs dispatch_execute.
  4. Wraps the resulting literals back into Promise objects or a VoidPromise.

Implementation Gotchas​

  • Nested Functions: The default task resolver cannot handle nested or local functions because they cannot be easily imported by the container. Tasks must be accessible at the module level.
  • Async and Dynamic: AsyncPythonFunctionTask does not support DYNAMIC execution mode. Eager and dynamic behaviors are currently mutually exclusive.
  • Manual Output Names: If a task returns multiple values, Flyte typically names them o0, o1, etc. To provide meaningful names, use typing.NamedTuple as the return type annotation.
  • IgnoreOutputs: You can raise the IgnoreOutputs exception within a task to indicate that the generated outputs should be discarded, which is useful in distributed training scenarios.