The Task Execution Lifecycle

Defining the task object

When you decorate an asynchronous function with TaskEnvironment.task, the decorator does not leave the function as a plain callable. It constructs an AsyncFunctionTaskTemplate, records it in the environment's _tasks dictionary, and returns that task template. The template carries the callable together with the environment's image, resources, cache policy, retries, timeout, secrets, and other task metadata.

env = flyte.TaskEnvironment(name="my_env", image="my_image", resources=Resources(cpu="1", memory="1Gi"))

@env.task
async def my_task():
    pass

This is the usage shown in TaskEnvironment's definition in _task_environment.py. Internally, TaskEnvironment.task derives the task name as self.name + "." + func.__name__, builds the native interface with NativeInterface.from_callable(func), and instantiates AsyncFunctionTaskTemplate with func=func. If the environment has a plugin_config, TaskPluginRegistry.find(config_type=type(self.plugin_config)) selects a registered subclass instead; an unregistered configuration type raises ValueError.

AsyncFunctionTaskTemplate is publicly re-exported by flyte.extend, so it is also the extension point for task-plugin implementations:

from flyte.extend import AsyncFunctionTaskTemplate


class HookedTask(AsyncFunctionTaskTemplate):
    async def pre(self, *args, **kwargs):
        return {"started": True}

    async def post(self, return_vals):
        return return_vals

The inherited TaskTemplate.pre and TaskTemplate.post methods are asynchronous extension points. The default pre returns {}, and the default post returns its argument. A plugin subclass can override them (and can be selected by registering it with TaskPluginRegistry.register).

Invocation chooses local or controller execution

The same task object follows two different paths depending on the current Flyte context:

no task context  -> AsyncFunctionTaskTemplate.__call__ -> forward -> func
inside a task   -> AsyncFunctionTaskTemplate.__call__ -> controller.submit
                                                   -> execute -> pre -> func -> post

TaskTemplate.__call__ checks internal_ctx().is_task_context(). Outside a task context, it calls forward. AsyncFunctionTaskTemplate.forward only calls the wrapped function:

def forward(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
    return self.func(*args, **kwargs)

For an asynchronous function, forward returns the coroutine without awaiting it; the caller's await performs the await. It does not call pre or post. Consequently, directly calling a task during ordinary local Python execution exercises the function body but not the execution hooks.

Inside a Flyte task context, __call__ obtains the controller with get_controller() and submits the task. For an ordinary coroutine function it returns controller.submit(self, *args, **kwargs). When AsyncFunctionTaskTemplate.__post_init__ finds that func is not a coroutine function, it sets _call_as_synchronous = True; __call__ then uses controller.submit_sync and waits for the resulting future. This permits TaskEnvironment.task to accept both the FunctionTypes supported by the decorator, while keeping synchronous downstream calls usable from an async task.

The local Run path also creates a task context and submits through a local controller. _run.py checks the same _call_as_synchronous flag, using submit_sync for a synchronous task and submit otherwise. This is different from directly invoking the task outside a run: a task executed through the local controller reaches the controller's task-dispatch path and therefore uses the task's execution lifecycle.

Do not call execute as a replacement for a normal task invocation. AsyncFunctionTaskTemplate.execute asserts that internal_ctx().data.task_context is present and raises an assertion failure otherwise. The runtime controller is responsible for invoking it in the appropriate context.

The pre–function–post lifecycle

AsyncFunctionTaskTemplate.execute is the method that runs the wrapped function in a task context. Its ordering is explicit:

1. Read the current internal context.
2. Await self.pre(*args, **kwargs).
3. Replace the current task context's data with the dictionary returned by pre.
4. Call the wrapped function, awaiting it when inspect.iscoroutinefunction(self.func) is true.
5. Await self.post(v) inside the replacement context.
6. Return v.

The core implementation in _task.py is:

ctx = internal_ctx()
assert ctx.data.task_context is not None, "Function should have already returned if not in a task context"
ctx_data = await self.pre(*args, **kwargs)
tctx = ctx.data.task_context.replace(data=ctx_data)
with ctx.replace_task_context(tctx):
    if iscoroutinefunction(self.func):
        v = await self.func(*args, **kwargs)
    else:
        v = self.func(*args, **kwargs)
    await self.post(v)
return v

This makes pre useful for producing task-context data that the function reads through Flyte's internal context. The returned dictionary replaces the existing data payload; it is not merged with that payload. Return a dictionary (or {}), matching the TaskTemplate.pre contract.

post receives the function result after the function has completed. Although the base hook returns its argument and its extension-point documentation permits inspection or transformation, execute does not assign the result of await self.post(v) back to v. The value returned by execute is the original v. Use post for post-function observation or side effects unless a task subclass also changes execute to consume a transformed value. If the function raises, control does not reach the post call, so post is not an exception fallback.

A synchronous function wrapped in AsyncFunctionTaskTemplate follows the same execute ordering; only the function-call branch changes from await self.func(...) to self.func(...). The class sets _call_as_synchronous for that case during __post_init__, after calling TaskTemplate.__post_init__ to normalize metadata such as `