Skip to main content

Managing Data I/O in Container Tasks

In flyte-sdk, the ContainerTask class provides a mechanism for executing arbitrary containerized workloads. It manages the mapping between Flyte task inputs/outputs and the container's filesystem, allowing you to run existing binaries or scripts within a Flyte workflow.

Defining a Container Task​

To create a container task, instantiate flyte.extras.ContainerTask. You must specify the container image, the command to run, and the expected inputs and outputs.

from flyte.extras import ContainerTask
from flyte.image import Image
import pathlib

# Define a task that runs a shell script inside a container
my_container_task = ContainerTask(
    name="bash-task",
    image=Image.from_base("ubuntu:latest"),
    command=["/bin/bash", "-c"],
    arguments=["echo 'hello world' > /var/outputs/res"],
    inputs={"val": int},
    outputs={"res": str},
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    metadata_format="JSON"
)

Configuring Input and Output Directories​

ContainerTask uses specific directories inside the container to exchange data with the Flyte engine. By default, these are:

  • input_data_dir: /var/inputs
  • output_data_dir: /var/outputs

When the task executes, Flyte mounts the input data into the input_data_dir. The task is expected to write its results into the output_data_dir.

Data Loading Formats​

The metadata_format parameter determines how Flyte serializes and deserializes the data. Supported formats include:

  • JSON (Default)
  • YAML
  • PROTO

This configuration is captured in the DataLoadingConfig during task serialization:

# Internal representation of data loading configuration
config = my_container_task.data_loading_config(serialization_context)
# Result: tasks_pb2.DataLoadingConfig(input_path="/var/inputs", output_path="/var/outputs", format=JSON)

Handling File and Directory Inputs​

When working with flyte.io.File or flyte.io.Dir types, flyte-sdk handles the volume mounting automatically. However, you must use a specific path-like syntax in your command to reference these inputs.

Correct Path Syntax​

Instead of using template syntax like {{.inputs.my_file}}, you must reference the file by its expected path within the input_data_dir.

from flyte.io import File

file_task = ContainerTask(
    name="file-processor",
    image=Image.from_base("alpine"),
    command=["cat"],
    # Use the path where the file will be mounted: input_data_dir / input_name
    arguments=["/var/inputs/data_file"],
    inputs={"data_file": File},
    input_data_dir="/var/inputs"
)

If you use template syntax for File or Dir types, flyte-sdk will raise an AssertionError during command rendering to prevent execution failures in the Flyte backend.

Local Execution with Docker​

For local testing, ContainerTask.execute uses the local Docker daemon to run the container.

  1. Volume Binding: flyte-sdk creates a temporary local directory and binds it to the container's output_data_dir.
  2. Image Management: It automatically pulls the required image if it is not present locally using client.images.pull(uri).
  3. Log Streaming: If local_logs=True (default), container logs are printed to the console with the prefix [Local Container].
# Local execution requires Docker to be installed
# result = await my_container_task.execute(val=10)

Troubleshooting and Constraints​

  • Docker Dependency: Local execution will fail with an ImportError if the docker Python package is not installed, or an error if the Docker daemon is not reachable.
  • Template Syntax: Template syntax {{.inputs.key}} is only supported for primitive types (strings, integers, etc.). For File and Dir types, you must use the explicit path /var/inputs/<key>.
  • Execution Timeout: The current implementation of execute waits indefinitely for the container to finish (container.wait()). Ensure your containerized process has its own internal timeouts if necessary.
  • Output Parsing: flyte-sdk parses outputs by reading files from the output_data_dir named after the output keys. For example, if outputs={"res": str}, the task must write to /var/outputs/res.