Skip to main content

Task authoring and execution

Flyte tasks are the atomic units of computation in any workflow. Understanding how flytekit models them — and the layered class hierarchy beneath the @task decorator — unlocks everything from basic caching to writing your own distributed training plugins.

The Class Hierarchy​

Every task in flytekit inherits from one of several base classes, each adding a different layer of capability:

Task                          (flytekit/core/base_task.py)
 └── PythonTask               (flytekit/core/base_task.py)
      └── PythonAutoContainerTask  (flytekit/core/python_auto_container.py)
           └── PythonFunctionTask  (flytekit/core/python_function_task.py)

Task is the lowest-level base, closest to the FlyteIDL TaskTemplate. It holds task_type, name, a TypedInterface (IDL-level, not Python-native), TaskMetadata, a security context, and documentation. Every Task instance appends itself to FlyteEntities.entities at construction time so the serialization toolchain can discover it later. The abstract methods dispatch_execute(), pre_execute(), and execute() define the execution contract.

PythonTask extends Task with a Python-native Interface (inputs and outputs typed as actual Python types, not IDL literals), deck support, environment variables, and task_config. It provides the full dispatch_execute() implementation that drives execution: translating input literals to Python values, calling pre_execute → execute → post_execute, converting outputs back to a LiteralMap via TypeEngine, and writing Deck HTML. Task subclasses like SQLTask and ContainerTask — which don't have a Python function body — stop here.

PythonFunctionTask extends PythonAutoContainerTask (itself a PythonTask) and is the class produced by the @task decorator. Its constructor calls transform_function_to_interface() to auto-detect the decorated function's signature, and extract_task_module() to derive the task's fully-qualified name. All @task-decorated functions are ultimately instances of PythonFunctionTask or a registered subclass.

Configuring a Task with TaskMetadata​

TaskMetadata (in flytekit/core/base_task.py) is the dataclass that controls retries, caching, timeouts, and related settings. You don't usually instantiate it directly — the @task decorator creates it for you — but knowing its fields is essential:

from flytekit import task
import datetime

@task(
    retries=3,
    timeout=datetime.timedelta(hours=2),
    interruptible=True,
    deprecated="Use the new_process task instead.",
    cache=True,
    cache_version="v2",
    cache_serialize=True,
)
def process(data: str) -> str:
    ...

The TaskMetadata.__post_init__ method enforces several validation rules at decoration time:

  • cache=True requires a non-empty cache_version; omitting it raises ValueError: "Caching is enabled cache=True but cache_version is not set."
  • cache_serialize=True requires cache=True
  • cache_ignore_input_vars also requires cache=True
  • timeout accepts either an int (seconds, auto-converted to timedelta) or a datetime.timedelta directly

The Cache Object​

The legacy cache_version string approach requires you to manually bump the version whenever the function body changes. The newer Cache object (in flytekit/core/cache.py) can compute a version hash automatically from the function source and container image:

from flytekit import task
from flytekit.core.cache import Cache

@task(cache=Cache(serialize=True))
def expensive_transform(x: int) -> int:
    return x * 2

Cache.get_version() runs one or more CachePolicy implementations and SHA-256 hashes the combined result. You can still pin a version explicitly with Cache(version="v1.2.3"), and you can exclude specific inputs from the cache key with Cache(ignored_inputs=("timestamp",)).

When you pass cache=True (the boolean) with no cache_version, the @task wrapper automatically constructs a Cache() object with default policies — so the shorthand still works, it just triggers automatic versioning.

The @task Decorator​

The @task decorator in flytekit/core/task.py is the primary user-facing entry point. Its simplest form requires only a type-annotated function:

from flytekit import task

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

The decorator extracts the function's signature through transform_function_to_interface(), constructs a TaskMetadata from the decorator arguments, looks up the appropriate task class via TaskPlugins.find_pythontask_plugin(type(task_config)), and instantiates it. The resulting object is a PythonFunctionTask instance — add is no longer a plain Python function after decoration.

Common Decorator Parameters​

Beyond the basic signature, @task accepts a rich set of parameters that all feed into either TaskMetadata or the PythonAutoContainerTask constructor:

from flytekit import task, Resources, Secret
from flytekit.image_spec import ImageSpec

my_image = ImageSpec(packages=["pandas", "scikit-learn"], registry="myrepo")

@task(
    container_image=my_image,          # ImageSpec or string URI
    requests=Resources(cpu="2", mem="4Gi"),
    limits=Resources(cpu="4", mem="8Gi"),
    # OR use resources= for same request and limit:
    # resources=Resources(cpu="2", mem="4Gi"),
    environment={"PYTHONPATH": "/app"},
    retries=2,
    timeout=300,                       # seconds, or timedelta
    interruptible=True,
    enable_deck=True,
    secret_requests=[Secret(group="my-secrets", key="api-key")],
)
def train_model(data_path: str) -> float:
    ...

When container_image is omitted, PythonAutoContainerTask falls back to the FLYTE_INTERNAL_IMAGE environment variable.

The resources shorthand sets both request and limit to the same value. When you need different request and limit, use requests and limits separately. For Pod plugin tasks, these values apply only to the primary container.

Async Functions​

If you decorate an async def function with @task, flytekit detects this via inspect.iscoroutinefunction() and substitutes AsyncPythonFunctionTask for PythonFunctionTask:

@task
async def fetch_data(url: str) -> str:
    import aiohttp
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as resp:
            return await resp.text()

AsyncPythonFunctionTask (in flytekit/core/python_function_task.py) overrides __call__ to return a coroutine via async_flyte_entity_call_handler, and its async_execute() method simply awaits the task function directly.

The Execution Lifecycle​

When flytekit runs a task — whether locally or on the Flyte platform — the path goes through PythonTask.dispatch_execute():

  1. pre_execute(user_params) — called before input conversion. The default implementation in PythonTask returns user_params unchanged. Plugin tasks override this to set up resources (e.g., the Spark plugin establishes a SparkSession here so that it's available before any type transformers run).

  2. Input conversion — _literal_map_to_python_input() calls TypeEngine.literal_map_to_kwargs() to convert the protobuf LiteralMap to a Python dict.

  3. execute(**kwargs) — the actual user function runs here. In PythonFunctionTask, the DEFAULT execution mode simply calls self._task_function(**kwargs).

  4. post_execute(user_params, rval) — called with the raw return value. The default is a no-op (return rval). Plugin tasks can use this for cleanup or output transformation.

  5. Output conversion — _output_to_literal_map() uses TypeEngine.async_to_literal() (gathered with asyncio) to convert Python values back to a LiteralMap.

For local execution, Task.local_execute() adds an additional wrapping layer: it handles the optional LocalTaskCache lookup/store (controlled by the FLYTE_LOCAL_CACHE_ENABLED and FLYTE_LOCAL_CACHE_OVERWRITE environment variables), then delegates to sandbox_execute() → dispatch_execute().

IgnoreOutputs​

In distributed scenarios, tasks running on non-primary workers may not produce meaningful outputs. Raising IgnoreOutputs (from flytekit/core/base_task.py) signals flytekit's entrypoint (flytekit/bin/entrypoint.py) to skip writing the outputs protobuf entirely:

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

@task
def worker_task(rank: int, data: str) -> str:
    result = run_computation(rank, data)
    if rank != 0:
        raise IgnoreOutputs()  # only rank 0 should produce outputs
    return result

The PyTorch plugin uses exactly this pattern — non-rank-0 workers raise IgnoreOutputs in post_execute() after the training run completes.

Three Execution Modes​

PythonFunctionTask.ExecutionBehavior is an enum with three values that fundamentally change how the function body is interpreted.

DEFAULT​

This is the standard @task. The function body runs as-is on the container, execute() calls self._task_function(**kwargs) directly, and the return values become the task outputs.

DYNAMIC​

A @dynamic task's function body runs at execution time on the Flyte cluster, but instead of producing output data, it produces a workflow specification (a DynamicJobSpec) describing which tasks to run next. The @dynamic decorator in flytekit/core/dynamic_workflow_task.py is simply:

dynamic = functools.partial(task.task, execution_mode=PythonFunctionTask.ExecutionBehavior.DYNAMIC)

This means the full @task decorator signature applies — you can pass container_image, resources, etc.

from flytekit import task, dynamic
import typing

@task
def square(n: int) -> int:
    return n * n

@dynamic
def compute_squares(count: int) -> typing.List[int]:
    return [square(n=i) for i in range(count)]

When compute_squares runs on the Flyte platform (mode TASK_EXECUTION), PythonFunctionTask.dynamic_execute() calls compile_into_workflow(), which compiles the function body into a PythonFunctionWorkflow and serializes it to a DynamicJobSpec. The platform then executes that workflow spec as a sub-workflow, running square for each value of i.

Locally, the same function runs like a normal workflow — the PythonFunctionWorkflow actually executes in LOCAL_DYNAMIC_TASK_EXECUTION mode.

Two important restrictions apply to dynamic tasks:

  • Keep node counts low. The dynamic_workflow_task.py module explicitly warns to keep dynamic workflows under 50 tasks. Each loop iteration that calls a task creates a node.
  • You can iterate over Python values, but not over Promise objects. In the @dynamic body, range(count) works because count is a real Python int at execution time. Iterating over the output of another task call — which would be a Promise — would fail.

EAGER​

Eager workflows (@eager) are a third mode where Python becomes the orchestrator. Every task or workflow called from inside an @eager function triggers a live Flyte execution, and the result is awaited asynchronously. The EagerAsyncPythonFunctionTask class (in flytekit/core/python_function_task.py) handles this.

from flytekit import task, eager

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

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

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

At runtime on the Flyte platform, EagerAsyncPythonFunctionTask.execute() constructs a Controller (the worker queue), sets up SIGINT/SIGTERM signal handlers, and calls loop_manager.run_sync(self.async_execute, ...). Each add_one(x=x) call inside the async function body creates a real execution on the cluster and awaits its result.

When the eager workflow itself is registered, get_as_workflow() wraps it in an ImperativeWorkflow with an EagerFailureHandlerTask as the on_failure handler. If the eager task fails, the failure handler queries FlyteAdmin for all running sub-executions tagged with the parent execution ID and terminates them.

A few constraints apply to eager workflows:

  • Eager and dynamic modes don't mix: AsyncPythonFunctionTask.async_execute() raises NotImplementedError for DYNAMIC execution mode.
  • The execution_mode kwarg is silently removed by EagerAsyncPythonFunctionTask.__init__ — it always uses EAGER mode.
  • Flyte conditionals aren't supported inside @eager; use plain Python if statements instead.

TaskPlugins: Extending the Task System​

When you write a Flyte plugin — like the Ray, Spark, or PyTorch plugins — you register a new task class that handles a specific task_config type. TaskPlugins (in flytekit/core/task.py) is the global registry:

from flytekit.extend import TaskPlugins, PythonFunctionTask

class MyConfig:
    def __init__(self, cluster: str):
        self.cluster = cluster

class MyFunctionTask(PythonFunctionTask[MyConfig]):
    _MY_TASK_TYPE = "my-plugin-task"

    def __init__(self, task_config: MyConfig, task_function, **kwargs):
        super().__init__(
            task_config=task_config,
            task_function=task_function,
            task_type=self._MY_TASK_TYPE,
            **kwargs,
        )

    def get_custom(self, settings) -> dict:
        return {"cluster": self.task_config.cluster}

    def pre_execute(self, user_params):
        # Set up anything needed before execution
        return user_params

TaskPlugins.register_pythontask_plugin(MyConfig, MyFunctionTask)

Once registered, @task(task_config=MyConfig("production")) automatically uses MyFunctionTask instead of the default PythonFunctionTask. The @task decorator calls TaskPlugins.find_pythontask_plugin(type(task_config)) to perform this lookup; if no match is found, it falls back to PythonFunctionTask.

Registration is global and collision raises TypeError — you cannot register two different classes for the same config type. Registering the same class twice is a no-op.

The Ray plugin demonstrates this pattern concisely:

# from plugins/flytekit-ray/flytekitplugins/ray/task.py
class RayFunctionTask(PythonFunctionTask):
    _RAY_TASK_TYPE = "ray"

    def __init__(self, task_config: RayJobConfig, task_function: Callable, **kwargs):
        super().__init__(
            task_config=task_config,
            task_function=task_function,
            task_type=self._RAY_TASK_TYPE,
            **kwargs,
        )

TaskPlugins.register_pythontask_plugin(RayJobConfig, RayFunctionTask)

PythonInstanceTask for Object-Based Tasks​

When the task isn't a decorated function but a module-level object (think a DBT run configuration, a DuckDB query, or a shell script), use PythonInstanceTask as the base class:

from flytekit.core.python_function_task import PythonInstanceTask

class MyInstanceTask(PythonInstanceTask[MyConfig]):
    def __init__(self, name: str, config: MyConfig, **kwargs):
        super().__init__(name=name, task_config=config, **kwargs)

    def execute(self, **kwargs):
        # Custom logic here
        ...

my_task = MyInstanceTask(name="my.task", config=MyConfig(cluster="gpu"))
# Invoke like any other task

The PythonInstanceTask extends PythonAutoContainerTask and is designed for the case where the task definition is the object itself, not a wrapped function.

TaskResolverMixin: Bridging Serialization and Runtime​

For tasks that run in containers (the common case), flytekit needs to serialize the task during registration and then reconstruct it when the container starts. TaskResolverMixin (in flytekit/core/base_task.py) defines this two-phase contract through abstract methods.

The container command that flytekit generates looks like:

pyflyte-execute --inputs s3://path/inputs.pb --output-prefix s3://outputs/location \
  --raw-output-data-prefix /tmp/data \
  --resolver flytekit.core.python_auto_container.default_task_resolver \
  -- \
  task-module repo_root.workflows.example task-name my_function

The DefaultTaskResolver (in flytekit/core/python_auto_container.py) implements this strategy:

  • loader_args(settings, task): records the module path and the task's name in the module
  • load_task(loader_args): calls importlib.import_module and looks up the function by name

This is why task functions must be module-level. If you decorate a nested or inner function with @task and the default resolver is in use, PythonFunctionTask.__init__ raises:

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 (names starting with test_) are exempt from this check. If you use custom decorators on task functions, apply functools.wraps to preserve the module-level lookup. To go further and use a completely different loading strategy, implement TaskResolverMixin and pass your resolver to @task(task_resolver=my_resolver).

Testing Tasks​

Two utilities in flytekit/core/testing.py make unit testing tasks straightforward without running the full task logic.

task_mock() is a context manager that replaces execute() with a MagicMock for the duration of the with block:

from flytekit.core.testing import task_mock

@task
def fetch_from_db(query: str) -> list:
    # expensive DB call
    ...

@workflow
def my_workflow(q: str) -> list:
    return fetch_from_db(query=q)

def test_workflow():
    with task_mock(fetch_from_db) as mock:
        mock.return_value = [{"id": 1}]
        result = my_workflow(q="SELECT *")
    assert result == [{"id": 1}]

patch() is the decorator equivalent, injecting a MagicMock as the first argument to the test function:

from flytekit.core.testing import patch

@patch(fetch_from_db)
def test_workflow_with_patch(mock_fetch):
    mock_fetch.return_value = [{"id": 2}]
    result = my_workflow(q="SELECT *")
    assert result == [{"id": 2}]

Both utilities work by temporarily swapping t.execute — the actual function that dispatch_execute() calls — with a mock, then restoring it when the block exits.

For local caching during tests, note that LocalTaskCache reads FLYTE_LOCAL_CACHE_ENABLED and FLYTE_LOCAL_CACHE_OVERWRITE fresh at each local_execute() call via LocalConfig.auto(). If your tests run sequentially and a cached result from one test bleeds into another, set FLYTE_LOCAL_CACHE_OVERWRITE=true or FLYTE_LOCAL_CACHE_ENABLED=false.

Common Pitfalls​

Nested functions with the default resolver. The constraint is enforced at decoration time, not at serialization time. You'll see the ValueError as soon as Python imports the module that defines the nested @task.

Cache version required for cache=True. TaskMetadata.__post_init__ raises immediately. The safest migration path is to switch to Cache() (no version argument), which auto-hashes the function source.

cache_serialize without cache. This also raises a ValueError at instantiation. The three cache-related parameters are validated together.

node_dependency_hints on a static task. PythonFunctionTask.__init__ raises ValueError if node_dependency_hints is set but execution_mode != DYNAMIC. The parameter is only meaningful on @dynamic tasks, where it allows pre-registering launch plans that flytekit can't discover statically.

disable_deck is deprecated. Since flytekit 1.10.0, use enable_deck=True instead. Passing both disable_deck and enable_deck raises ValueError. Accessing the disable_deck property on a PythonTask instance emits a DeprecationWarning.

Unrecognized @task kwargs raise immediately. The @task wrapper calls raise ValueError(f"Unrecognized argument(s) for task: {kwargs.keys()}") for any leftover keyword arguments. A common mistake is pod_template=... being spelled wrong (e.g., template=...).

pickle_untyped=True is for prototyping only. This flag tells PythonFunctionTask to pickle any outputs that lack type annotations instead of failing. It loses type safety and is explicitly not recommended for production use.