Push to xcom before KubernetesPodOperator deferral

[vchiapaikeo]

Currently, deferrable KPO will not allow users to see logs of the underlying base container until completion (even with get_logs=True). This is because logs are not streamed until the trigger event successfully fires and execute_complete is called. What's worse is that we do not know what the pod_name or the pod_namespace is until execute_complete is called as well. With synchronous KPO, this happens immediately after get_or_create_pod is called. We should do the same with the deferrable flow. Therefore, users can add extra_links to the KPO (retrieving pod_name and pod_namespace from xcom) and stream in-flight logs from their monitoring tool of choice (for us, it is Cloud Monitoring because we use GKE).

This moves the xcom_push calls up to execute_async. To test, I used this dag and observed xcom appear during deferral state.

from airflow import DAG

from airflow.providers.google.cloud.operators.kubernetes_engine import (
    GKEStartPodOperator,
)


DEFAULT_TASK_ARGS = {
    "owner": "gcp-data-platform",
    "start_date": "2021-04-20",
    "retries": 0,
    "retry_delay": 60,
}

with DAG(
    dag_id="test_gke_op",
    schedule_interval="@daily",
    max_active_runs=1,
    max_active_tasks=5,
    catchup=False,
    default_args=DEFAULT_TASK_ARGS,
) as dag:

    whoami = GKEStartPodOperator(
        task_id="whoami",
        name="whoami",
        cmds=["gcloud"],
        arguments=["auth", "list"],
        image="gcr.io/google.com/cloudsdktool/cloud-sdk:slim",
        project_id="redacted-project-id",
        namespace="airflow-default",
        location="us-central1",
        cluster_name="airflow-gke-cluster",
        service_account_name="default",
        is_delete_operator_pod=True,
    )

    sleep_3000 = GKEStartPodOperator(
        task_id="sleep_3000",
        name="sleep-3000",
        cmds=["sleep"],
        arguments=["3000"],
        image="gcr.io/google.com/cloudsdktool/cloud-sdk:slim",
        project_id="redacted-project-id",
        namespace="airflow-default",
        location="us-central1",
        cluster_name="airflow-gke-cluster",
        service_account_name="default",
        is_delete_operator_pod=True,
        deferrable=True,
    )

    sleep_3600 = GKEStartPodOperator(
        task_id="sleep_3600",
        name="sleep-3600",
        cmds=["sleep"],
        arguments=["3600"],
        image="gcr.io/google.com/cloudsdktool/cloud-sdk:slim",
        project_id="redacted-project-id",
        namespace="airflow-default",
        location="us-central1",
        cluster_name="airflow-gke-cluster",
        service_account_name="default",
        is_delete_operator_pod=True,
        deferrable=True,
    )

    sleep_4000 = GKEStartPodOperator(
        task_id="sleep_4000",
        name="sleep-4000",
        cmds=["sleep"],
        arguments=["4000"],
        image="gcr.io/google.com/cloudsdktool/cloud-sdk:slim",
        project_id="redacted-project-id",
        namespace="airflow-default",
        location="us-central1",
        cluster_name="airflow-gke-cluster",
        service_account_name="default",
        is_delete_operator_pod=True,
    )

    sleep_4800 = GKEStartPodOperator(
        task_id="sleep_4800",
        name="sleep-4800",
        cmds=["sleep"],
        arguments=["4800"],
        image="gcr.io/google.com/cloudsdktool/cloud-sdk:slim",
        project_id="redacted-project-id",
        namespace="airflow-default",
        location="us-central1",
        cluster_name="airflow-gke-cluster",
        service_account_name="default",
        is_delete_operator_pod=True,
    )

---

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

[vchiapaikeo]

@hussein-awala , @eladkal - thoughts on this when you have a moment please. This would help us out a lot on the observability front when moving KPO to deferrable mode and streaming logs of in flight pods.

[hussein-awala]

Looks good for me!

[hussein-awala]

FYI, I'm working on loading the logs in async mode without waiting for pod termination, it will be possible soon.

[vchiapaikeo]

@eladkal , @potiuk - can you help merge when you have a moment please?

[eladkal]

LGTM