Skip to main content

How to Define a Raw Container Task

When you need to execute non-Python binaries, pre-built vendor containers, shell scripts, or legacy CLI tools inside a workflow, standard @task function decorators require wrapping the execution in Python subprocess calls. The ContainerTask class in flyte-sdk solves this by defining a raw container task that executes a specific container image directly with custom commands and arguments.

Defining a Container Task

To create a raw container task, instantiate ContainerTask from flyte.extras (or flyte.extras._container). You must provide the task name, the container image, and the command to execute.

from flyte.extras import ContainerTask

simple_shell_task = ContainerTask(
name="simple-shell-task",
image="alpine:latest",
command=["echo", "Hello from raw container"],
)

You can pass images as URI strings, as "auto" (which resolves to Image.from_debian_base()), or as Image instances created via flyte.Image.from_base(...).

from flyte import Image
from flyte.extras import ContainerTask

custom_image = Image.from_base("ubuntu:22.04")

date_task = ContainerTask(
name="date-task",
image=custom_image,
command=["date"],
arguments=["-u"],
)

Passing Inputs to Container Tasks

ContainerTask accepts typed inputs via the inputs dictionary. However, the syntax used in command and arguments depends on the input data type.

Primitive Inputs

For primitive types such as int, float, str, and bool, reference the input name using the template interpolation syntax {{.inputs.<name>}}:

from flyte.extras import ContainerTask

calculate_task = ContainerTask(
name="calculate-task",
image="python:3.11-slim",
inputs={"x": int, "y": int, "prefix": str},
command=["python3", "-c"],
arguments=[
"import sys; "
"res = {{.inputs.x}} + {{.inputs.y}}; "
"print(f'{{.inputs.prefix}}: {res}')"
],
)

File and Directory Inputs

For flyte.io.File and flyte.io.Dir inputs, do not use template interpolation syntax like {{.inputs.infile}}. Flyte mounts file and directory inputs into the container's input_data_dir (which defaults to /var/inputs). You must reference them using the absolute container path syntax /var/inputs/<name>.

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

line_count_task = ContainerTask(
name="line-count-task",
image="alpine:latest",
inputs={"input_file": File},
command=["wc", "-l", "/var/inputs/input_file"],
)

If you use {{.inputs.<name>}} for a File or Dir input, ContainerTask raises an AssertionError during execution:

AssertionError: File and Directory commands should not use the template syntax like this: {{.inputs.infile}}
Please use a path-like syntax, such as: /var/inputs/infile.
This requirement is due to how Flyte Propeller processes template syntax inputs.

Capturing Outputs

ContainerTask captures outputs by reading files from the directory defined by output_data_dir (default: /var/outputs). To return data from the container:

  1. Define an outputs dictionary mapping output names to types (int, float, str, bool, datetime, timedelta, File, Dir).
  2. Have the container command write raw values into individual files named after each output key under /var/outputs/ (e.g., /var/outputs/result).
import datetime
from flyte.extras import ContainerTask
from flyte.io import File

process_and_output_task = ContainerTask(
name="process-and-output",
image="alpine:latest",
inputs={"seed": int},
outputs={
"computed_val": int,
"is_valid": bool,
"summary_file": File,
},
command=["/bin/sh", "-c"],
arguments=[
"mkdir -p /var/outputs && "
"echo $(( {{.inputs.seed}} * 42 )) > /var/outputs/computed_val && "
"echo true > /var/outputs/is_valid && "
"echo 'Processing completed successfully' > /var/outputs/summary_file"
],
)

Output Type Parsing Rules

When the container finishes execution, ContainerTask parses each output file according to its declared type:

  • bool: Parses as False if the content in the output file is case-insensitively "false". All other non-empty strings parse as True.
  • datetime.datetime: Converted from an ISO-8601 string via datetime.datetime.fromisoformat(...).
  • datetime.timedelta: Converted via a regex matching formatted durations (e.g., "1 days, 02:03:04.000000" or "00:15:30").
  • File and Dir: Instantiated asynchronously from local file system locations via File.from_local(output_path) and Dir.from_local(output_path).
  • Other types (such as int, float, str): Cast directly using the type constructor (output_type(output_val)).

_get_output returns outputs as a tuple, allowing Flyte to map the returned elements to the declared output interface.

Customizing Paths, Metadata, and Task Resources

You can customize input/output paths, serialization formats, and underlying TaskTemplate execution parameters:

import pathlib
from flyte.extras import ContainerTask
from flyte.models import Resources

custom_config_task = ContainerTask(
name="custom-config-task",
image="alpine:latest",
inputs={"threshold": float},
outputs={"status": str},
input_data_dir=pathlib.Path("/custom/inputs"),
output_data_dir=pathlib.Path("/custom/outputs"),
metadata_format="JSON", # Options: "JSON", "YAML", "PROTO"
local_logs=True,
command=["/bin/sh", "-c"],
arguments=[
"mkdir -p /custom/outputs && "
"echo 'Threshold: {{.inputs.threshold}}' > /custom/outputs/status"
],
# TaskTemplate kwargs
resources=Resources(cpu="2", memory="4Gi"),
retries=2,
)

Configuration Parameters

ParameterTypeDefaultDescription
namestrRequiredName identifier for the task.
imageUnion[str, Image]RequiredImage URI, "auto", or flyte.Image instance.
commandList[str]RequiredCommand to run in the container.
argumentsOptional[List[str]]NoneArguments appended to the command.
inputsOptional[Dict[str, Type]]NoneMapping of input names to Python types.
outputsOptional[Dict[str, Type]]NoneMapping of output names to Python types.
input_data_dirstr | pathlib.Path"/var/inputs"Target mount directory inside the container for inputs.
output_data_dirstr | pathlib.Path"/var/outputs"Directory inside the container where outputs are written.
metadata_formatLiteral["JSON", "YAML", "PROTO"]"JSON"Serialization format configured on tasks_pb2.DataLoadingConfig.
local_logsboolTrueStream Docker container stdout/stderr to the console during local execution.
**kwargsAnyAdditional TaskTemplate parameters such as resources, cache, retries, env_vars, secrets, timeout, and pod_template.

Execution Mechanisms

Remote Execution on Flyte Cluster

When compiled and serialized for a remote Flyte cluster, ContainerTask sets task_type="raw-container". ContainerTask.data_loading_config(...) produces a tasks_pb2.DataLoadingConfig specifying input_path, output_path, and the protobuf metadata format. FlytePropeller handles volume mounts, data loading, and container orchestration directly on Kubernetes.

Local Execution via Docker

When executed locally, ContainerTask.execute(**kwargs) orchestrates a Docker container on the host machine:

  1. Requires the Python docker package (pip install docker) and an active Docker daemon.
  2. Allocates a temporary local directory via storage.get_random_local_directory() and binds it to output_data_dir in read-write mode ("rw").
  3. Resolves File and Dir inputs by creating Docker volume mounts from the local storage paths to /var/inputs/<key>.
  4. Renders primitive input template variables ({{.inputs.<name>}}) in command strings.
  5. Pulls the target image via client.images.pull(...) if not already present locally.
  6. Runs the container, streams logs if local_logs=True, waits for completion, and extracts output values from the bound output directory.

Troubleshooting

Missing Docker Package during Local Run

  • Symptom: Running the task locally raises ImportError: Docker is not installed. Please install Docker by running 'pip install docker'.
  • Fix: Install the Docker SDK for Python in your local environment with pip install docker and verify that the Docker daemon is running.

Template Syntax Used for File or Dir

  • Symptom: Task fails with AssertionError: File and Directory commands should not use the template syntax like this: {{.inputs.infile}}.
  • Fix: Replace {{.inputs.<var>}} with /var/inputs/<var> (or <input_data_dir>/<var>) in command or arguments.

Missing Output Files

  • Symptom: Missing outputs result in output_val = None, causing errors when converted (e.g., TypeError: int() argument must be a string, a bytes-like object or a real number, not 'NoneType').
  • Fix: Ensure the container command creates the output directory (mkdir -p /var/outputs) and writes non-empty content to a file exactly matching every key defined in outputs.