Configuring Container Arguments for Remote Execution

When a task is scheduled for execution on a remote Flyte cluster, the control plane needs an exact specification of the container runtime: the entrypoint binary, input/output paths, code distribution bundles, image identifiers, and task loader directives. In flyte-sdk, task templates translate high-level Python functions or container commands into these CLI tokens and protobuf definitions during serialization.

Serialization Architecture and SerializationContext

During compilation and deployment, flyte-sdk converts task instances into protobuf tasks_pb2.Container specifications using SerializationContext (defined in flyte.models). Rather than hardcoding absolute remote storage locations or run identifiers at compilation time, flyte-sdk inserts Flyte Propeller template variables (such as {{.input}} and {{.outputPrefix}}). When Kubernetes pods or execution engines launch on the Flyte backend, Propeller substitutes these template tokens with concrete execution URIs.

from flyte.models import CodeBundle, ImageCache, SerializationContext

# SerializationContext supplies runtime placeholders and bundle metadata
sctx = SerializationContext(
    version="v1.0.0",
    project="my-project",
    domain="development",
    input_path="{{.input}}",
    output_path="{{.outputPrefix}}",
    code_bundle=CodeBundle(tgz="s3://my-bucket/code.tar.gz", destination="."),
)

Inside flyte._internal.runtime.task_serde, the serialization pipeline builds the protobuf container definition:

# flyte-sdk sets an empty command list and relies on container args for execution directives
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),
)

Standard Python Task Argument Generation

For tasks defined with @task or by subclassing AsyncFunctionTaskTemplate (in flyte.extend), remote execution is managed through the runtime CLI entrypoint a0. The task template's container_args(serialize_context) method generates the full CLI invocation passed to the primary container.

from flyte.extend import AsyncFunctionTaskTemplate
from flyte.models import SerializationContext

# AsyncFunctionTaskTemplate constructs a structured CLI invocation for the `a0` runtime entrypoint
template: AsyncFunctionTaskTemplate = ...
args = template.container_args(sctx)

Generated CLI Structure

AsyncFunctionTaskTemplate.container_args generates the command tokens as follows:

def container_args(self, serialize_context: SerializationContext) -> List[str]:
    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}}",
    ]
    # Image cache resolution
    if serialize_context.image_cache and serialize_context.image_cache.serialized_form:
        args = [*args, "--image-cache", serialize_context.image_cache.serialized_form]
    elif serialize_context.image_cache:
        args = [*args, "--image-cache", serialize_context.image_cache.to_transport]

    # Code distribution bundles
    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 '.'}"]]

    # Task resolution loader
    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),
            ],
        ]

    assert all(isinstance(item, str) for item in args), f"All args should be strings, non string item = {args}"
    return args

Argument Categories and Flags

Argument / Flag	Value Source	Purpose
a0	Static binary name	Runtime CLI entrypoint in the Flyte container.
--inputs	serialize_context.input_path	Input literal protobuf location (default: {{.input}}).
--outputs-path	serialize_context.output_path	Output directory destination (default: {{.outputPrefix}}).
--version	serialize_context.version	Task serialization version string.
--raw-data-path	{{.rawOutputDataPrefix}}	Destination prefix for offloaded blobs and files.
--checkpoint-path	{{.checkpointOutputPrefix}}	Destination path for task state checkpoints.
--prev-checkpoint	{{.prevCheckpointPrefix}}	Source path of previous checkpoint for state recovery.
--run-name	{{.runName}}	Current workflow execution run name.
--name	{{.actionName}}	Action/node identifier in the workflow graph.
--image-cache	serialize_context.image_cache	Serialized image cache mapping for fast container image lookup.
--tgz / --pkl	serialize_context.code_bundle	Path to downloaded archive bundle or pickled function payload.
--dest	code_bundle.destination	Code unpacking directory destination (defaults to .).
--resolver	DefaultTaskResolver.import_path	Resolver module path (flyte._internal.resolvers.default.DefaultTaskResolver).
mod <M> instance <T>	DefaultTaskResolver.loader_args	Python module and task object attribute names to import dynamically.

When a task uses a pickled code bundle (--pkl), flyte-sdk omits the --resolver arguments because the unpickling process re-inflates the task instance directly without importing from source modules.

Raw Container Tasks and CoPilot Data Loading

When running tasks in pre-built images without the Flyte Python SDK (such as C++, Rust, or Bash executables), flyte-sdk provides ContainerTask in flyte.extras (implemented in flyte.extras._container).

ContainerTask sets task_type="raw-container". Rather than invoking the a0 entrypoint, ContainerTask.container_args(sctx) returns the user-provided command and argument tokens unmodified:

from flyte.extras import ContainerTask

task = ContainerTask(
    name="process_data",
    image="ubuntu:22.04",
    command=["python3", "main.py"],
    arguments=["--threshold", "{{.inputs.threshold}}", "/var/inputs/data_file"],
    inputs={"threshold": float, "data_file": File},
    outputs={"result": int},
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    metadata_format="JSON",
)

CoPilot Data Loading Configuration

Flyte CoPilot handles side-loading input data to the container filesystem and reading outputs when the process exits. ContainerTask.data_loading_config produces the required tasks_pb2.DataLoadingConfig:

def data_loading_config(self, sctx: SerializationContext) -> tasks_pb2.DataLoadingConfig:
    literal_to_protobuf = {
        "JSON": tasks_pb2.DataLoadingConfig.JSON,
        "YAML": tasks_pb2.DataLoadingConfig.YAML,
        "PROTO": tasks_pb2.DataLoadingConfig.PROTO,
    }

    return tasks_pb2.DataLoadingConfig(
        input_path=str(self._input_data_dir) if self._input_data_dir else None,
        output_path=str(self._output_data_dir) if self._output_data_dir else None,
        enabled=True,
        format=literal_to_protobuf.get(self._metadata_format, "JSON"),
    )

Primitive Templates vs. File/Directory Path Syntax

In ContainerTask, input values are referenced differently depending on whether they are primitive values or file/directory references:

1. Primitive Inputs: Reference primitive types (e.g., int, str, float) using template notation {{.inputs.key}}.
2. File and Directory Inputs (File, Dir): Reference file or directory inputs using path-like syntax rooted at input_data_dir (e.g., /var/inputs/data_file).

Using template syntax {{.inputs.infile}} for File or Dir inputs raises an AssertionError during local volume preparation in _render_command_and_volume_binding:

if input_val and type(input_val) in [File, Dir]:
    if not path_k:
        raise AssertionError(
            "File and Directory commands should not use the template syntax "
            "like this: {{.inputs.infile}}\n"
            "Please use a path-like syntax, such as: /var/inputs/infile.\n"
            "This requirement is due to how Flyte Propeller processes template syntax inputs."
        )

Path-like arguments ensure Flyte CoPilot and local Docker executors can bind volume mounts directly to /var/inputs/<input_name>.

Extending and Customizing Container Execution

Custom execution environments can override how container commands, resource specifications, and runtime plugins are built by subclassing task templates and using the extension primitives exported in flyte.extend.

Creating a Custom Task Template

Subclass AsyncFunctionTaskTemplate to intercept argument generation, attach additional flags, or customize execution behavior:

from typing import Any, List
from flyte.extend import AsyncFunctionTaskTemplate
from flyte.models import SerializationContext


class CustomEngineTaskTemplate(AsyncFunctionTaskTemplate):
    def container_args(self, serialize_context: SerializationContext) -> List[str]:
        # Generate baseline standard CLI arguments
        base_args = super().container_args(serialize_context)

        # Append custom runtime engine configuration
        return base_args + ["--custom-flag", "custom-value"]

Registering Plugins via TaskPluginRegistry

When tasks accept custom plugin configurations (plugin_config), TaskEnvironment resolves the matching template class using TaskPluginRegistry (exported in flyte.extend):

from dataclasses import dataclass
from flyte.extend import TaskPluginRegistry


@dataclass
class CustomEngineConfig:
    concurrency_limit: int = 4


# Register the template class for the custom config type
TaskPluginRegistry.register(CustomEngineConfig, CustomEngineTaskTemplate)

# TaskEnvironment automatically selects CustomEngineTaskTemplate
# whenever plugin_config=CustomEngineConfig(...) is used

Pod Spec Helpers

When building pod specifications or custom backend templates, flyte.extend provides pod_spec_from_resources and PRIMARY_CONTAINER_DEFAULT_NAME ("primary") to construct Kubernetes Pod templates that align with Flyte backend expectations:

from flyte._pod import PodSpec
from flyte._resources import Resources
from flyte.extend import PRIMARY_CONTAINER_DEFAULT_NAME, pod_spec_from_resources

resources = Resources(cpu="2", mem="4Gi")
pod_spec: PodSpec = pod_spec_from_resources(
    resources=resources,
    primary_container_name=PRIMARY_CONTAINER_DEFAULT_NAME,
)

This ensures container resource requests, limits, and runtime flags are structured uniformly across both standard Python tasks and specialized container extensions.