Skip to main content

Local Execution and Testing

Local Execution and Testing of Container Tasks​

This section guides you through understanding how flyte-sdk executes containerized tasks locally using Docker, focusing on the internal ContainerTask class. While ContainerTask is an internal representation and not directly instantiated by users, comprehending its underlying mechanisms is crucial for debugging and effectively configuring your Flyte tasks. You will learn how flyte-sdk handles image pulling, volume binding for I/O, and streaming container logs during local execution.

Prerequisites​

To follow along and run the examples, ensure you have:

  • Docker installed and running on your local machine.
  • The flyte-sdk installed in your Python environment. You may need to install the docker library explicitly: pip install docker.

Understanding ContainerTask's Role​

In flyte-sdk, every @task decorated function that runs within a container is internally represented by a ContainerTask instance. This class is responsible for orchestrating the container's lifecycle, from preparing its execution environment to capturing its outputs. It acts as the blueprint for how your task's container will behave.

You configure the ContainerTask indirectly through parameters provided to the @task decorator. Key parameters that influence local execution include:

  • container_image: The Docker image to use for the task (e.g., "ubuntu:latest").
  • container_command: The entrypoint command for the container (e.g., ["bash", "-cx"]).
  • container_args: Arguments passed to the container_command. These can include templated inputs like {{.inputs.my_input}}.
  • input_data_dir: The directory inside the container where input data will be mounted. Defaults to /var/inputs.
  • output_data_dir: The directory inside the container where output data should be written. Defaults to /var/outputs.
  • local_logs: A boolean flag (default True) that, when set, streams container logs directly to your console during local execution.

Local Execution Flow​

When you execute a Flyte workflow locally, flyte-sdk simulates the Flyte engine's behavior using your local Docker daemon. The ContainerTask.execute method orchestrates this process:

  1. Input and Output Directory Normalization: The input_data_dir and output_data_dir paths are normalized. A temporary local directory is created to serve as the host-side mount point for the container's output.

  2. Command and Volume Preparation: The _prepare_command_and_volumes method processes the container_command and container_args. It performs two critical actions:

    • Input Templating: It replaces placeholders like {{.inputs.my_input}} (as seen in _render_command_and_volume_binding) with the actual local file paths where flyte-sdk has staged your task's inputs.
    • Volume Binding: For flytekit.io.File and flytekit.io.Dir inputs, _render_command_and_volume_binding identifies these and creates Docker volume bindings. This maps a local path on your host machine (where the input file/directory resides) to a corresponding path inside the container (e.g., within input_data_dir). Similarly, the temporary local output directory is bound to the container's output_data_dir.

  3. Image Pulling: The _pull_image_if_not_exists method checks if the specified container_image is available locally. If not, it pulls the image from Docker Hub (or your configured registry), ensuring the container has the necessary environment to run.

  4. Container Execution: A Docker container is launched using docker.from_env().containers.run. The prepared command, arguments, and volume bindings are passed to the Docker client. The container is run in detached mode (detach=True) and configured to be removed automatically upon completion (remove=True).

  5. Log Streaming: If local_logs is True, the execute method continuously streams logs from the running container to your console, providing real-time feedback on your task's progress.

  6. Output Retrieval: Once the container finishes, _get_output reads the results from the output_data_dir (which is volume-mounted to a local temporary directory). It then converts these raw outputs into the Python types specified in your task's signature.

Example: Reversing a String in a Container​

Let's create a simple task that takes a string, writes it to a file, and then uses a container to reverse the string and write the result to another file. This demonstrates input/output handling and command execution within a custom container.

First, define your task and workflow:

import os
from flytekit import task, workflow
from flytekit.core.image import Image
from flytekit.io import File

@task(
    container_image="ubuntu:latest",
    container_command=["bash", "-cx"],
    container_args=[
        "echo 'Hello from container!';",
        "cat {{.inputs.input_str_file}} > /tmp/input.txt;",
        "rev /tmp/input.txt > {{.outputs.output_str_file}};",
        "echo 'Container finished processing.'",
    ],
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    local_logs=True,
)
def reverse_string_task(input_str_file: File) -> File:
    # The body of this function is not executed during local container task execution.
    # It primarily defines the task's interface (inputs and outputs).
    pass

@workflow
def reverse_string_workflow(input_string: str) -> str:
    # When running locally, Flytekit automatically handles converting the input_string
    # into a File object that the container task expects.
    reversed_file = reverse_string_task(input_str_file=input_string)
    # Similarly, Flytekit reads the content of the output File and returns it as a string.
    return reversed_file

if __name__ == "__main__":
    # To execute the workflow locally, simply call it as a regular Python function.
    print("Running reverse_string_workflow locally...")
    result = reverse_string_workflow(input_string="hello flyte")
    print(f"Local execution result: {result}")

When you run this Python script, flyte-sdk will:

  1. Create a temporary local file containing "hello flyte".
  2. Launch a Docker container using the ubuntu:latest image.
  3. Mount the temporary input file into the container at a path like /var/inputs/input_str_file.
  4. Execute the bash command, which reads the input, reverses it, and writes the reversed string to a file within the container's /var/outputs directory.
  5. Stream the container's stdout and stderr to your console because local_logs is True.
  6. Mount the container's /var/outputs directory to another temporary local directory.
  7. Read the reversed string from the output file in the temporary local directory.
  8. Return the reversed string "etyfl olleh" as the result of the workflow.

Next Steps​

  • Experiment with different container_command and container_args to see how the container's behavior changes.
  • Try setting local_logs=False in the @task decorator to observe the difference in log output during local execution.
  • Explore how flytekit.io.Dir can be used to pass entire directories as inputs or outputs to your container tasks. This will involve ContainerTask creating additional volume bindings for the directories. This is handled by the _render_command_and_volume_binding method, which detects File and Dir types and configures the appropriate Docker mounts.