Task authoring and execution

Task Authoring and Execution in Flytekit

Tasks are the fundamental unit of computation in Flyte. Each task represents a versioned, independently executable piece of work with a strongly typed interface. Understanding how tasks are declared, configured, and executed is essential for building reliable Flyte workflows.

The @task Decorator

The primary way to define a Flyte task is with the @task decorator. Decorate any Python function with type annotations to create a task:

from flytekit import task
import datetime

@task
def process_data(data: list) -> dict:
    return {"count": len(data), "sum": sum(data)}

The decorator automatically extracts the function signature—input types from parameters, output types from the return annotation—to build the task's Interface. This happens in transform_function_to_interface() in flytekit/core/interface.py, which uses typing.get_type_hints() to read annotations and inspect.signature() to extract default values.

Configuration Options

The @task decorator accepts numerous configuration parameters:

from flytekit import task, Cache, Secret

@task(
    cache=True,
    cache_version="1.0.0",
    retries=3,
    timeout=datetime.timedelta(minutes=5),
    interruptible=True,
    container_image="my-image:tag",
    secret_requests=[Secret(group="api", key="key")],
    enable_deck=True,
)
def my_task(x: int) -> str:
    return str(x * 2)

Caching requires cache_version — setting cache=True without a version raises ValueError in TaskMetadata.__post_init__():

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

The modern approach uses the Cache object instead of separate parameters:

from flytekit import task, Cache

@task(cache=Cache(version="1.0.0", serialize=True))
def cached_task(x: int) -> int:
    return x * 2

Task Class Hierarchy

Flytekit defines a hierarchy of task classes. Understanding this hierarchy helps when you need to extend or customize task behavior.

Task (Abstract Base)

Located in flytekit/core/base_task.py, the Task class is the root of the hierarchy. It:

• Defines the abstract methods dispatch_execute(), pre_execute(), and execute()
• Stores task metadata (interface, name, task type, security context)
• Registers itself with FlyteEntities.entities on construction for serialization discovery
• Handles local execution caching through local_execute()

The __call__ method delegates to flyte_entity_call_handler() in flytekit/core/promise.py, which determines the execution mode—compilation (building a workflow graph), local execution, or remote execution.

PythonTask (Python-Native Tasks)

PythonTask extends Task with Python-specific interface handling. It introduces:

• Interface object (flytekit/core/interface.py) that represents the Python function signature
• pre_execute() and post_execute() hooks for setup and cleanup
• Literal-to-Python type conversion via _literal_map_to_python_input() and _output_to_literal_map()
• Deck generation for visualization (_write_decks())

The dispatch_execute() method in PythonTask performs the core execution lifecycle:

# Simplified from flytekit/core/base_task.py lines 727-794
new_user_params = self.pre_execute(ctx.user_space_params)
native_inputs = self._literal_map_to_python_input(input_literal_map, exec_ctx)
native_outputs = self.execute(**native_inputs)
native_outputs = self.post_execute(new_user_params, native_outputs)
literals_map, _ = run_sync(self._output_to_literal_map, native_outputs, exec_ctx)
return literals_map

PythonFunctionTask (Function-Based Tasks)

PythonFunctionTask extends PythonAutoContainerTask and is the concrete class for @task-decorated functions. It:

• Extracts the interface from the wrapped function using transform_function_to_interface()
• Validates that the task function is not nested (the default resolver cannot handle it)
• Routes execution based on ExecutionBehavior:
  • DEFAULT: Calls _task_function(**kwargs) directly
  • DYNAMIC: Invokes dynamic_execute() to compile and run a workflow at runtime
  • EAGER: Creates remote Flyte executions for nested task calls

TaskMetadata

TaskMetadata (in flytekit/core/base_task.py) is a dataclass containing execution configuration:

@dataclass
class TaskMetadata:
    cache: bool = False
    cache_serialize: bool = False
    cache_version: str = ""
    cache_ignore_input_vars: Tuple[str, ...] = ()
    interruptible: Optional[bool] = None
    deprecated: str = ""
    retries: int = 0
    timeout: Optional[Union[datetime.timedelta, int]] = None
    pod_template_name: Optional[str] = None
    generates_deck: bool = False
    is_eager: bool = False

Validation in __post_init__() enforces invariant combinations—caching requires a version, serialization requires caching, and ignored variables require caching.

Task Execution Modes

Tasks behave differently depending on the execution context, managed through ExecutionState.Mode:

Mode	Purpose
TASK_EXECUTION	Remote or Flyte backend execution
LOCAL_TASK_EXECUTION	Purely local task execution
LOCAL_WORKFLOW_EXECUTION	Local workflow execution
DYNAMIC_TASK_EXECUTION	Compiling a workflow at task runtime
EAGER_EXECUTION	Eager workflow with remote task executions
EAGER_LOCAL_EXECUTION	Eager workflow running locally
LOCAL_DYNAMIC_TASK_EXECUTION	Dynamic workflow running locally

The flyte_entity_call_handler() function in flytekit/core/promise.py implements the dispatch logic:

1. If compiling (inside a workflow), create a node and return Promise objects
2. If already in local execution, call local_execute() with the appropriate mode
3. Otherwise, start a new local execution context

Local Execution Flow

When you call a task directly (outside a workflow), local_execute() runs:

# Simplified from flytekit/core/base_task.py lines 282-362
literals = translate_inputs_to_literals(
    ctx,
    incoming_values=kwargs,
    flyte_interface_types=self.interface.inputs,
    native_types=self.get_input_types(),
)
input_literal_map = _literal_models.LiteralMap(literals=literals)

# Check local cache if enabled
if self.metadata.cache and local_config.cache_enabled:
    outputs_literal_map = LocalTaskCache.get(
        self.name, self.metadata.cache_version, input_literal_map,
        self.metadata.cache_ignore_input_vars
    )
    if outputs_literal_map is None:
        outputs_literal_map = self.sandbox_execute(ctx, input_literal_map)
        LocalTaskCache.set(self.name, self.metadata.cache_version,
                          input_literal_map, self.metadata.cache_ignore_input_vars,
                          outputs_literal_map)
else:
    outputs_literal_map = self.sandbox_execute(ctx, input_literal_map)

# Wrap outputs as Promises
return create_task_output(vals, self.python_interface)

sandbox_execute() sets up a task sandbox directory and calls dispatch_execute().

Remote Execution Flow

On the Flyte backend, dispatch_execute() receives a serialized LiteralMap of inputs:

# From flytekit/core/base_task.py lines 714-794
new_user_params = self.pre_execute(ctx.user_space_params)
native_inputs = self._literal_map_to_python_input(input_literal_map, exec_ctx)
native_outputs = self.execute(**native_inputs)  # User code runs here
native_outputs = self.post_execute(new_user_params, native_outputs)
literals_map, _ = run_sync(self._output_to_literal_map, native_outputs, exec_ctx)
return literals_map

Dynamic Tasks

Dynamic tasks (@dynamic) are tasks that compile a workflow at runtime. They're created using execution_mode=PythonFunctionTask.ExecutionBehavior.DYNAMIC:

from flytekit import dynamic, task
import typing

@task
def t1(a: int) -> int:
    return a + 1

@dynamic
def my_dynamic_subwf(n: int) -> typing.List[int]:
    results = []
    for i in range(n):
        results.append(t1(a=i))
    return results

In local execution, dynamic_execute() compiles and executes the workflow directly. On the backend, it produces a DynamicJobSpec containing the compiled workflow template.

Eager Workflows

Eager workflows (@eager) use async Python to execute nested tasks on the Flyte backend:

from flytekit import task, eager

@task
def add_one(x: int) -> int:
    return x + 1

@eager
async def eager_workflow(x: int) -> int:
    out = add_one(x=x)
    return add_one(x=out)

EagerAsyncPythonFunctionTask (in flytekit/core/python_function_task.py) manages this behavior:

• Sets TaskMetadata(is_eager=True) automatically
• During local execution, runs the async function with a Controller that submits nested tasks to the Flyte backend via FlyteRemote
• Uses a WorkerQueue to manage concurrent task submissions
• The EagerFailureHandlerTask cleans up running child executions if the eager workflow fails

Task Resolvers

Task resolvers (TaskResolverMixin) enable task rehydration from serialized container arguments. When a container starts, it receives arguments like:

pyflyte-execute --inputs s3://path/inputs.pb --output-prefix s3://outputs/ \
  --resolver flytekit.core.python_auto_container.default_task_resolver -- \
  task-module my_module task-name my_task

The DefaultTaskResolver (in flytekit/core/python_auto_container.py) uses:

• load_task(loader_args): Calls importlib.import_module() and looks up the task by name
• loader_args(settings, task): Returns ["task-module", m, "task-name", t]

Custom resolvers can store tasks in databases, class variables (ClassStorageTaskResolver), or pickle files (DefaultNotebookTaskResolver).

Testing Tasks

Use task_mock() from flytekit/core/testing.py to mock task execution in tests:

from flytekit import task
from flytekit.core.testing import task_mock

@task
def t1(i: int) -> int:
    return i * 2

def test_task():
    with task_mock(t1) as m:
        m.return_value = 42
        result = t1(10)
        assert result == 42

task_mock() temporarily replaces the execute method with a mock, restoring the original afterward.

Extending Tasks

Custom Task Plugins

Register custom task types via TaskPlugins:

from flytekit.core.task import TaskPlugins

TaskPlugins.register_pythontask_plugin(MyConfig, MyCustomTask)

When @task(task_config=MyConfig(...)) is used, the factory finds the matching plugin.

Custom Executors (Shim Tasks)

For external container tasks, use ExecutableTemplateShimTask and ShimTaskExecutor:

from flytekit.core.shim_task import ShimTaskExecutor, ExecutableTemplateShimTask

class MyExecutor(ShimTaskExecutor):
    def execute_from_model(self, task_template, **kwargs):
        # Business logic using task_template.custom
        pass

class MyContainerTask(ExecutableTemplateShimTask, PythonTask):
    def __init__(self, **kwargs):
        super().__init__(tt=None, executor_type=MyExecutor, **kwargs)
    
    def get_custom(self, settings):
        return {"my_config": self.task_config.value}

Decorator Gotchas

Nested Functions Not Allowed

The default task resolver cannot handle nested functions. This raises ValueError:

@task
def outer():
    @task  # Error: nested function
    def inner():
        pass
    return inner()

Test modules (test_*.py) are exempt. Custom decorators must use functools.wraps or functools.update_wrapper.

Output Tuples Not Supported

Tasks cannot return plain tuples—use typing.NamedTuple for multiple outputs:

from typing import NamedTuple

@task
def multi() -> NamedTuple("Result", [("x", int), ("y", str)]):
    return (1, "hello")  # Works with NamedTuple

@task
def bad() -> tuple:  # TypeError at runtime
    return (1, "hello")

Async Task Detection

The @task decorator automatically detects async functions and wraps them with AsyncPythonFunctionTask:

@task
async def async_task(x: int) -> int:
    return x * 2

Reference Tasks

Reference tasks (@reference_task) point to remotely registered tasks without a local implementation:

import typing
from flytekit import reference_task

@reference_task(
    project="flytesnacks",
    domain="development",
    name="recipes.aaa.simple.join_strings",
    version="553018f39e519bdb2597b652639c30ce16b99c79",
)
def ref_t1(a: typing.List[str]) -> str:
    """Empty function acts as skeleton to match remote interface."""
    return "hello"

Reference tasks cannot be run locally—they must be mocked with task_mock().