Skip to main content

Task Containerization Explained

A decorated function becomes a task specification whose container arguments tell the Flyte runtime where to read inputs, where to write outputs, how to identify the action, and how to load the task code. In flyte-sdk, AsyncFunctionTaskTemplate is the task-template type used for ordinary functions; despite its name, it accepts both asynchronous and synchronous callables and packages either one for remote execution.

From a decorated function to a task template​

TaskEnvironment.task is the normal construction path. Its decorator computes the task name from the environment and function names, derives a short name, creates a typed interface with NativeInterface.from_callable, and stores the resulting template in TaskEnvironment._tasks:

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

@env.task
async def my_task():
    pass

For an environment without a plugin configuration, the decorator instantiates AsyncFunctionTaskTemplate with the wrapped func, environment image and resources, cache and retry settings, timeout, reusable-environment settings, environment variables, secrets, pod template, parent environment, interface, reporting setting, short name, plugin configuration, and inline-I/O limit. The generated task name is self.name + "." + func.__name__, while short_name defaults to the function name. The parent environment is retained through a weak reference, allowing wire serialization to include the environment name.

The template is a keyword-only dataclass with two fields specific to this layer:

@dataclass(kw_only=True)
class AsyncFunctionTaskTemplate(TaskTemplate[P, R]):
    func: FunctionTypes
    plugin_config: Optional[Any] = None

TaskTemplate.__post_init__ runs first. The inherited initialization normalizes task configuration such as the image, cache, retries, and short name. AsyncFunctionTaskTemplate.__post_init__ then marks a non-coroutine callable with _call_as_synchronous = True. Thus the class name does not mean that only async def functions are supported: the same template represents synchronous functions, but the controller uses the synchronous submission path for them.

If an environment supplies plugin_config, TaskEnvironment.task performs an exact-type lookup in TaskPluginRegistry and constructs the registered subclass instead. Decoration raises a ValueError when no class is registered for that configuration type. Plugin subclasses are typed as extensions of AsyncFunctionTaskTemplate, so plugin-specific configuration travels with the same task-template and serialization flow.

Local forwarding versus task execution​

Packaging for a container is separate from ordinary local invocation. The template's forward method directly calls the wrapped 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)

For an asynchronous function this returns a coroutine for the caller to await; for a synchronous function it returns the value directly. The remote/runtime method is execute. It requires an existing task context, runs the asynchronous pre hook, temporarily replaces the task context with the hook's data, invokes the wrapped function with the appropriate async or sync call, and awaits post afterward:

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

execute returns the original v; the return value of post is not assigned back. It is therefore not the general-purpose local-call API, and a post hook cannot currently replace the task result. The source also contains a TODO about the eventual pre/post execution arrangement.

Serialization context and container construction​

Remote runs and deployments create a SerializationContext and translate the task with translate_task_to_wire. The context supplies the task version, input and output paths, optional code bundle, optional image cache, and resolver root. Code packaging can produce a tar-style bundle or a pickle bundle; the resulting artifact and its destination are represented in the context before container_args is called.

AsyncFunctionTaskTemplate.container_args is the central packaging boundary. It does not construct the image itself. Instead, it turns serialization metadata and runtime placeholders into the argument vector for the generated container, in a fixed order.

The initial arguments are the a0 entrypoint, input and output locations, version, and Flyte-provided execution metadata:

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

The paths and version come from SerializationContext. The raw-output, checkpoint, previous-checkpoint, run-name, and action-name values remain runtime placeholders for Flyte to resolve when the container runs. serialize_context.version is used for --version; the source has an inline comment questioning whether the code-bundle version would be preferable, so callers should not assume those values are interchangeable.

If an image cache is present, the method appends --image-cache. It prefers image_cache.serialized_form when available and otherwise uses image_cache.to_transport:

if serialize_context.image_cache and serialize_context.image_cache.serialized_form:
    args = [*args, "--image-cache", serialize_context.image_cache.serialized_form]
else:
    if serialize_context.image_cache:
        args = [*args, "--image-cache", serialize_context.image_cache.to_transport]

The task image itself is supplied separately when the protobuf container is built. For example, an image with Python dependencies can be declared on a task:

@flyte.task(image=(flyte.Image.from_debian_base().with_pip_packages("requests", "numpy")))
def my_task(x: int) -> int:
    import numpy as np
    return np.sum([x, 1])

The image builder and serialization flow resolve that image before the container specification is emitted; container_args adds image-cache transport only when the serialization context contains it.

Code bundles and resolver arguments​

When SerializationContext.code_bundle exists, container_args selects the available artifact. A tar bundle is passed as --tgz; a pickle bundle is passed as --pkl. Both forms add --dest, using the bundle destination or . when no destination is set:

if serialize_context.code_bundle:
    if serialize_context.code_bundle.tgz:
        args = [*args, "--tgz", f"{serialize_context.code_bundle.tgz}"]
    elif serialize_context.code_bundle.pkl:
        args = [*args, "--pkl", f"{serialize_context.code_bundle.pkl}"]
    args = [*args, "--dest", f"{serialize_context.code_bundle.destination or '.'}"]

A pickle bundle is the exceptional case: it contains the task object, so no resolver arguments are appended. A tar bundle, and a run with no code bundle, both append the default resolver import path and loader arguments. The loader receives this task and serialize_context.root_dir, allowing the default resolver to derive how to import the wrapped function.

The complete decision is:

if not serialize_context.code_bundle or not serialize_context.code_bundle.pkl:
    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),
    ]

This means --tgz does not eliminate the resolver: only --pkl does. The method finishes with an assertion that every argument is a string. Consequently, custom serialization-context values for paths, versions, image-cache transport, and bundle locations must already be string-valued.

How arguments become a Flyte container​

_internal.runtime.task_serde._get_urun_container places the generated argument list beside the image and the other task configuration:

return tasks_pb2.Container(
    image=img_uri,
    command=[],
    args=task_template.container_args(serialize_context),
    resources=resources,
    env=env,
    data_config=task_template.data_loading_config(serialize_context),
    config=task_template.config(serialize_context),
)

Before constructing this protobuf, the function converts task resources with get_proto_resources, maps env_vars to protobuf key/value pairs, obtains the normalized image URI, and rejects a remaining string image with RuntimeSystemError("BadConfig", "Image is not a valid image"). The resulting container therefore carries more than the command-line packaging arguments: it includes the task image, resource requirements, environment variables, data-loading configuration, and task or plugin configuration. get_proto_task places that container in the serialized task template alongside interface, retries, timeout, cache metadata, secrets, pod information, and extended resources.

For example, resource settings are task metadata inherited into this serialization path:

@task(resources=Resources(cpu=1, memory="1GiB", gpu="T4:1"))
def my_task() -> int:
    return 42

The interface created at decoration time is also serialized, and remote call arguments are converted through that native interface before execution. Container arguments describe how the runtime starts and loads the task; they do not replace the typed input/output interface.

Runtime loading inside the container​

The runtime counterpart is _internal.runtime.entrypoints._download_and_load_task. If a bundle has a tar or pickle location, the runtime downloads it first. A pickle bundle is loaded directly with load_pkl_task. For a tar bundle, the runtime requires the resolver and resolver arguments, then calls load_task. With no bundle, it also requires the resolver pair:

if code_bundle and (code_bundle.tgz or code_bundle.pkl):
    code_bundle = await download_code_bundle(code_bundle)
    if code_bundle.pkl:
        return load_pkl_task(code_bundle)

    if not resolver or not resolver_args:
        raise flyte.errors.RuntimeSystemError(
            "MalformedCommand", "Resolver and resolver args are required. for task"
        )
    return load_task(resolver, *resolver_args)
if not resolver or not resolver_args:
    raise flyte.errors.RuntimeSystemError("MalformedCommand", "Resolver and resolver args are required. for task")
return load_task(resolver, *resolver_args)

The default resolver's task-module logic inspects task.func to identify the function's module and entity name. This is why source and import metadata matter for tar and no-bundle execution. source_file exposes the wrapped callable's __code__.co_filename when available and returns None when the callable has no code object. Functions without an importable module or file can therefore fail during resolver-based reconstruction, whereas pickle loading uses the serialized task object directly.

Constraints visible at deployment time​

Containerization is shared by remote execution and deployment. Remote execution builds or reuses the image and code bundle, creates the serialization context, translates the task to wire format, and converts call arguments through the native interface. Deployment creates one context for the environment and serializes each registered task, so the same container_args rules become part of the deployed task specification.

Several configuration boundaries are important when diagnosing a generated command:

  • A missing resolver or resolver arguments is a runtime RuntimeSystemError with code MalformedCommand for tar and no-bundle paths.
  • A pickle bundle does not receive resolver arguments; a tar bundle does.
  • The code-bundle destination defaults to . in the generated arguments.
  • Image-cache mismatches can fall back to computing the task image URI, while serialization logs a warning that differing Flyte SDK versions can produce inconsistent image resolution.
  • Synchronous functions are valid inputs, but __post_init__ marks them for synchronous controller submission and execute calls them without await.
  • Reusable environments cannot set a pod template; reusable environments with concurrency greater than one reject synchronous functions and require an async function.
  • Plugin configuration requires an exact TaskPluginRegistry registration before decoration can construct a plugin task template.

The net effect is a deliberately split package: image, resources, environment, data-loading settings, and plugin configuration occupy structured fields in the Flyte container protobuf, while input/output locations, runtime metadata, code-bundle selection, and resolver loading instructions are passed as command-line arguments generated by AsyncFunctionTaskTemplate.container_args.