Skip to main content

Tasks: Base Classes and Execution Model

The @task decorator is how nearly every Flyte task comes into existence, but beneath it is a layered class hierarchy that handles everything from type conversion to container serialization to remote execution dispatch. Understanding those layers tells you both how to configure tasks precisely and how to extend them for custom task types.

The Class Hierarchy​

Flytekit's task abstraction is a four-level stack:

Task                        # flytekit/core/base_task.py — IDL-level, no Python types
  └─ PythonTask             # flytekit/core/base_task.py — adds Python-native Interface
       └─ PythonAutoContainerTask  # flytekit/core/python_auto_container.py — adds container/image resolution
            ├─ PythonFunctionTask  # flytekit/core/python_function_task.py — wraps a Python callable
            └─ PythonInstanceTask  # flytekit/core/python_function_task.py — for platform-defined execute()

Each layer adds something concrete:

  • Task holds a TypedInterface (Flyte's IDL-level interface), the TaskMetadata, task_type string, SecurityContext, and Documentation. It also registers itself in the global FlyteEntities.entities list at construction time — this is how the serialization tooling discovers tasks at registration time. Task declares the abstract contract: dispatch_execute(), pre_execute(), and execute().

  • PythonTask replaces the IDL TypedInterface with a Python-native Interface (typed with actual Python types). It's generic over T (the task config type) and adds environment variables, deck support, and the concrete dispatch_execute() pipeline.

  • PythonAutoContainerTask (in flytekit/core/python_auto_container.py) adds container image resolution, resource requests/limits, pod template handling, and the TaskResolverMixin that determines how the task reconstructs itself at runtime.

  • PythonFunctionTask and PythonInstanceTask are the leaves. PythonFunctionTask wraps a Python callable and auto-detects its interface via transform_function_to_interface(). PythonInstanceTask is for classes where you override execute() yourself (used by the shell task, dbt task, etc.).

Configuring Tasks with TaskMetadata​

TaskMetadata (a dataclass in flytekit/core/base_task.py) holds the per-task configuration that maps directly to the Flyte IDL TaskMetadata proto. Its fields and defaults are:

@dataclass
class TaskMetadata(object):
    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

The __post_init__ validates several constraints at decoration time so you find out immediately rather than at runtime:

  • cache=True without cache_version raises ValueError
  • cache_serialize=True without cache=True raises ValueError
  • cache_ignore_input_vars set without cache=True raises ValueError
  • timeout as an int is converted to datetime.timedelta(seconds=timeout) automatically

When you pass timeout=300 to @task, TaskMetadata stores it as timedelta(seconds=300).

The @task Decorator Builds TaskMetadata Internally​

You don't construct TaskMetadata directly when using @task. Instead, the decorator's arguments map directly onto the dataclass fields:

@task(
    cache=True,
    cache_version="1.0",
    cache_serialize=True,
    retries=3,
    timeout=300,
    interruptible=True,
)
def my_task(x: int) -> int:
    return x * 2

Inside flytekit/core/task.py, the wrapper() function constructs this:

_metadata = TaskMetadata(
    cache=cache,
    cache_serialize=cache_serialize,
    cache_version=cache_version,
    cache_ignore_input_vars=cache_ignore_input_vars,
    retries=retries,
    interruptible=interruptible,
    deprecated=deprecated,
    timeout=timeout,
)

You can inspect the metadata on any task:

assert my_task.metadata.cache is True
assert my_task.metadata.cache_version == "1.0"
assert my_task.metadata.retries == 3

Caching: Legacy vs. Cache Object​

The cache parameter accepts either a bool (legacy) or a Cache object (defined in flytekit/core/cache.py):

from flytekit.core.cache import Cache

# Legacy: explicit version string
@task(cache=True, cache_version="v1.2.3")
def legacy_cached(x: int) -> int:
    return x + 1

# New: Cache object with auto-versioning via policies
@task(cache=Cache(version="v1.2.3", serialize=True, ignored_inputs=("timestamp",)))
def modern_cached(x: int, timestamp: str) -> int:
    return x + 1

When cache=True is passed as a bare boolean (without a cache_version), the decorator automatically constructs a Cache() object using the platform's default cache policies. Passing cache_serialize, cache_version, or cache_ignore_input_vars alongside a Cache object raises ValueError — use either interface, not both.

Local Caching​

The same cache=True flag also controls local execution caching. During local runs (calling a task in Python), flytekit uses LocalTaskCache (backed by diskcache at ~/.flyte/local-cache). The second call with the same inputs returns the cached result without re-running the function:

@task(cache=True, cache_version="v1")
def is_even(n: int) -> bool:
    return n % 2 == 0

is_even(n=4)  # executes
is_even(n=4)  # returns cached result, function not called again
is_even(n=5)  # executes (different input)

Disable local caching with FLYTE_LOCAL_CACHE_ENABLED=false, or force re-execution with FLYTE_LOCAL_CACHE_OVERWRITE=true.

The @task Decorator and Plugin Dispatch​

After building TaskMetadata, the decorator looks up the appropriate task class via TaskPlugins:

task_plugin = TaskPlugins.find_pythontask_plugin(type(task_config))
if inspect.iscoroutinefunction(fn):
    if task_plugin is PythonFunctionTask:
        task_plugin = AsyncPythonFunctionTask
task_instance = task_plugin(task_config, decorated_fn, metadata=_metadata, ...)

TaskPlugins._PYTHONFUNCTION_TASK_PLUGINS is a dict from config type → task class. When task_config=None, find_pythontask_plugin(type(None)) falls back to PythonFunctionTask. Plugins like PyTorch or Spark register their own config types:

TaskPlugins.register_pythontask_plugin(config_object_type, plugin_object_type)

When the decorated function is an async def, the decorator substitutes AsyncPythonFunctionTask for the default PythonFunctionTask — or asserts that the registered plugin is already a subclass of AsyncPythonFunctionTask.

The Task Execution Contract​

The Task base class defines three abstract methods that form the execution contract:

  • pre_execute(user_params) — called before input conversion. Returns (possibly modified) ExecutionParameters. Used by plugins like Spark to set up a session before type transformers run.
  • execute(**kwargs) — called with Python-native inputs. Returns Python-native outputs. This is where your code runs.
  • dispatch_execute(ctx, input_literal_map) — orchestrates the entire pipeline from Flyte literals to Python and back. Called both locally (via local_execute → sandbox_execute → dispatch_execute) and at runtime (directly by pyflyte-execute).

post_execute(user_params, rval) is a concrete extension point (not abstract) on PythonTask. It runs after execute() and can modify outputs or trigger cleanup. By default it's a no-op.

Two Execution Paths Through Task​

Local path (calling a task from Python, within a test or workflow local run):

task(**kwargs)
  → Task.__call__  → flyte_entity_call_handler
  → Task.local_execute          # checks LocalTaskCache, translates kwargs → LiteralMap
  → Task.sandbox_execute        # wraps ctx with task sandbox
  → Task.dispatch_execute       # the actual work
  → Promise objects returned    # wrapped for workflow composition

The local_execute step in flytekit/core/base_task.py is also where local cache hits short-circuit execution: if self.metadata.cache is True and LocalConfig.auto().cache_enabled is True, it checks LocalTaskCache.get(...) before calling sandbox_execute.

Runtime path (inside a running Flyte container, invoked by pyflyte-execute):

pyflyte-execute --inputs ... --resolver ... -- module-path task-name
  → entrypoint._dispatch_execute
  → task_def.dispatch_execute(ctx, input_literal_map)
  → outputs written to output_prefix

How PythonTask.dispatch_execute Works​

PythonTask (in flytekit/core/base_task.py) provides the concrete dispatch_execute(). Here is its pipeline, step by step:

  1. pre_execute is called. For most tasks this is a no-op, but plugin tasks (e.g., Spark) configure external resources here.

  2. Deck setup — if enable_deck=True is configured, timeline and other deck fields are attached to the execution context.

  3. Input translation — _literal_map_to_python_input(input_literal_map, exec_ctx) calls TypeEngine.literal_map_to_kwargs(...) to convert Flyte Literal values back to Python objects.

  4. execute(**native_inputs) — your function runs.

  5. post_execute — any post-processing or cleanup.

  6. Output translation — _output_to_literal_map(native_outputs, exec_ctx) converts Python values back to Literal objects. This is async internally: it uses asyncio.create_task for each output in parallel. If a conversion fails, the error message names the specific output key and the Python type involved.

  7. _write_decks — writes deck HTML if enabled.

The return value is a LiteralMap for normal tasks, or a DynamicJobSpec for dynamic tasks (see below).

Exceptions from execute() are wrapped in FlyteUserRuntimeException when running on the platform (keeping the original for local runs) so the backend can distinguish user errors from system errors.

PythonFunctionTask and ExecutionBehavior​

PythonFunctionTask (in flytekit/core/python_function_task.py) is the concrete class produced by @task. Its ExecutionBehavior enum controls what execute() does:

class ExecutionBehavior(Enum):
    DEFAULT = 1   # call self._task_function(**kwargs) directly
    DYNAMIC = 2   # compile function body into DynamicJobSpec at runtime
    EAGER = 3     # run as an eager workflow (Python-as-propeller)

DEFAULT is straightforward — execute() calls the stored _task_function with the already-translated Python kwargs.

DYNAMIC is set by @dynamic. When execute() is called on the platform, it invokes dynamic_execute(), which calls compile_into_workflow(). That function runs the user's code in a compilation context (not executing sub-tasks, but recording the DAG they form), and returns a DynamicJobSpec. The entrypoint writes this as futures.pb instead of outputs.pb, and Flyte propeller continues executing the specified sub-workflow.

Locally, dynamic_execute() detects that it's in local execution mode and runs the function body directly (executing sub-tasks inline):

def dynamic_execute(self, task_function: Callable, **kwargs) -> Any:
    ctx = FlyteContextManager.current_context()
    if ctx.execution_state and ctx.execution_state.is_local_execution():
        # runs sub-tasks inline, returns LiteralMap
        ...
    if ctx.execution_state and ctx.execution_state.mode == ExecutionState.Mode.TASK_EXECUTION:
        return self.compile_into_workflow(ctx, task_function, **kwargs)

The Nested Function Restriction​

PythonFunctionTask raises a ValueError at construction time if the task function is nested (not importable at module level) and the default resolver is in use:

if (
    not istestfunction(func=task_function)
    and isnested(func=task_function)
    and not is_functools_wrapped_module_level(task_function)
):
    raise ValueError(
        "TaskFunction cannot be a nested/inner or local function. "
        "It should be accessible at a module level for Flyte to execute it."
        ...
    )

Test modules (where the module name starts with test_) are exempt. If you need to wrap your task function with a custom decorator, use functools.wraps so the inner function's __module__ and __qualname__ remain accessible.

Async Tasks​

When @task decorates an async def function, the decorator automatically instantiates AsyncPythonFunctionTask instead of PythonFunctionTask. This class overrides __call__ to be async (using async_flyte_entity_call_handler), and provides async_execute() which awaits the task function:

async def async_execute(self, *args, **kwargs) -> Any:
    if self.execution_mode == self.ExecutionBehavior.DEFAULT:
        return await self._task_function(**kwargs)
    elif self.execution_mode == self.ExecutionBehavior.DYNAMIC:
        raise NotImplementedError

execute on AsyncPythonFunctionTask is set to loop_manager.synced(async_execute) — a synchronous wrapper that runs the coroutine in a managed event loop, so the task can still be called from the synchronous dispatch_execute pipeline.

AsyncPythonFunctionTask does not support DYNAMIC execution mode.

Eager Tasks​

EagerAsyncPythonFunctionTask (also in flytekit/core/python_function_task.py) implements the eager workflow pattern, where Python code is the orchestrator — each task invocation becomes a full Flyte execution rather than a in-process function call.

The @eager decorator creates EagerAsyncPythonFunctionTask instances, which always set metadata.is_eager=True and always use ExecutionBehavior.EAGER. Any execution_mode kwarg passed in is silently dropped.

At runtime on a Flyte backend, execute() sets up a Controller (an async worker queue), installs signal handlers for SIGINT/SIGTERM, and runs the async task function via run_with_backend(). Each sub-task or sub-workflow call the user's code awaits dispatches a new Flyte execution, with the results collected asynchronously.

Locally, EagerAsyncPythonFunctionTask uses ExecutionState.Mode.EAGER_LOCAL_EXECUTION and runs the function directly.

On failure, EagerAsyncPythonFunctionTask.get_as_workflow() wraps the eager task in an ImperativeWorkflow with an EagerFailureHandlerTask as the on_failure handler. That cleanup task terminates any still-running sub-executions tagged with the parent execution's ID.

TaskResolverMixin: Task Rehydration at Runtime​

When Flyte runs your task, it spins up a container and invokes pyflyte-execute. That command needs to know which Python object to run. TaskResolverMixin (in flytekit/core/base_task.py) is the abstract interface that handles this round-trip.

A resolver implements four abstract methods:

  • location — importable dotted path to the resolver singleton instance
  • loader_args(settings, task) — returns a list of strings that identify this task to load_task
  • load_task(loader_args) — reconstructs the Task object from those strings
  • get_all_tasks() — returns all tasks known to this resolver

The get_default_command() method in PythonAutoContainerTask combines these into the container's args field:

container_args = [
    "pyflyte-execute",
    "--inputs", "{{.input}}",
    "--output-prefix", "{{.outputPrefix}}",
    ...
    "--resolver", self.task_resolver.location,
    "--",
    *self.task_resolver.loader_args(settings, self),
]

default_task_resolver (in flytekit/core/python_auto_container.py) looks up the task module and function name when loader_args is called, and uses importlib.import_module + attribute lookup in load_task. This is why module-level functions are required.

For tasks defined inside a workflow function body, flytekit uses the workflow object itself as the resolver. The args become an index:

# Serialized args for tasks nested inside my_wf:
["--resolver", "myproject.workflows.my_wf", "--", "0"]
["--resolver", "myproject.workflows.my_wf", "--", "1"]

my_wf (as a WorkflowBase) implements TaskResolverMixin and tracks all tasks created during its execution, indexed by position.

kwtypes for Non-Function Tasks​

When building a non-function task (like SQLTask) that doesn't auto-detect its interface from a Python function signature, use kwtypes to define the typed interface:

from flytekit.core.base_task import kwtypes
from flytekit.types.schema import FlyteSchema

sql = SQLTask(
    "my-query",
    query_template="SELECT * FROM table WHERE ds = '{{ .Inputs.ds }}' LIMIT 10",
    inputs=kwtypes(ds=datetime.datetime),
    outputs=kwtypes(results=FlyteSchema),
    metadata=TaskMetadata(retries=2, cache=True, cache_version="0.1"),
)

kwtypes(**kwargs) simply returns an OrderedDict[str, Type] preserving argument order. It's a convenience wrapper for the Interface constructor's inputs= and outputs= arguments.

Skipping Output Upload with IgnoreOutputs​

In distributed training jobs, only the rank-0 worker's outputs are meaningful. Workers with rank > 0 can raise IgnoreOutputs from their execute() method to skip the output upload step entirely:

from flytekit import task
from flytekit.core.base_task import IgnoreOutputs

@task
def distributed_worker(rank: int) -> str:
    # ... do training work ...
    if rank != 0:
        raise IgnoreOutputs(f"Worker {rank} outputs ignored")
    return "primary result"

In flytekit/bin/entrypoint.py, _dispatch_execute catches IgnoreOutputs specifically and returns early without writing outputs.pb:

except FlyteUserRuntimeException as e:
    if isinstance(e.value, IgnoreOutputs):
        logger.warning(f"User-scoped IgnoreOutputs received! Outputs.pb will not be uploaded. reason {e}!!")
        return

IgnoreOutputs must propagate out of execute() as-is — dispatch_execute catches it at the post_execute step, and the entrypoint then catches it wrapped inside FlyteUserRuntimeException.

Runtime Execution Flow​

When pyflyte-execute is invoked by the Flyte platform, the execution flow through the entrypoint (flytekit/bin/entrypoint.py) is:

pyflyte-execute CLI arguments parsed
  → _execute_task() → setup_execution()
  → _dispatch_execute(ctx, load_task, inputs_path, output_prefix)
      Step 1: Download inputs.pb from blob store → LiteralMap
      Step 2: task_def.dispatch_execute(ctx, literal_map)
               → PythonTask.dispatch_execute:
                   pre_execute → _literal_map_to_python_input
                   → execute(**native_inputs)
                   → post_execute → _output_to_literal_map
      Step 3a: LiteralMap  → write outputs.pb to output_prefix
      Step 3b: DynamicJobSpec → write futures.pb to output_prefix
      Step 3c: IgnoreOutputs raised → return, no outputs.pb
      Step 3d: Exception → write error.pb (ContainerError proto)

By default, pyflyte-execute exits with code 0 after writing error.pb (Flyte propeller reads the error file). Set FLYTE_FAIL_ON_ERROR=true to make it exit with code 1 instead — this is required by Databricks and AWS Batch plugins where the orchestrator infers failure only from the exit code.

Configuration Reference​

Environment VariableDefaultEffect
FLYTE_LOCAL_CACHE_ENABLEDtrueEnables local task caching during local execution
FLYTE_LOCAL_CACHE_OVERWRITEfalseForces re-execution even on cache hit
FLYTE_INTERNAL_IMAGE—Overrides the default container image for all tasks
FLYTE_FAIL_ON_ERRORfalseExit code 1 on task failure (required by some cloud batch backends)
_F_EE_ROOT—Internal: passes root eager execution name to nested eager tasks

The local cache location (~/.flyte/local-cache) is managed by LocalTaskCache in flytekit/core/local_cache.py and controlled via LocalConfig (read from environment or config file).