Understanding Task Templates

The common task-template contract

When a Flyte task is called, flyte-sdk has to support two different operations: running code locally during development and describing work for a Flyte controller to execute. TaskTemplate is the shared contract for both operations. It stores the task name, interface, task_type, image, resources, cache and retry settings, environment variables, secrets, timeout, pod-template information, and report settings. Its __post_init__ normalizes an "auto" or string image into an Image, converts string cache requests into Cache objects, converts an integer retry count into a RetryStrategy, and defaults an empty short_name to the task name (flyte/_task.py).

The base class separates execution hooks from serialization hooks. pre returns task-context data before execution, execute is the execution method a subclass implements, and post runs after the result. config, custom_config, data_loading_config, and container_args provide the corresponding task-specific serialization points. The base execute raises NotImplementedError, while the default container_args, config, and custom_config return empty values.

The call path is context-sensitive. TaskTemplate.__call__ checks internal_ctx(); outside a Flyte task context it calls forward, while inside a task context it obtains a controller and submits the task. A task marked _call_as_synchronous is submitted with controller.submit_sync and its result is returned synchronously; other tasks use controller.submit. If no controller is available in a task context, the call raises a RuntimeSystemError. TaskTemplate.aio provides the analogous awaitable path and wraps a synchronous controller future with asyncio.wrap_future.

override returns a dataclass replacement rather than mutating the existing template. It permits changes such as resources, cache, retries, timeout, environment variables, secrets, and pod template, but explicitly rejects attempts to override name, image, docs, or interface. Reusable tasks also impose restrictions: overriding resources, environment variables, or secrets while reusability remains enabled raises ValueError.

Python-function tasks

Construction through TaskEnvironment.task

The normal function path begins with TaskEnvironment.task in flyte/_task_environment.py. The decorator derives the task name as the environment name followed by the function name, and uses the function name as the default short name. Without plugin configuration, it selects AsyncFunctionTaskTemplate; it builds the interface with NativeInterface.from_callable(func) and passes the environment's image, resources, cache, reusable policy, documentation, environment variables, secrets, pod template, and other settings to the template.

The following is the construction logic used by flyte-sdk, including the default and plugin-selected template choice:

if self.plugin_config is not None:
    from flyte.extend import TaskPluginRegistry

    task_template_class = TaskPluginRegistry.find(
        config_type=type(self.plugin_config)
    )
    if task_template_class is None:
        raise ValueError(
            f"No task plugin found for config type {type(self.plugin_config)}. "
            f"Please register a plugin using flyte.extend.TaskPluginRegistry.register() api."
        )
else:
    task_template_class = AsyncFunctionTaskTemplate

tmpl = task_template_class(
    func=func,
    name=task_name,
    image=self.image,
    resources=self.resources,
    cache=cache or self.cache,
    retries=retries,
    timeout=timeout,
    reusable=self.reusable,
    docs=docs,
    env_vars=self.env_vars,
    secrets=self.secrets,
    pod_template=pod_template or self.pod_template,
    parent_env=weakref.ref(self),
    interface=NativeInterface.from_callable(func),
    report=report,
    short_name=short,
    plugin_config=self.plugin_config,
    max_inline_io_bytes=max_inline_io_bytes,
)
self._tasks[task_name] = tmpl

AsyncFunctionTaskTemplate is publicly available from flyte.extend, although its implementation is in flyte/_task.py. Its func field retains the original callable, and its source_file property returns func.__code__.co_filename when the callable has code available. This gives function tasks source-location information without requiring a separate source-file setting.

Despite its name, the template can wrap either an async or synchronous function. During __post_init__, it calls TaskTemplate.__post_init__ and sets _call_as_synchronous when iscoroutinefunction(self.func) is false. That flag changes controller dispatch: synchronous wrapped functions use submit_sync, whereas coroutine functions use the asynchronous submission path. TaskEnvironment.task also rejects a synchronous function in a reusable environment whose concurrency is greater than one.

Local calls versus in-task execution

For a local call, TaskTemplate.__call__ reaches AsyncFunctionTaskTemplate.forward. forward simply calls the original function and does not await it:

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

Consequently, a local call to a wrapped async function returns its coroutine to the caller; the async caller is responsible for awaiting it. A synchronous function returns its ordinary value. Calling execute directly is different and is not the normal local entry point. AsyncFunctionTaskTemplate.execute asserts that a task context exists, runs pre, temporarily installs the returned data in the task context, invokes the original function, awaits it only when it is a coroutine function, then runs post:

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 gives the function template two visibly different execution contracts: forward is the transparent local call, while execute is the context-aware implementation used when a task execution controller invokes the task.

Serialized Python runner arguments

At serialization time, _internal/runtime/task_serde.py treats task templates polymorphically. It creates a protobuf tasks_pb2.TaskTemplate, places the task's native interface into its typed interface, and obtains the container's image, arguments, resources, environment, data-loading configuration, and task configuration through template methods. For cache version calculation, function tasks use VersionParameters(func=task.func, image=task.image); other templates use func=None.

AsyncFunctionTaskTemplate.container_args describes the Python task runner rather than the user's function's shell command. Its initial arguments include the runner name, serialized input and output paths, version, raw-output and checkpoint placeholders, and run/action names:

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}}",
]

It appends image-cache data when present. A code bundle contributes either --tgz or --pkl, followed by --dest. If there is no code bundle, or the bundle is not a pickle bundle, the method appends the default resolver import path and resolver loader arguments. It finally asserts that every argument is a string. The resolver's function-module extraction path specifically expects an AsyncFunctionTaskTemplate and uses its wrapped function; non-function templates are not supported by that extraction route.

Extending and attaching templates

The public flyte.extend.TaskPluginRegistry is a type-keyed registry. TaskEnvironment.task looks up a plugin using the exact type(self.plugin_config), not a string name or a base-class search. Registration stores the configuration type and template class directly:

class _Registry:
    def register(self, config_type: Type, plugin: Type[T]):
        self._plugins[config_type] = plugin

    def find(self, config_type: Type) -> typing.Optional[Type[T]]:
        return self._plugins.get(config_type)

A plugin therefore needs to be registered before decorating a function whose environment has that configuration, and the selected class must be compatible with the function-task template contract. If lookup returns no class, decoration raises the explicit ValueError shown above.

Templates that are not created by the decorator can be added with TaskEnvironment.add_task. The method requires a unique task name, stores the object in the environment's task map, assigns a weak parent-environment reference, and returns the same template. This is the integration point for manually constructed templates such as ContainerTask.

Raw-container tasks

flyte.extras.ContainerTask represents the other execution model: the task's executable work is a raw container command rather than a Python callable. Its constructor requires name, image, and command; command is typed as List[str]. It accepts optional input and output type dictionaries, command arguments, input and output directories, a metadata format, and local-log behavior.

The constructor sets task_type="raw-container" and creates a NativeInterface from the input and output mappings. Input entries become (type, None) pairs, while outputs are passed through as declared. String images are converted with Image.from_base; the special string "auto" selects Image.from_debian_base. String data directories become pathlib.Path objects. The defaults are /var/inputs, /var/outputs, metadata format "JSON", and local_logs=True.

The following is an illustrative construction using the actual public constructor signature. The repository has no checked-in ContainerTask(...) call site; attach the resulting template to a TaskEnvironment with add_task when integrating it into an environment.

from flyte.extras import ContainerTask

container_task = ContainerTask(
    name="raw_container_task",
    image="auto",
    command=["/bin/sh", "-c"],
    arguments=["echo {{.inputs.message}} > /var/outputs/result"],
    inputs={"message": str},
    outputs={"result": str},
)

Command rendering and file inputs

ContainerTask renders both the command and arguments with _prepare_command_and_volumes. Scalar inputs use the template form {{.inputs.key}}; the matching value is converted to a string and substituted into the command. File and directory inputs use a separate path-based form. A command containing a path beginning with the configured input directory, such as /var/inputs/infile, yields a bind mount from the local File.path or Dir.path to /var/inputs/infile in the container.

The implementation explicitly rejects a File or Dir passed through template syntax such as {{.inputs.infile}}. It raises an assertion and asks for path-like syntax instead. The exact type check is type(input_val) in [File, Dir], so subclasses do not enter this bind-mount branch. Each command and argument is processed independently, and the resulting volume dictionaries are merged.

Local Docker execution and outputs

ContainerTask.execute is a local Docker path, not a direct Python-function call. It first imports the optional docker package and raises an installation error if the import fails. It normalizes the configured input and output directories, creates a random local output directory with storage.get_random_local_directory(), prepares the command and input bindings, and adds a binding from that temporary directory to the container's output directory.

It then obtains a client with docker.from_env(), requires the normalized image to be an Image object, pulls the image when it is not listed locally, and starts a detached removable container with the rendered command and volume bindings. When local_logs is true, it streams the container logs and prints each line with the [Local Container] prefix. It waits for the container, then reads declared output files from the temporary directory.

Output conversion follows the declared output type. Booleans are true for every string except case-insensitive "false"; datetimes use datetime.datetime.fromisoformat; timedeltas use the class's day/time regular expression; File and Dir use from_local; all other types are called with the file contents. Missing files are represented as None and still passed to the converter, so a missing output can fail conversion. Outputs are returned as a tuple, preserving the mapping's order as separate Flyte outputs rather than one list output.

There is no special daemon error handling after docker.from_env(): a missing or unreachable Docker daemon surfaces from the Docker client. The method also contains a TODO for a container wait timeout and does not pass the task's general timeout field to container.wait().

CoPilot and wire serialization

For raw-container execution, ContainerTask.container_args returns the raw command followed by its arguments:

def container_args(self, sctx: SerializationContext) -> List[str]:
    return self._cmd + (self._args if self._args else [])

Its data_loading_config enables Flyte CoPilot data loading, serializes the input and output paths, and maps "JSON", "YAML", and "PROTO" to the corresponding tasks_pb2.DataLoadingConfig values. An unrecognized value falls back to "JSON" in the protobuf constructor rather than being explicitly validated at runtime. This differs from the Python-function template, whose container_args point to the Python runner and its code or resolver loading mechanism.

Kubernetes and runtime integration

Both task forms ultimately flow through _internal.runtime.task_serde. The serializer calls the same template methods for image, command arguments, resources, environment, data_loading_config, and task configuration, then writes a tasks_pb2.TaskTemplate with the task type, typed interface, retry and timeout metadata, cache metadata, security context, and extended resources. The distinction is supplied by the template: Python tasks have the default Python task type and runner arguments, while ContainerTask declares raw-container and returns its user command and arguments.

A task with a Kubernetes PodTemplate takes a different serialization branch. _get_k8s_pod requires a container whose name equals pod_template.primary_container_name; otherwise it raises ValueError. For that primary container, serialization replaces the command and args with the serialized task container values, uses the task image only when the pod container has no image, copies non-empty resource requirements, and prepends task environment variables to the pod container's existing environment. Other pod containers remain in the pod specification.

The flyte.extend.pod_spec_from_resources helper creates a Kubernetes V1PodSpec containing one primary V1Container. It maps Resources.cpu, memory, gpu, and ephemeral_storage to Kubernetes resource names, using nvidia.com/gpu for GPUs by default. It rejects non-singular resource values, and when only requests or only limits are supplied it uses the supplied map for both. It does not map the Resources.shm field. The helper imports Kubernetes client types when called, so the Kubernetes client is an optional dependency for this path.

The split between AsyncFunctionTaskTemplate and ContainerTask is therefore more than two constructors for the same task: the former retains a callable and participates in function resolution, async/synchronous dispatch, lifecycle hooks, and Python code-bundle loading; the latter retains a container command, performs Docker-oriented local execution, binds file inputs, materializes output files, and advertises CoPilot data-loading paths. TaskTemplate and the runtime serializer keep those execution models interoperable at the task-template boundary without making raw containers pretend to be Python functions.