KPO pods can hang if sidecar never exits #58968

from airflow import __version__ as airflow_version
from airflow.decorators import dag
from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
from kubernetes.client import models as k8s
from pendulum import datetime


def sleep_container(name: str, duration: int | str) -> k8s.V1Container:
    return k8s.V1Container(
        name=name,
        image="busybox:latest",
        command=["/bin/sh", "-c"],
        args=[f"sleep {duration}"],
    )


def healthy_pod_spec() -> k8s.V1PodSpec:
    return k8s.V1PodSpec(
        containers=[
            sleep_container("base", "20"),
        ],
        termination_grace_period_seconds=86400,
        restart_policy="Never",
    )


def zombie_pod_spec() -> k8s.V1PodSpec:
    return k8s.V1PodSpec(
        containers=[
            sleep_container("base", "5"),
            sleep_container("sidecar", "infinity"),
        ],
    )


@dag(
    dag_id="kpo_zombie_dag",
    schedule=None,
    start_date=datetime(2024, 1, 1),
    catchup=False,
    tags=["test", "kpo", "zombie"],
)
def kpo_zombie_dag():
    zombie_task = KubernetesPodOperator(
        task_id="zombie_kpo_task",
        full_pod_spec=k8s.V1Pod(
            metadata=k8s.V1ObjectMeta(),
            spec=zombie_pod_spec(),
        ),
    )
    healthy_task = KubernetesPodOperator(
        task_id="healthy_kpo_task",
        full_pod_spec=k8s.V1Pod(
            metadata=k8s.V1ObjectMeta(),
            spec=healthy_pod_spec(),
        ),
    )

    tasks = [zombie_task, healthy_task]

    # Create 3 mapped healthy KPO tasks (only for Airflow >= 2.3.0)
    # Mapped tasks were introduced in Airflow 2.3.0
    try:
        # Parse version string (e.g., "2.3.0" or "2.3.0+astro.1")
        version_parts = airflow_version.split("+")[0].split(".")
        major = int(version_parts[0])
        minor = int(version_parts[1]) if len(version_parts) > 1 else 0

        supports_mapped_tasks = major > 2 or (major == 2 and minor >= 3)

        if supports_mapped_tasks:
            # Use expand_kwargs to create 2 mapped task instances (map_index 0, 1)
            # one zombie and one healthy
            # Each dict in the list creates one mapped instance
            mapped_healthy_task = KubernetesPodOperator.partial(
                task_id="mapped_kpo_task",
            ).expand_kwargs(
                [
                    {
                        "full_pod_spec": k8s.V1Pod(
                            metadata=k8s.V1ObjectMeta(),
                            spec=zombie_pod_spec(),
                        ),
                    },
                    {
                        "full_pod_spec": k8s.V1Pod(
                            metadata=k8s.V1ObjectMeta(),
                            spec=healthy_pod_spec(),
                        ),
                    },
                ]
            )
            tasks.append(mapped_healthy_task)
    except (ValueError, IndexError):
        # If version parsing fails, skip mapped tasks to be safe
        pass

    return tasks


kpo_dag = kpo_zombie_dag()