Skip to main content

Creating Custom Container Tasks

Create a raw-container task​

Use extras.ContainerTask when the work should be performed by an arbitrary Docker image and command rather than by a Python function:

from flyte.extras import ContainerTask

word_count = ContainerTask(
    name="word-count",
    image="python:3.11-slim",
    command=["sh", "-c"],
    arguments=["echo {{.inputs.value}} > /var/outputs/result"],
    inputs={"value": str},
    outputs={"result": str},
)

ContainerTask is publicly re-exported from flyte.extras; extras._container contains the implementation but is an internal module. The constructor requires name, image, and command. Pass command and argument tokens as lists: execute combines command and arguments, and container_args returns that combined list for serialization.

The constructor creates a TaskTemplate with task type raw-container and a NativeInterface built from the inputs and outputs dictionaries. String images are converted to Image objects with Image.from_base; the special string "auto" uses Image.from_debian_base. An Image object can also be supplied directly.

The default container paths are /var/inputs and /var/outputs. The container must write each declared output to a file below the configured output path. In the example, the task writes /var/outputs/result.

Configure primitive inputs with template substitution​

For primitive inputs, use the exact {{.inputs.<name>}} syntax in a command or argument token:

from flyte.extras import ContainerTask

format_value = ContainerTask(
    name="format-value",
    image="python:3.11-slim",
    command=["sh", "-c"],
    arguments=["printf '%s' '{{.inputs.value}}' > /var/outputs/formatted"],
    inputs={"value": int},
    outputs={"formatted": str},
)

The implementation finds input names matching letters, digits, and underscores, obtains the corresponding value from kwargs, converts it with str(), and replaces the template text. Substitution is direct string replacement; ContainerTask does not shell-quote or escape the value. The command is also printed during local execution, so avoid putting sensitive values into command strings.

_prepare_command_and_volumes applies this rendering independently to every token in the combined command-and-arguments list. It preserves the list token boundaries and merges any volume bindings produced while rendering.

Mount File and Dir inputs with path-like references​

File and Dir inputs use a path under the configured input directory, not a template expression:

from flyte.extras import ContainerTask
from flyte.io import File

copy_file = ContainerTask(
    name="copy-file",
    image="python:3.11-slim",
    command=["sh", "-c"],
    arguments=["cat /var/inputs/infile > /var/outputs/copy"],
    inputs={"infile": File},
    outputs={"copy": File},
)

When a command token contains /var/inputs/infile, the path parser extracts infile. For an input whose exact runtime type is File or Dir, ContainerTask binds the input object's .path on the host to /var/inputs/infile in the container, with Docker mode rw. The same behavior applies when the configured input_data_dir is different:

from flyte.extras import ContainerTask
from flyte.io import Dir

inspect_directory = ContainerTask(
    name="inspect-directory",
    image="python:3.11-slim",
    command=["sh", "-c"],
    arguments=["find /data/in/records -type f > /data/out/listing"],
    inputs={"records": Dir},
    outputs={"listing": str},
    input_data_dir="/data/in",
    output_data_dir="/data/out",
)

Do not write {{.inputs.infile}} for a File or Dir input. The implementation raises an AssertionError and asks for a path-like reference instead. The special branch checks type(input_val) in [File, Dir], so subclasses or other compatible wrappers do not use this bind-mount behavior.

Define the output-file contract​

Declare outputs as a name-to-type dictionary. At runtime, _get_output looks for one file per declared name at <output_data_dir>/<name>, reads an existing file as text, converts that text to the declared type, and returns the values as a tuple in dictionary iteration order:

from flyte.extras import ContainerTask

produce_values = ContainerTask(
    name="produce-values",
    image="python:3.11-slim",
    command=["sh", "-c"],
    arguments=[
        "printf '%s' '42' > /var/outputs/count; "
        "printf '%s' 'true' > /var/outputs/ready"
    ],
    outputs={"count": int, "ready": bool},
)

The output conversion rules implemented by ContainerTask are:

  • bool: every value other than case-insensitive exact "false" becomes True.
  • datetime.datetime: parsed with datetime.datetime.fromisoformat.
  • datetime.timedelta: parsed by _string_to_timedelta using the supported days, H:MM:SS[.microseconds]-style form.
  • File: reconstructed with File.from_local(output_path).
  • Dir: reconstructed with Dir.from_local(output_path).
  • Other types: constructed by calling the type with the text value, as in output_type(output_val).

For File and Dir outputs, the output path is passed to from_local; the container still needs to create the corresponding named output path. Missing output files are represented internally as None, which can fail when the declared type conversion is applied. Text is read without removing a trailing newline, so a command such as echo can produce a newline in a string or numeric conversion input.

Run locally with Docker​

ContainerTask.execute is the local Docker path. It imports the Python Docker SDK only when execution starts, so local execution requires the docker package and a Docker daemon address understood by docker.from_env().

The execution sequence is:

  1. Normalize the input and output directories.
  2. Allocate a random host-side output directory with storage.get_random_local_directory().
  3. Render the combined command and arguments and add a bind mount from that random directory to output_data_dir.
  4. Create a Docker client with docker.from_env().
  5. Resolve the image URI, pulling it if it is not already listed locally.
  6. Start the container detached with remove=True and the generated volume bindings.
  7. Stream logs when local_logs=True, then wait for the container.
  8. Read and convert the declared output files.

The command is printed regardless of local_logs; that flag only controls streamed log output prefixed with [Local Container]:

from flyte.extras import ContainerTask

quiet_task = ContainerTask(
    name="quiet-task",
    image="python:3.11-slim",
    command=["sh", "-c"],
    arguments=["printf '%s' done > /var/outputs/status"],
    outputs={"status": str},
    local_logs=False,
)

The image is normally an Image after construction. execute defensively raises an assertion if a raw string remains at that point, and uses the image object's .uri. If the Docker SDK is unavailable, it raises an ImportError with the installation command pip install docker. The image must already be available locally or be pullable from its registry.

Serialize the task for Flyte​

The runtime serializer handles ContainerTask through its TaskTemplate hooks. _internal.runtime.task_serde.get_proto_task calls _get_urun_container, which creates a FlyteIDL Container whose image is the resolved image URI, whose args come from container_args, and whose data_config comes from data_loading_config:

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),
)

ContainerTask serializes as task type raw-container. data_loading_config sets enabled=True, serializes the input and output paths, and maps metadata_format values "JSON", "YAML", and "PROTO" to the corresponding tasks_pb2.DataLoadingConfig enum values. The configured paths therefore participate in Flyte's serialized raw-container data-loading configuration as well as in local Docker volume setup.

If a concrete pod template is configured through inherited TaskTemplate options, get_proto_task passes the generated primary container to _get_k8s_pod. The pod's primary container must exist under primary_container_name; its image is filled from the task when the pod does not specify one, while the task's command, args, resources, and environment are applied to the primary container. Other pod containers remain part of the pod template.

Configuration and operational caveats​

  • arguments defaults to None and is appended after command; command itself is annotated as List[str], so provide token lists even though the docstring mentions a single string.
  • input_data_dir and output_data_dir accept strings or pathlib.Path values. String values are converted to Path objects by the constructor and normalized with os.path.normpath during local execution.
  • metadata_format defaults to "JSON"; the declared supported values are "JSON", "YAML", and "PROTO".
  • The inherited TaskTemplate keyword options are forwarded through **kwargs, including task-level settings such as resources, retries, cache, environment, secrets, timeout, and pod template. The inherited timeout is serialized at the Flyte task level, but it does not impose a timeout on the local container.wait() call; the source contains an explicit TODO for a container timeout.
  • container.wait() is called without checking its returned status. A failed container can therefore proceed to output parsing without a task-specific failure raised by execute.
  • Docker mounts both File/Dir input bindings and the output binding with mode rw. The container is removed after exit, but execute does not visibly remove the random host output directory.
  • The local Docker path is distinct from serialization: local execution needs a Docker daemon, while serialization produces FlyteIDL container and data-loading fields for Flyte runtime handling.
  • The repository snapshot provides no direct ContainerTask construction sites, tests, README examples, or examples directory; the constructor and implementation in extras/_container.py are the available usage contract.