Defining Container Images

When building tasks in flyte-sdk, you often need to specify a custom container image that includes your project's dependencies, code, and specific configurations. The flyte.Image class provides a fluent and declarative API to define these container images, allowing you to start from a base image and progressively add layers.

Building Images with flyte.Image

The flyte.Image class follows a two-step pattern for image specification:

1. Start with a base image using one of the from_* class methods.
2. Customize the image by chaining with_* instance methods, each of which returns a new Image object with the added layer.

This immutable, fluent API ensures that each customization step is explicit and easy to reason about.

Base Image Constructors (from_*)

flyte.Image offers several constructors to define your starting point:

Image.from_debian_base()

This is a common starting point for Python-based projects, providing a Debian-based image with a specified Python version. It automatically handles the initial setup for a Python environment.

import flyte

my_image = (
    flyte.Image.from_debian_base(python_version=(3, 12)) # Start with Debian and Python 3.12
    .with_pip_packages("pandas", "scikit-learn") # Add Python packages
    .with_apt_packages("curl", "git") # Add system packages
)

# Define a TaskEnvironment instance to associate the image with tasks
my_env = flyte.TaskEnvironment(name="my_tasks_env", image=my_image)

@my_env.task
def my_task(x: int) -> int:
    import pandas as pd
    df = pd.DataFrame({'col': [x]})
    return df['col'].sum()

Image.from_base()

Use Image.from_base() when you want to start with an arbitrary pre-built container image by its full URI (e.g., myregistry.com/myimage:v1.0). This image is assumed to already exist in a registry.

import flyte

# Start with a pre-existing image
base_image = flyte.Image.from_base("ubuntu:22.04")

# You can still add Python-specific layers if the base image has Python installed
my_custom_image = base_image.with_pip_packages("requests")

Important: Images created with from_base() are generally not intended for further system-level modifications using methods like with_apt_packages(). Attempting to do so will raise a ValueError:

import flyte

with_apt_error_image = flyte.Image.from_base("myregistry.com/myimage:v123")

# This will raise a ValueError
try:
    with_apt_error_image.with_apt_packages("vim")
except ValueError as e:
    print(f"Error: {e}")

Image.from_uv_script()

If your project uses uv scripts for dependency management, you can build an image directly from a uv script file. The script's header defines the Python version and dependencies.

Consider a my_script.py file:

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.12"
# dependencies = ["httpx"]
# ///

import httpx

def main():
    response = httpx.get("https://example.com")
    print(response.status_code)

if __name__ == "__main__":
    main()

You can then create an image from this script:

import flyte
from pathlib import Path

# Assuming my_script.py is in the current directory
script_path = Path("my_script.py")

uv_image = flyte.Image.from_uv_script(
    script=script_path,
    name="my-uv-app",
    registry="localhost",
    python_version=(3, 12)
)

# Define a TaskEnvironment instance
uv_env = flyte.TaskEnvironment(name="uv_script_env", image=uv_image)

@uv_env.task
def run_uv_script():
    # The script will be available in the container
    pass

Image.from_dockerfile()

For full control over your image build process, you can provide a custom Dockerfile. Note that flyte-sdk does not parse or understand the contents of your Dockerfile, so subsequent with_* methods for adding layers are not supported.

Consider a Dockerfile named Dockerfile.custom:

FROM python:3.9-slim-buster
WORKDIR /app
COPY requirements.txt .
RUN pip install -r requirements.txt
COPY . .
CMD ["python", "app.py"]

import flyte
from pathlib import Path

# Assuming Dockerfile.custom is in the current directory
dockerfile_path = Path("Dockerfile.custom")

dockerfile_image = flyte.Image.from_dockerfile(
    file=dockerfile_path,
    name="my-dockerfile-app",
    registry="localhost"
)

# Define a TaskEnvironment instance
dockerfile_env = flyte.TaskEnvironment(name="dockerfile_env", image=dockerfile_image)

@dockerfile_env.task
def run_dockerfile_app():
    # Your Dockerfile defines the environment
    pass

Image.from_ref_name()

This method allows you to reference a pre-built image by a symbolic name, which can be configured externally (e.g., via CLI or Flyte configuration). This is useful for decoupling image URIs from your code.

import flyte

# Reference an image named 'my_configured_image'
# The actual URI for 'my_configured_image' would be resolved at runtime
ref_image = flyte.Image.from_ref_name("my_configured_image")

# Define a TaskEnvironment instance
ref_env = flyte.TaskEnvironment(name="ref_image_env", image=ref_image)

@ref_env.task
def task_using_ref_image():
    pass

Customization Methods (with_*)

Once you have a base image, you can chain with_* methods to add various layers and configurations. Each with_* method returns a new Image object, preserving the immutability of the image definition.

with_pip_packages()

Installs Python packages using pip. You can specify multiple packages and provide custom index URLs or extra arguments.

import flyte

my_image = (
    flyte.Image.from_debian_base()
    .with_pip_packages("requests", "numpy==1.23.0") # Install specific versions
)

# Define a TaskEnvironment instance
pip_env = flyte.TaskEnvironment(name="pip_packages_env", image=my_image)

@pip_env.task
def fetch_data():
    import requests
    response = requests.get("https://api.example.com")
    return response.json()

To install private packages, you can mount secrets during the build process:

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: 'flyte.Secret' accepts at most 0 positional argument(s)

import flyte

private_package_image = (
    flyte.Image.from_debian_base()
    .with_pip_packages(
        "git+https://$GITHUB_PAT@github.com/myorg/private-repo.git",
        secret_mounts=[flyte.Secret("GITHUB_PAT")] # Mount GITHUB_PAT as an environment variable
    )
)

with_apt_packages()

Installs system-level packages using apt-get (for Debian-based images).

import flyte

apt_image = (
    flyte.Image.from_debian_base()
    .with_apt_packages("build-essential", "libgdal-dev") # Install C/C++ build tools and GDAL development files
)

with_env_vars()

Sets environment variables within the container.

import flyte

env_var_image = (
    flyte.Image.from_debian_base()
    .with_env_vars({"MY_API_KEY": "some_value", "DEBUG_MODE": "true"})
)

# Define a TaskEnvironment instance
env_var_env = flyte.TaskEnvironment(name="env_var_env", image=env_var_image)

@env_var_env.task
def check_env_var() -> str:
    import os
    return os.getenv("MY_API_KEY", "Not Found")

with_requirements()

Installs Python packages from a requirements.txt file. This is an alternative to with_pip_packages() for managing dependencies.

# Assuming a requirements.txt file exists:
# pandas==2.0.0
# matplotlib

import flyte
from pathlib import Path

# Create a dummy requirements.txt for the example
Path("requirements.txt").write_text("pandas==2.0.0\nmatplotlib")

req_image = (
    flyte.Image.from_debian_base()
    .with_requirements("requirements.txt")
)

# Define a TaskEnvironment instance
req_env = flyte.TaskEnvironment(name="req_env", image=req_image)

@req_env.task
def plot_data():
    import matplotlib.pyplot as plt
    plt.plot([1, 2, 3])
    # ... save plot ...

Important: The file specified in with_requirements() must exist at the time of image definition, otherwise a FileNotFoundError will be raised.

with_uv_project() and with_poetry_project()

These methods integrate with uv and Poetry-based Python projects, respectively. They handle copying pyproject.toml, lock files, and installing project dependencies.

import flyte
from pathlib import Path

# Assuming pyproject.toml and uv.lock exist for a uv project
uv_proj_image = (
    flyte.Image.from_debian_base()
    .with_uv_project(pyproject_file=Path("path/to/my_uv_project/pyproject.toml"))
)

# Assuming pyproject.toml and poetry.lock exist for a Poetry project
poetry_proj_image = (
    flyte.Image.from_debian_base()
    .with_poetry_project(pyproject_file=Path("path/to/my_poetry_project/pyproject.toml"))
)

with_source_folder() and with_source_file()

These methods allow you to copy local directories or individual files into the container image. This is crucial for including your application code or configuration files.

import flyte
from pathlib import Path

# Create dummy files/folders for the example
Path("my_module.py").write_text("def hello(): return 'Hello from module'")
Path("config/").mkdir(exist_ok=True)
Path("config/settings.yaml").write_text("key: value")

# Copy a single file
file_image = (
    flyte.Image.from_debian_base()
    .with_source_file(src=Path("my_module.py"), dst="/app/")
)

# Copy an entire folder
folder_image = (
    flyte.Image.from_debian_base()
    .with_source_folder(src=Path("config/"), dst="/etc/app_config/")
)

# Define a TaskEnvironment instance for file_image
file_env = flyte.TaskEnvironment(name="file_copy_env", image=file_image)

@file_env.task
def use_copied_module() -> str:
    import sys
    sys.path.append("/app/") # Ensure the copied path is in PYTHONPATH
    import my_module
    return my_module.hello()

with_workdir()

Sets the working directory inside the container. This affects where commands are executed and where relative paths resolve.

import flyte

workdir_image = (
    flyte.Image.from_debian_base()
    .with_workdir("/app/my_project")
)

with_dockerignore()

Specifies a .dockerignore file to exclude certain files or directories from being copied into the build context, optimizing build times and image size.

import flyte
from pathlib import Path

# Create a dummy .dockerignore file
Path(".dockerignore").write_text("*.log\n__pycache__/")

dockerignore_image = (
    flyte.Image.from_debian_base()
    .with_dockerignore(path=Path(".dockerignore"))
)

with_commands()

Executes arbitrary shell commands during the image build process. This is useful for tasks not covered by other with_* methods, such as installing custom software or configuring system settings.

import flyte

commands_image = (
    flyte.Image.from_debian_base()
    .with_commands([
        "apt-get update",
        "apt-get install -y some-utility",
        "mkdir -p /data"
    ])
)

Integrating Images with flyte.TaskEnvironment

While flyte.Image defines the container specification, flyte.TaskEnvironment is responsible for encapsulating the execution environment for a group of tasks. This includes the container image, compute resources, environment variables, and other configurations.

flyte-sdk uses a three-level hierarchy for task configuration:

1. TaskEnvironment: Sets default configurations for all tasks within that environment (most general).
2. @env.task decorator: Overrides settings for a specific task.
3. task.override(): Provides overrides at the invocation time of a task (most specific).

When you define an Image object, you typically assign it to the image parameter of a TaskEnvironment.

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: 'flyte.Resources' does not accept keyword(s): cpu, memory

import flyte

# Define a custom image
my_env_image = (
    flyte.Image.from_debian_base(python_version=(3, 11))
    .with_pip_packages("requests")
)

# Create a TaskEnvironment and assign the image
my_env = flyte.TaskEnvironment(
    name="my_custom_environment",
    image=my_env_image,
    resources=flyte.Resources(cpu="1", memory="1Gi"),
    env_vars={"ENV_SETTING": "production"}
)

@my_env.task
async def my_api_task():
    import requests
    response = requests.get("https://api.example.com/data")
    return response.json()

In this example, my_api_task will execute within a container built from my_env_image, inheriting the specified resources and environment variables from my_env.

Cloning TaskEnvironment with clone_with()

If you need multiple environments that share a common base configuration but differ in a few settings, you can use the clone_with() method. This avoids repetitive definitions.

import flyte

base_env = flyte.TaskEnvironment(
    name="base_env",
    image=flyte.Image.from_debian_base().with_pip_packages("numpy"),
    resources=flyte.Resources(),
)

# Clone the base environment and override some settings
dev_env = base_env.clone_with(
    name="dev_env",
    image=flyte.Image.from_debian_base(python_version=(3, 10)).with_pip_packages("pytest"),
    env_vars={"DEBUG": "true"},
    cache="disable" # Override cache policy
)

@dev_env.task
async def dev_task():
    # This task runs in 'dev_env' with its specific image and settings
    pass

Important: When you clone a TaskEnvironment using clone_with(), any tasks previously defined within the original environment are not automatically transferred to the new cloned environment. You must explicitly re-define or add tasks to the cloned environment if needed.

Gotchas and Warnings

• TaskEnvironment.name Validation: The name parameter for TaskEnvironment must follow snake_case or kebab-case conventions. Invalid names will result in a ValueError.
• Strict Type Checking: Many parameters in TaskEnvironment (e.g., image, resources, secrets) have strict type requirements. Providing incorrect types will lead to a TypeError.
• reusable and pod_template Conflict: If a TaskEnvironment is configured with reusable containers, you cannot also specify a pod_template for individual tasks within that environment. Doing so will raise a ValueError.