Fix @task.kubernetes to receive input and send output

[vchiapaikeo]

closes: #28933

@task.kubernetes currently does not handle receiving input nor does it properly return output. This PR addresses these issues by passing input to the K8's pod (via an env var) and output where the K8's xcom sidecar expects it (as /airflow/xcom/result.json).

It also seemed like the initial implementation did not properly base64 encode / decode the values. Since there might be unusual characters / line breaks in both the function and inputs, I added base64 encoding to the env var setting (similar to how it's handled in docker).

Testing

To test, I started Airflow w/ breeze (breeze --python 3.9 --backend postgres start-airflow), added a google_cloud_default connection_id to my sandbox GCP project, and copied my kube_config file to /files/.kube/config. Because I am using GCP, I needed to install gcloud. I do so with the following commands in the breeze shell:

curl -sSL https://sdk.cloud.google.com > /tmp/gcloud_installer && bash /tmp/gcloud_installer --install-dir=$AIRFLOW_HOME --disable-prompts
source /root/airflow/google-cloud-sdk/completion.bash.inc
source /root/airflow/google-cloud-sdk/path.bash.inc

I then ran this test DAG:

import os

from airflow import DAG
from airflow.decorators import task

DEFAULT_TASK_ARGS = {
    "owner": "gcp-data-platform",
    "start_date": "2022-12-16",
    "retries": 0,
}

@task.kubernetes(
    image="python:3.8-slim-buster",
    namespace=os.getenv("AIRFLOW__KUBERNETES_EXECUTOR__NAMESPACE", "airflow-default"),
    in_cluster=False,
    config_file="/files/.kube/config",
)
def k8s_basic() -> str:
    import time
    time.sleep(1)


@task.kubernetes(
    image="python:3.8-slim-buster",
    namespace=os.getenv("AIRFLOW__KUBERNETES_EXECUTOR__NAMESPACE", "airflow-default"),
    in_cluster=False,
    config_file="/files/.kube/config",
    # multiple_outputs=True,
)
def k8s_func(val: str = "a") -> str:
    import time
    time.sleep(1)

    print(f"Got val: {val}")
    return {"a": val, "d": [1, 2, 3], "f": {1: val}}


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

    basic = k8s_basic()

    no_input_or_output = k8s_func()

    with_input = k8s_func.override(task_id="with_input")("b")

    with_input_and_output = k8s_func.override(task_id="with_input_and_output", do_xcom_push=True)("b")

    no_input_and_output = k8s_func.override(task_id="no_input_and_output", do_xcom_push=True)()

Results

With Input:

With Input and Output:

No input and Output:

With multiple_outputs = 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]

Hi @eladkal, do you know who might be a good person to review this?

[vchiapaikeo]

Modified with b64 encoded contents sent straight to the Py commands and input serialized with pickle instead of json. Retested and results are as expected:

[vchiapaikeo]

Oops - sorry for pinging everyone. Forgot that I had done a rebase on the UI and didn't pull first. Feel free to unassign here!

[vchiapaikeo]

The shlex.quote here isn't really necessary anymore since they don't capture user input but I left it in anyways. Thought that made it a bit more readable but I can remove if you prefer.

[vchiapaikeo]

Hey @uranusjr , anything else I can do here?

[uranusjr]

lgtm

[vchiapaikeo]

Can someone help merge this?

[vchiapaikeo]

@uranusjr , can you help merge this please?

[fletchjeff]

Hey @vchiapaikeo, I was just looking in to this very problem and saw this fix. Nice work! Thank you :)

[okulbida]

@fletchjeff Do you know when this gonna be released?

[vchiapaikeo]

@okulbida , this went out with cncf providers release 5.2.1

https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/stable/index.html#id4

[asafm]

Just verifying here. It seems that the input as passed via ENV VAR to the KubernetesPodOperator, meaning if you used XCom Backend like S3 to avoid the problem of a large xcom, this won't solve this, since upon executing your shell command, you will probably fail on exec {...your command name...} argument list too long due to a very large value of the ENV VAR.