Skip to main content

Task authoring and execution

Flyte tasks are the fundamental building blocks of a workflow, representing a single unit of execution. In flytekit, tasks are strongly typed, versioned, and independently executable.

Declaring Tasks​

The most common way to define a task is using the @task decorator from flytekit.core.task. This decorator transforms a standard Python function into a PythonFunctionTask.

from flytekit import task

@task
def add(a: int, b: int) -> int:
    return a + b

When you decorate a function, flytekit automatically inspects the function signature to determine the task's interface (inputs and outputs). This is handled by transform_function_to_interface in flytekit.core.interface.

Task Configuration​

The @task decorator accepts several parameters to control execution behavior:

  • Retries: Use retries to specify how many times the task should be retried on failure.
  • Timeout: Use timeout (as an int of seconds or datetime.timedelta) to limit the execution duration.
  • Caching: Enable caching by setting cache=True and providing a cache_version.
  • Resources: Specify requests and limits for CPU, memory, and other resources using the Resources class.
from datetime import timedelta
from flytekit import task, Resources

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

Internally, these configurations are stored in a TaskMetadata object (defined in flytekit.core.base_task), which is then converted to the Flyte IDL representation via to_taskmetadata_model.

Core Task Abstractions​

Flytekit uses a class hierarchy to manage different task behaviors:

  1. Task: The base class in flytekit.core.base_task. It captures the Flyte IDL TaskTemplate and does not have a Python-native interface.
  2. PythonTask: A subclass of Task that introduces a python_interface (an Interface object). It handles the mapping between Python types and Flyte literals.
  3. PythonFunctionTask: The primary class for tasks defined by Python functions. It manages the execution of the user's code and supports different execution modes like DYNAMIC and EAGER.

The Plugin System​

Flytekit supports various task types (e.g., Spark, SQL, Pods) through TaskPlugins. When you provide a task_config to the @task decorator, flytekit uses TaskPlugins.find_pythontask_plugin to locate the appropriate subclass of PythonFunctionTask.

# Example of a task using a specific config (e.g., for a Spark plugin)
@task(task_config=Spark(spark_conf={"spark.driver.memory": "2g"}))
def my_spark_task(a: int) -> int:
    ...

Task Execution​

Tasks can be executed in different environments, and flytekit handles the translation of data accordingly.

Local Execution​

When you call a task directly in a Python script, flytekit triggers local_execute. This method:

  1. Translates Python native inputs into Flyte literals using translate_inputs_to_literals.
  2. Checks the local cache if caching is enabled.
  3. Invokes sandbox_execute, which eventually calls dispatch_execute.
  4. Translates the resulting literals back into Python native values or Promise objects.

Remote Execution​

On the Flyte platform, the container starts with a command generated by a TaskResolver. The default resolver (default_task_resolver) uses the module and function name to rehydrate the task object.

The entry point calls dispatch_execute, which:

  1. Converts input literals to Python native types via _literal_map_to_python_input.
  2. Calls the user's execute method (the decorated function).
  3. Converts the return values back to literals using _output_to_literal_map.

Dynamic and Eager Tasks​

Flytekit supports advanced execution patterns:

  • Dynamic Tasks: Declared with @dynamic. These tasks return a DynamicJobSpec during execution, which Flyte Propeller then uses to expand the workflow graph at runtime. This is implemented in PythonFunctionTask.dynamic_execute.
  • Eager Tasks: Declared by setting is_eager=True in metadata (or using the @eager decorator). These tasks allow for more flexible, imperative-style execution where tasks are awaited. They are implemented by EagerAsyncPythonFunctionTask and use a Controller to manage sub-executions on the Flyte backend.

Important Constraints​

  • Nested Functions: The default_task_resolver cannot handle tasks defined inside other functions. Tasks must be accessible at the module level so they can be imported and executed in the container.
  • Caching Requirements: If cache=True is set, a cache_version must also be provided, or a ValueError will be raised during task initialization in TaskMetadata.__post_init__.
  • Async Tasks: Tasks defined with async def are instantiated as AsyncPythonFunctionTask. These tasks are executed using a loop_manager to handle the asynchronous lifecycle.