Skip to main content

Kubernetes Pod Customization

Kubernetes Pod customization in this SDK is managed through the PodTemplate class, which provides a mechanism to inject low-level Kubernetes configurations into the execution environment of Flyte tasks. This allows developers to define sidecar containers, mount custom volumes, and attach metadata like labels and annotations that are not exposed through the standard task decorators.

The PodTemplate Class​

The PodTemplate class (located in src/flyte/_pod.py) acts as a wrapper around the standard Kubernetes V1PodSpec. It consists of four primary fields:

  • pod_spec: A kubernetes.client.V1PodSpec object defining the desired state of the Pod (containers, volumes, affinity, etc.).
  • primary_container_name: The name of the container within the pod_spec that Flyte should treat as the main execution engine. This defaults to "primary".
  • labels: A dictionary of Kubernetes labels to apply to the Pod metadata.
  • annotations: A dictionary of Kubernetes annotations to apply to the Pod metadata.

When a task is executed, Flyte uses the to_k8s_pod method to serialize these components into a K8sPod message, ensuring the pod_spec is sanitized for the Kubernetes API.

Applying PodTemplates​

PodTemplates can be applied at three different levels of specificity, following the standard configuration hierarchy in TaskEnvironment.

Environment-Wide Configuration​

You can set a default PodTemplate for all tasks within a TaskEnvironment. This is useful for infrastructure-level requirements like image pull secrets or common sidecars.

import flyte
from kubernetes.client import V1PodSpec, V1LocalObjectReference

# Define a template with image pull secrets
common_pod_template = flyte.PodTemplate(
    pod_spec=V1PodSpec(
        image_pull_secrets=[V1LocalObjectReference(name="my-registry-creds")],
    ),
    labels={"team": "data-science"}
)

env = flyte.TaskEnvironment(
    name="production-env",
    pod_template=common_pod_template,
    image="my-custom-image:latest"
)

Per-Task Configuration​

The @env.task decorator allows overriding the environment's default template for a specific task.

@env.task(pod_template=specialized_template)
async def specialized_task(x: int) -> int:
    return x + 1

Dynamic Overrides​

Tasks can have their PodTemplate overridden at invocation time using the .override() method. This is particularly useful for workflows that need to adjust infrastructure dynamically based on input data.

@env.task
async def my_task(data: str):
    ...

# Override the template for a specific call
custom_call = my_task.override(pod_template=dynamic_template)(data="input")

Advanced Use Cases​

Mounting Volumes (GCSFuse Example)​

One of the most common uses for PodTemplate is mounting specialized volumes, such as GCSFuse for high-performance cloud storage access. This requires both annotations for the CSI driver and volumes/volume_mounts in the pod_spec.

from kubernetes.client import V1Container, V1Volume, V1VolumeMount, V1CSIVolumeSource

gcs_fuse_template = flyte.PodTemplate(
    primary_container_name="primary",
    annotations={
        "gke-gcsfuse/volumes": "true",
        "gke-gcsfuse/cpu-limit": "4",
    },
    pod_spec=V1PodSpec(
        containers=[
            V1Container(
                name="primary",
                volume_mounts=[
                    V1VolumeMount(
                        name="gcs-fuse",
                        mount_path="/mnt/gcs",
                        read_only=False,
                    )
                ],
            )
        ],
        volumes=[
            V1Volume(
                name="gcs-fuse",
                csi=V1CSIVolumeSource(
                    driver="gcsfuse.csi.storage.gke.io",
                    volume_attributes={
                        "bucketName": "my-data-bucket",
                        "mountOptions": "implicit-dirs,uid=65532,gid=65532",
                    },
                ),
            )
        ],
    ),
)

Sidecar Containers​

You can define sidecar containers by adding multiple entries to the containers list in the V1PodSpec. Flyte will inject the task's code and logic into the container specified by primary_container_name, while the other containers run alongside it.

sidecar_template = flyte.PodTemplate(
    primary_container_name="primary",
    pod_spec=V1PodSpec(
        containers=[
            V1Container(name="primary"),
            V1Container(
                name="logging-agent",
                image="fluentd:latest",
            )
        ]
    )
)

Constraints and Requirements​

Primary Container Naming​

The primary_container_name field is critical. If you provide a pod_spec with a list of containers, one of those containers must match the name specified in primary_container_name (which defaults to "primary"). Flyte merges task-specific configurations (like the image, command, and environment variables) into this specific container.

Incompatibility with Reusable Environments​

In src/flyte/_task_environment.py, the SDK enforces a strict separation between PodTemplate and reusable environments. If a TaskEnvironment is initialized with reusable=True (or any ReusePolicy), attempting to set a pod_template will raise a ValueError.

# This will raise a ValueError
env = flyte.TaskEnvironment(
    name="invalid-env",
    reusable=flyte.ReusePolicy(concurrency=1),
    pod_template=my_template
)

This restriction exists because reusable environments rely on long-lived containers where the Pod lifecycle is managed differently than standard task executions.

Named Pod Templates​

In addition to passing PodTemplate objects, the pod_template parameter in TaskEnvironment and .override() also accepts a string. This string refers to a named PodTemplate that has been pre-registered on the Flyte cluster by an administrator. This allows for centralized management of complex infrastructure configurations.