Skip to main content

Task authoring and execution

Task authoring starts with a typed callable​

When you write a regular Python function and call it inside a workflow, Flyte needs both the function’s Python behavior and a Flyte interface describing its inputs and outputs. The @task decorator creates a PythonFunctionTask for that purpose and derives its interface from the function’s type annotations.

from flytekit import task

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

PythonFunctionTask.__init__ calls transform_function_to_interface for the function, removes any names listed in ignore_input_vars, and derives the task name with extract_task_module. The resulting task is callable: Task.__call__ routes calls through flyte_entity_call_handler, which allows the same task object to participate in workflow compilation and local execution.

The task hierarchy determines how much behavior you get:

Task
└── PythonTask
    └── PythonAutoContainerTask
        └── PythonFunctionTask
  • Task is the FlyteIDL-oriented base abstraction. It stores the task type, name, typed interface, metadata, security context, and documentation, and registers the instance in FlyteEntities.entities. It does not provide a Python-native interface.
  • PythonTask adds an Interface, task configuration, environment variables, and Deck settings. It also defines the Python-value-to-Flyte-literal execution path.
  • PythonAutoContainerTask adds automatic container and resource configuration; it is the intermediate class used by PythonFunctionTask and PythonInstanceTask.
  • PythonFunctionTask associates the task with a user callable and supplies the default, dynamic, and eager execution behaviors.

Use PythonFunctionTask directly when you are implementing a task type or plugin rather than decorating a function. Its constructor requires a task configuration and a callable:

from flytekit.core.python_function_task import PythonFunctionTask


def foo(i: int):
    pass

pytask = PythonFunctionTask(None, foo, None, environment={"BAZ": "baz"})

For ordinary user tasks, prefer @task; the decorator also selects an async task class for coroutine functions and looks up a registered task plugin based on the type of task_config.

Function location and task names​

The default resolver rehydrates a serialized task by importing its module and looking up the task name. Consequently, a normal PythonFunctionTask must refer to a module-level function. With the default resolver, PythonFunctionTask.__init__ rejects nested or local functions unless the function is in a test module or has been handled with a supported functools wrapper. The error explicitly recommends a module-level function, functools.wraps/update_wrapper, or a custom TaskResolverMixin.

A task’s name is derived from its defining module and function. If the task has been instantiated in another module, the name property can prefix the instantiated location, so choose globally unique task names for the Flyte project, domain, and version in which they are registered.

Configure execution metadata​

Task metadata belongs on the task, not in the function body. For example, enable local and hosted caching with a non-empty cache version:

from flytekit import task

@task(cache=True, cache_serialize=True, cache_version="1.0")
def expensive_step(i: str):
    print(i)

TaskMetadata.__post_init__ validates the relationships between these fields:

  • cache=True requires cache_version.
  • cache_serialize=True requires cache=True.
  • cache_ignore_input_vars requires cache=True.

cache_ignore_input_vars is a tuple of input names omitted when the cache key is calculated. retries sets the retry count, timeout accepts either integer seconds or a datetime.timedelta, and interruptible records whether the task may run on interruptible/preemptible capacity. pod_template_name identifies an existing cluster PodTemplate, while generates_deck and is_eager are carried into the serialized task metadata.

The metadata object converts these Python settings to Flyte’s task model through TaskMetadata.to_taskmetadata_model(). That method maps cache to discoverable, cache_version to discovery_version, retries to a RetryStrategy, and forwards timeout, interruptibility, deprecation text, Deck generation, PodTemplate name, ignored cache inputs, and eager state.

Local caching is performed by Task.local_execute. It first translates the call arguments into a LiteralMap; when metadata.cache is true and LocalConfig.auto().cache_enabled is true, it checks LocalTaskCache using the task name, cache version, literal inputs, and ignored input names. A cache hit returns the stored literal outputs without calling sandbox_execute; a miss executes the task and stores the resulting literal map. The FLYTE_LOCAL_CACHE_ENABLED setting controls whether that local cache is enabled.

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

With the same local input, the first call is a cache miss and later calls use the cached output. A different input produces a different cache key unless that input is listed in cache_ignore_input_vars.

Configure the container and runtime​

@task forwards container and execution settings to PythonAutoContainerTask. A task can specify an image, environment, resources, secrets, an accelerator, shared memory, a PodTemplate, and other container settings. A concrete task declaration from the test suite uses both a custom image and a PodTemplate:

from flytekit import task
from flytekit.models.pod import PodTemplate
from kubernetes.client import V1Container, V1PodSpec

@task(
    container_image="repo/image:0.0.0",
    pod_template=PodTemplate(
        primary_container_name="primary",
        labels={"lKeyA": "lValA"},
        annotations={"aKeyA": "aValA"},
        pod_spec=V1PodSpec(
            containers=[V1Container(name="primary")]
        ),
    ),
    pod_template_name="A",
)
def func_with_pod_template(i: str):
    print(i + "a")

Use either the combined resources form or separate request/limit settings. The combined form accepts tuples for request and limit values:

from flytekit import Resources, task

@task(resources=Resources(cpu=("1", "2"), mem=(1024, 2048), gpu=1))
def my_task():
    pass

PythonAutoContainerTask rejects resources when it is combined with separate requests or limits. Environment variables supplied to the task are merged with the environment in SerializationSettings; in the container test, {"FOO": "bar"} from serialization settings and {"BAZ": "baz"} from the task both appear in the generated container.

The generated container command includes the default resolver and loader arguments. For the PodTemplate example, the command has the form:

pyflyte-execute --inputs {{.input}} --output-prefix {{.outputPrefix}} \
  --resolver flytekit.core.python_auto_container.default_task_resolver -- \
  task-module tests...test_python_function_task \
  task-name func_with_pod_template

The exact module argument is the defining module of the task. TaskResolverMixin.loader_args() supplies these identifying arguments, while the resolver’s load_task() uses them to reconstruct the task in the execution container.

Understand the execution lifecycle​

Do not call dispatch_execute directly in application code. Call the task object normally, or use its local execution path through the normal Flyte call handling. At runtime, flytekit.bin.entrypoint invokes dispatch_execute after the task has been rehydrated.

For a PythonTask, dispatch_execute performs this sequence:

input LiteralMap
  -> _literal_map_to_python_input (TypeEngine)
  -> pre_execute
  -> execute(**native_inputs)
  -> post_execute
  -> _output_to_literal_map (TypeEngine)
  -> output LiteralMap

PythonTask._literal_map_to_python_input calls TypeEngine.literal_map_to_kwargs with the task’s Python input types. The default pre_execute and post_execute methods return their input unchanged, but plugin tasks override pre_execute to initialize runtime systems. The Ray plugin, for example, initializes Ray using its task configuration and returns the supplied ExecutionParameters.

PythonFunctionTask.execute is the function boundary. In default mode it simply calls the wrapped function:

def execute(self, **kwargs) -> Any:
    if self.execution_mode == self.ExecutionBehavior.DEFAULT:
        return self._task_function(**kwargs)

After execution, _output_to_literal_map maps Python results to the declared output names and uses TypeEngine.async_to_literal for conversion. It handles zero outputs, single outputs, multiple outputs, and the single-element NamedTuple convention. Conversion failures include the task name and output position in the raised error.

For local calls, Task.local_execute first translates native arguments or Promise values into literals, then calls sandbox_execute. sandbox_execute changes the execution parameters to a task sandbox and delegates to dispatch_execute. The local path converts the returned literals back into Promise objects; a task with no outputs returns a VoidPromise.

Suppress outputs explicitly​

IgnoreOutputs is an exception used by tasks whose generated outputs may safely be ignored, such as distributed workers that are not responsible for publishing the result:

from flytekit.core.base_task import IgnoreOutputs

raise IgnoreOutputs()

The task entrypoint recognizes this specific exception and returns without writing output data. It must be raised in the execution path where the entrypoint can recognize it; wrapping it in another unrelated exception changes that behavior.

Choose an execution mode​

PythonFunctionTask.ExecutionBehavior has three values: DEFAULT, DYNAMIC, and EAGER. The decorators provide the supported user-facing forms for the latter two.

Default tasks​

A default task executes its wrapped function once with native Python inputs after Flyte has performed literal conversion:

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

Dynamic tasks​

Use @dynamic when the function body creates task calls and the graph depends on runtime values:

from flytekit import dynamic, task

@task
def t1(a: int) -> str:
    a = a + 2
    return "fast-" + str(a)

@dynamic
def my_subwf(a: int) -> list[str]:
    s = []
    for i in range(a):
        s.append(t1(a=i))
    return s

@dynamic creates a PythonFunctionTask with execution_mode=DYNAMIC. In a real task execution, PythonFunctionTask.dynamic_execute compiles the function into a workflow. compile_into_workflow serializes the generated workflow and returns a DynamicJobSpec containing its task templates, nodes, outputs, and subworkflows. During local execution, dynamic_execute runs the generated PythonFunctionWorkflow and converts its native results into a literal map instead.

Only dynamic tasks may provide node_dependency_hints; the constructor rejects hints on a default-mode task because static workflows already expose their dependencies.

Async and eager tasks​

AsyncPythonFunctionTask is used for async functions. Its asynchronous __call__ uses async_flyte_entity_call_handler, and its async_execute awaits the wrapped function. Async dynamic execution is not supported: async_execute raises NotImplementedError when its execution mode is DYNAMIC.

Use @eager for an async function that orchestrates Flyte entities through backend executions:

from flytekit import eager, task

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

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

EagerAsyncPythonFunctionTask forces EAGER execution and sets TaskMetadata.is_eager=True. Its remote execute method creates a Controller worker queue when necessary; run_with_backend runs the user coroutine with eager execution state and the controller in context. The FLYTE_EAGER_ROOT environment variable can override the root execution tag used for nested eager executions.

get_as_workflow() adapts an eager task to an ImperativeWorkflow and adds an EagerFailureHandlerTask as its failure handler. That handler uses EagerFailureTaskResolver and, during remote dispatch, queries Admin for unfinished child executions tagged with the parent eager execution and terminates them. Its execute method is not intended to run; cleanup is implemented by dispatch_execute.

Use instance tasks when there is no user function​

A platform-defined task with a Python interface but no user function should derive from PythonInstanceTask, not PythonFunctionTask:

x = MyInstanceTask(name="x", task_config=..., task_type="python-task")
result = x(a=5)

PythonInstanceTask exists for task classes whose platform-specific execute implementation is supplied by the class. It inherits the automatic container behavior and captures the module-level variable so the loader can rehydrate the correct instance. PythonFunctionTask, by contrast, is specifically the base for extensions that execute a user-provided callable.

At the lowest level, custom task classes implement Task.dispatch_execute, Task.pre_execute, and Task.execute. Python-native task classes normally inherit the literal conversion and lifecycle implementation from PythonTask, overriding execute and, when needed, the hooks or serialization methods such as get_custom.

Extend function tasks with plugins and resolvers​

Task plugins associate a configuration type with a PythonFunctionTask subclass. A plugin follows the same pattern as the Ray integration:

class RayFunctionTask(PythonFunctionTask):
    _RAY_TASK_TYPE = "ray"

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

    def pre_execute(self, user_params):
        init_params = {"address": self._task_config.address}
        ray.init(**init_params)
        return user_params

The plugin registers its configuration and task class with TaskPlugins.register_pythontask_plugin(RayJobConfig, RayFunctionTask). The @task decorator can then select RayFunctionTask when its task_config is a RayJobConfig. A pre_execute override must return the original or modified ExecutionParameters; PythonTask.dispatch_execute uses that returned value when it creates the execution context.

Use TaskResolverMixin when module-and-name loading is insufficient, such as for a custom task representation or a task embedded in a special execution environment. The mixin requires a resolver location, a name, load_task(loader_args), loader_args(settings, task), and get_all_tasks(). The resolver location is written into the container command after --resolver, and load_task reconstructs the task from the following loader arguments.

Test task behavior locally​

Tasks are ordinary callable objects in local execution, so you can exercise their behavior without invoking the hosted entrypoint. For isolated unit tests, flytekit.core.testing.task_mock temporarily replaces a task’s execute method with a MagicMock and restores it when the context exits:

from flytekit.core.testing import task_mock

with task_mock(add_one) as mocked:
    add_one(x=1)
    mocked.assert_called_once()

The context manager accepts a PythonTask, workflow, or reference entity and rejects unrelated objects. Local cache tests should account for the FLYTE_LOCAL_CACHE_ENABLED setting: even a task declared with cache=True uses LocalTaskCache only when the local cache is enabled and LocalConfig.auto() reports it as enabled.