Feature similar to return.json but for pulling XCom in KubernetesPodOperator

[orphefs]

I recently came across a limitation where I need to pull an XCom result (a JSON string) into a KubernetesPodOperator task. The final goal is to be able to serialize this result into the EFS volume the pod has mounted. Then I can use that file in downstream tasks.

My only way of interacting with the XCom as of now is to use a Jinja template and have it render at runtime in the pod terminal:

echo -e {{{{ ti.xcom_pull('task_id') }}}} >> /efs_volume/file.json

When the contents of the XCom become large enough, the terminal command hits the OS character limit. My workaround right now is to push the file into some temporary network storage (i.e. S3), pass the URI into XCom, then use the URI to download the file into the pod. Obviously, this is not an ideal solution.

I was thinking that this can be done in a better way: introduce a new feature symmetric to the return.json concept, but for pulling XComs. E.g. introducing an argument do_xcom_pull_ids in the KubernetesPodOperator constructor, which specify which upstream XComs to make available inside the sidecar container, then serialize those files during setup so that they are available as e.g. task_id_X.json, task_id_Y.json, etc. for the user to pull.

If you also think this is a feature worth introducing, I'd be happy to give a it a go.

[potiuk]

I do not think this is needed.

If you want to pull XCom values in your KPO (or Docker operator) you can use the image with airflow to the Pod and you could simply run (providing the environment variables are passed):

 XCom.get_value(key=self.key, ti_key=ti_key)

Where `ti_key' is the compound key of what you want to pull. This will even make the values lazy - so if you are using Custom XCom backend, you will ony retrieve the value when you use it.

We even have a chapter in our docs describing how you can check if XCom is correctly initialized in the container you run:

https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/xcoms.html#working-with-custom-xcom-backends-in-containers

[orphefs]

Thanks for your reply @potiuk and apologies for my belated response.

I managed to put together an minimum working example that achieves injecting XCom via an environment variable INPUT_DATA, reading from that XCom via os.getenv() and writing to xcom.json file in a python script which is passed to the pod arguments.

Is this what you meant in your reply, or close to it? If not, would you kindly be able to provide an example?

# airflow config
global_trigger_rule = "none_failed"

default_args: Dict = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': airflow.utils.dates.days_ago(10),
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 0,
    'retry_delay': timedelta(minutes=60),
}

# The following task arguments are reused often

global_task_args: Dict = {
    "image_pull_secrets": [k8s.V1LocalObjectReference("regcred")],
    "is_delete_operator_pod": True,
    "namespace": NAMESPACE,
    "service_account_name": SERVICE_ACCOUNT_NAME,
    "log_events_on_failure": True,
}
dag_obj: DAG = DAG('dag_kpo_with_xcom_injection',
                   default_args=default_args,
                   description='Demonstration of using KubernetesPodOperator with XCom injection on Airflow',
                   catchup=False,
                   render_template_as_native_obj=True,
                   tags=['docker', 'k8s'], )


def escape_characters(string: str) -> str:
    return "'" + string.replace('"', '\"') + "'"


def push_some_xcom(**kwargs):
    d = {k: v for k, v in zip(range(0, 5), range(0, 5))}
    json_string = json.dumps(d)
    return escape_characters(json_string)


task_push_some_xcom: PythonOperator = PythonOperator(
    task_id='push_some_xcom',
    python_callable=push_some_xcom,
    do_xcom_push=True,
    provide_context=True,
    trigger_rule=global_trigger_rule,
    dag=dag_obj
)

python_script = '''
import os
import json

f = open("xcom.json", "a")
f.write(json.dumps(json.loads(os.getenv("INPUT_DATA"))))
f.close()

f = open("xcom.json", "r")
pulled_dict = json.loads(f.read())
print(pulled_dict)
print(type(pulled_dict))
f.close()

'''

task_pull_xcom = KubernetesPodOperator(
    name="pull_xcom",
    image="apache/airflow:2.5.0",
    cmds=["bash"],
    arguments=["-c",
               f"echo $INPUT_DATA && python -c '{python_script}' && cat xcom.json && sleep 10 && curl -sf -XPOST http://127.0.0.1:15020/quitquitquit"],
    task_id="pull_xcom",
    do_xcom_push=False,
    env_vars=[k8s_client.V1EnvVar(name="INPUT_DATA",
                                  value="{{ ti.xcom_pull(task_ids='push_some_xcom') }}")],
    trigger_rule=global_trigger_rule,
    dag=dag_obj,
    **global_task_args
)

task_push_some_xcom >> task_pull_xcom

[potiuk]

Roughly speaking - yes. But I was thinking more that you could (rather than having the script and passing it to bash, have your python script to read Xcom embedded in your image that you run as KPO and execute that python code directly there.

[potiuk]

Like - it really makes little sense to use KPO to execute code that is part of the DAG. For that Kubernetes Executor seems much better choice- provides the same isolation and it executes your Python code from DAG/task directly. KPO you usualy use to execute code that is "out there" - part of the image, rather than passing the code from the DAG to be executed there.

[orphefs]

Roughly speaking - yes. But I was thinking more that you could (rather than having the script and passing it to bash, have your python script to read Xcom embedded in your image that you run as KPO and execute that python code directly there.

I am not sure how to read Xcom embedded in the image, is there an example of that somewhere? Is the method you are suggesting still employing passing the actual XCom through env vars?

Passing large ( >600 kB ) environment variables make K8s raise an exception Exception when attempting to create Namespaced Pod. Lazily evaluating the XCom just when the script is being executed inside the pod would be ideal.

[potiuk]

Not the XCom "embedded" but the "Code to read XCom embedded". What I mean is that simply rather than sending the code to execute via string in the "python_script", the same code should be baked in the custom image you are running. It means that you have to have pipeline that rebuilds and publishes the image every time you want to change that code. But the advantage of it is that your command to run will be just "run this script embedded in the image" and it will take care about XCom exchange then

[orphefs]

Ok, I see, thank you for clarifying this. Does this method still imply the use of injecting environment variable to pass the XCom?

[potiuk]

I belive for now yes. But there are already attempts to bake in some Airflow "environment" into KPO (for example #33680). I think it might be a good idea to extend that to all Airflow configuration and simply be able to run KPO with "same configuration as Airflow" in terms of automatically passing all (or subset of):

• Airflow Variables
• Airflow Connections
• configuration of Airflow in the form of Variables

They could all be bassed as generated and mounted secrets - very similarly to what was discussed in #33680 - @SamWheating - since you were thinking about it - how about extending the idea of just passing the connections to KPO, to extend it with a complete set of env variables that Airflow needs. That would enable a "Python" KPO to be able to Pull/Push XComs, and generally do the same operations that regular Python tasks can do.