Skip to main content

Understanding Task Templates

Task templates at a glance​

When you decorate a Python callable with @env.task, you get a callable task object rather than the original function alone. For ordinary tasks, TaskEnvironment.task creates an AsyncFunctionTaskTemplate and stores it in TaskEnvironment._tasks. The public extension-facing name is extend.AsyncFunctionTaskTemplate; its concrete implementation is in _task.py, where it subclasses the generic TaskTemplate[P, R].

A minimal definition is:

import flyte
from flyte import Resources

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

@env.task
async def my_task():
    pass

The decorator derives the task name as environment.name + "." + func.__name__, builds a NativeInterface from the callable, and copies environment metadata such as the image, resources, cache policy, retries, timeout, environment variables, secrets, pod template, reporting setting, and inline-I/O limit into the template. The task decorator also accepts task-level short_name, cache, retries, timeout, docs, pod_template, report, and max_inline_io_bytes values.

TaskTemplate holds the common task identity and execution metadata. Its __post_init__ converts the special image value "auto" with Image.from_debian_base(), converts other string images with Image.from_base, turns cache strings into Cache objects, converts integer retries to RetryStrategy(count=...), and uses the task name when short_name is empty. AsyncFunctionTaskTemplate then checks inspect.iscoroutinefunction(self.func). A synchronous callable is accepted too; initialization marks it with _call_as_synchronous = True.

One template, three execution paths​

The same template is used for a direct local call, a run-managed local call, and a serialized task submitted for execution elsewhere. The distinction is made by TaskTemplate.__call__ and the caller's task context.

Direct local calls use forward​

Outside a Flyte task context, TaskTemplate.__call__ calls forward. AsyncFunctionTaskTemplate.forward simply invokes self.func:

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

It deliberately does not await the result. An asynchronous function therefore returns its coroutine to the caller, while a synchronous function returns its value. An async caller can await the async task's result normally; the task template does not add a second event loop or await operation in forward.

For a run-managed invocation, the documented entry point is flyte.run:

import flyte

env = flyte.TaskEnvironment("example")

@env.task
async def example_task(x: int, y: str) -> str:
    return f"{x} {y}"

if __name__ == "__main__":
    flyte.run(example_task, 1, y="hello")

Inside a task context, TaskTemplate.__call__ obtains a controller. Async templates use controller.submit(self, *args, **kwargs). Templates marked _call_as_synchronous use controller.submit_sync; __call__ waits for that future's result, while aio wraps the future so synchronous tasks remain awaitable in an async parent. The built-in migration pattern is:

@env.task
def my_legacy_task(x: int) -> int:
    return x

@env.task
async def my_new_parent_task(n: int) -> List[int]:
    collect = []
    for x in range(n):
        collect.append(my_legacy_task.aio(x))
    return asyncio.gather(*collect)

TaskTemplate.aio is therefore different from forward: it is the async-compatible entry point for both synchronous and asynchronous task templates. Calling a task from a task context without an initialized controller raises RuntimeSystemError("BadContext", "Controller is not initialized.").

Run-managed local execution reaches execute​

Runner._run_local creates a LocalController, installs a TaskContext, and dispatches synchronous templates through submit_sync and asynchronous templates through submit. LocalController.submit converts native arguments using _task.native_interface, calls direct_dispatch, and converts the captured outputs back to native values.

The worker-side implementation is AsyncFunctionTaskTemplate.execute. Its sequence is explicit:

async def execute(self, *args: P.args, **kwargs: P.kwargs) -> R:
    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 method cannot be treated as a direct-call replacement: it asserts that a task context already exists. It runs the inherited asynchronous pre hook, replaces the current task context with a copy containing the hook data, invokes the wrapped function using the appropriate sync/async form, and then calls post. The value returned by post is not assigned to v; the function result is returned unchanged. The base TaskTemplate.pre returns {}, and its base post returns its input, while subclasses can provide hook behavior.

Serialization supplies the container task​

During task serialization, _internal.runtime.task_serde.get_proto_task builds the Flyte task specification from the common template fields. It serializes the typed native interface, retry and timeout metadata, reporting flag, secrets, resources, pod information, custom configuration, and the container returned by _get_urun_container. That container obtains its command arguments from task.container_args(serialize_context).

For a task with caching enabled, the serializer uses a pickle bundle's computed version when one is present. Otherwise, when the task is an AsyncFunctionTaskTemplate, it computes the cache version from VersionParameters(func=task.func, image=task.image). Thus the wrapped callable and the resolved image participate in the cache-version input for the ordinary Python template.

How the standard container command is assembled​

AsyncFunctionTaskTemplate.container_args starts the runtime with a0 and passes the serialized input/output locations and task-run placeholders:

args = [
    "a0",
    "--inputs",
    serialize_context.input_path,
    "--outputs-path",
    serialize_context.output_path,
    "--version",
    serialize_context.version,
    "--raw-data-path",
    "{{.rawOutputDataPrefix}}",
    "--checkpoint-path",
    "{{.checkpointOutputPrefix}}",
    "--prev-checkpoint",
    "{{.prevCheckpointPrefix}}",
    "--run-name",
    "{{.runName}}",
    "--name",
    "{{.actionName}}",
]

If SerializationContext.image_cache is present, the method adds --image-cache using serialized_form when available, or to_transport otherwise. A code bundle adds --tgz for a tarball or --pkl for a pickle, followed by --dest using the bundle destination or ..

Source-based and pickle bundles take different paths. Whenever there is no code bundle, or the bundle does not contain a pickle, container_args appends the default resolver and loader arguments:

from flyte._internal.resolvers.default import DefaultTaskResolver

_task_resolver = DefaultTaskResolver()
args = [
    *args,
    "--resolver",
    _task_resolver.import_path,
    *_task_resolver.loader_args(task=self, root_dir=serialize_context.root_dir),
]

DefaultTaskResolver.import_path is flyte._internal.resolvers.default.DefaultTaskResolver. Its loader_args calls extract_task_module(task, root_dir) and emits ['mod', module, 'instance', function_name]. For an AsyncFunctionTaskTemplate, extract_task_module obtains the module from task.func, requires a module file, derives a module name relative to source_dir, and uses task.func.__name__ as the entity name. A task loaded from __main__ is converted to the stem of the main script. The method finally asserts that every generated container argument is a string.

This means a source-based task must be reloadable from its function module and the supplied source root. If the function has no module file, or its module file is not relative to the source directory, module extraction raises an error (or the relative-path operation fails). A pickle bundle avoids the resolver branch and uses --pkl instead.

Extension and customization points​

AsyncFunctionTaskTemplate carries plugin_config: Optional[Any] in addition to func. Normally TaskEnvironment.task selects the built-in template. If the environment has plugin_config, it calls TaskPluginRegistry.find(config_type=type(self.plugin_config)) and instantiates the registered template class instead. The registry stores a mapping from an exact configuration type to a template class:

TaskPluginRegistry.register(config_type, plugin)

Lookup is exact; a registration for a base class does not match a derived configuration type. If no matching registration exists, decoration raises a ValueError directing you to register a plugin with flyte.extend.TaskPluginRegistry.register(). Reusable environments cannot specify plugin_config.

The base TaskTemplate exposes extension methods that plugins can override: config, custom_config, data_loading_config, container_args, and sql. The base configuration methods return empty dictionaries, container_args returns an empty list, and sql returns None; AsyncFunctionTaskTemplate supplies the standard Python container arguments described above. data_loading_config is the hook used for raw-container/Flyte CoPilot data-loading configuration.

For per-invocation metadata changes, use TaskTemplate.override. It creates a dataclass replacement while preserving the original task object and interface. The method accepts values such as short_name, resources, cache, retries, timeout, reusable, env_vars, secrets, max_inline_io_bytes, and pod_template, but rejects attempts to override name, image, docs, or interface. With reusability enabled, resource, environment-variable, and secret overrides are rejected; pass reusable="off" when the source permits disabling reusability before changing those fields.

Operational boundaries to keep in mind​

  • Receive ordinary instances from @env.task; although the class is exported as extend.AsyncFunctionTaskTemplate, its implementation and runtime integrations are in the private _task module.
  • Do not call execute as a standalone local-function API. It requires an initialized task context; direct local calls use __call__ and forward instead.
  • Preserve the sync/async distinction when using reusable environments. TaskEnvironment.task rejects a synchronous function in a reusable environment whose concurrency is greater than one, and it rejects a pod template when the environment is reusable.
  • Keep post side effects separate from output transformation. execute awaits post(v) but returns v, so a post hook cannot replace the task result through its return value.
  • Treat source resolution as a deployment/runtime requirement. The default resolver imports the derived module and retrieves the function by name, so module-file and source-root assumptions apply to source-based bundles.
  • Remember that deployment diagnostics can use AsyncFunctionTaskTemplate.source_file: it returns func.__code__.co_filename when the wrapped callable exposes code and otherwise returns None.

Together, TaskTemplate and AsyncFunctionTaskTemplate form the boundary between a Python callable and a Flyte task specification: TaskEnvironment.task supplies the callable and metadata, __call__ selects direct or controller-backed execution, execute runs the worker lifecycle, and container_args gives the serialized container the a0 runtime plus either a code bundle or a resolver.