KubernetesPodOperator "env_vars" field is documented to be templated, but it doesn't work

[barrywhart]

Apache Airflow version

main (development)

What happened

The KubernetesPodOperator env_vars field is documented to be templated, but it doesn't work. If I pass a JSON string to the constructor, it throws an exception:

  File "/opt/airflow/airflow/models/baseoperator.py", line 410, in apply_defaults
    result = func(self, **kwargs, default_args=default_args)
  File "/opt/airflow/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py", line 236, in __init__
    self.env_vars = convert_env_vars(env_vars) if env_vars else []
  File "/opt/airflow/airflow/providers/cncf/kubernetes/backcompat/backwards_compat_converters.py", line 105, in convert_env_vars
    raise AirflowException(f"Expected dict or list, got {type(env_vars)}")
airflow.exceptions.AirflowException: Expected dict or list, got <class 'str'>

What you think should happen instead

If it receives a string, the constructor should convert it to Python (probably using ast.literal_eval()).

How to reproduce

I've written a test to demonstrate the behavior. It could be added to kubernetes_tests/test_kubernetes_pod_operator.py.

    def test_env_vars_are_templatized(self):
        # WHEN
        env_vars = """
            [
                {
                    "name": "ENV1",
                    "value": "val1",
                    "value_from": null
                },
                {
                    "name": "ENV2",
                    "value": "val2",
                    "value_from": null
                },
                {
                    "name": "ENV3",
                    "value": null,
                    "value_from": {
                        "config_map_key_ref": null,
                        "field_ref": {
                            "api_version": null,
                            "field_path": "status.podIP"
                        },
                        "resource_field_ref": null,
                        "secret_key_ref": null
                    }
                }
            ]
        """
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            env_vars=env_vars,
            labels={"foo": "bar"},
            name="test-" + str(random.randint(0, 1000000)),
            task_id="task" + self.get_current_task_name(),
            in_cluster=False,
            do_xcom_push=False,
        )
        # THEN
        context = create_context(k)
        actual_pod = self.api_client.sanitize_for_serialization(k.build_pod_request_obj(context))
        self.expected_pod['spec']['containers'][0]['env'] = [
            {'name': 'ENV1', 'value': 'val1'},
            {'name': 'ENV2', 'value': 'val2'},
            {'name': 'ENV3', 'valueFrom': {'fieldRef': {'fieldPath': 'status.podIP'}}},
        ]
        assert self.expected_pod == actual_pod

Operating System

Debian GNU/Linux 11 (bullseye) (I'm running the Breeze dev environment)

Versions of Apache Airflow Providers

I'm using the Breeze development environment with the latest main code as of August 19, 2022:

commit 0254f30a5a90f0c3104782525fabdcfdc6d3b7df (HEAD -> main, origin/main, origin/HEAD)
Author: Ash Berlin-Taylor <[email protected]>
Date:   Fri Aug 19 14:15:43 2022 +0100

    Let timetables control generated run_ids. (#25795)

Deployment

Other

Deployment details

Breeze development environment

Anything else

The problem happens every time.

I would like to submit a PR to address this issue. I would also like (if it sounds reasonable -- please let me know) to add similar templating support for these other fields:

• resources
• tolerations

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[potiuk]

I thing you misunderstood what "templated" means. You do not have to use string to make templates work.

If you look at env_vars type definition, you will not find "string":

 env_vars: Optional[List[k8s.V1EnvVar]] = None,

What you should provide is list of k8s.V1EnvVar and template names/values in it (those are strings).

Would that work for you?

[barrywhart]

This will not support our needs. We'd like KubernetesPodOperator to support more flexible templating:

• Dynamic number of environment variables provided via DAG JSON configuration (dag_run.conf)
• If provided, propagate other environment variable(s) from one Kubernetes task to another.

To provide more context, this is how our MLOPs team currently uses env_vars templating using some custom KubernetesPodOperator extensions we've developed.

Our MLOps team needs very flexible DAGs that support arbitrary ML code written by data scientists. The data scientists work on ML mode but do not touch Airflow code. Hence, there is one DAG called mlops_experiment with the following structure:

On each PR commit in any ML repo, the CI job runs this DAG, with the pod images and env vars defined by the repo code. To support this, the mlops_experiment "train" task defines env_vars as:

env_vars='dag_run.conf["train"]'

This simple Jinja string allows data scientists to define an arbitrary number of environment variables. We also need to track some metrics data across the "train" and "predict_batch" tasks, which requires passing a "run ID" generated by the "train" pod to the "predict_batch" pod. Thus, the "predict_batch" task defines env_vars as:

env_vars='dag_run.conf["predict_batch"] | propagate_run_id(task_instance)'

We believe other Airflow users may have the same needs, and we'd like to contribute this enhancement to the Airflow community.

[potiuk]

@jedcunningham @dimberman - let me know what you think. I think narrowing this down and solving just env_vars is not reallly a good solution - because there are a number of other parameters that only take k8s objects.

I kinda suspected what @barrywhart will come up with - the variable number vars was the first thing that came to my mind when I thought why this might be limiting,

But I think adding a "string" here might not be the best idea to solve it (serialization/deserialization issues?) and there is a good reason why we limited the parameters to only accept objects and not strings (but maybe I am wrong?).

Of course this dynamic nujmber of env vars could be workarounded by providing big number of variables and indexing them ({{var0}}, {{val0}}, {{var1}}{{val2}} .....). Then @barrywhart could do what they want (but it is rather nasty workaround).

But maybe we could do something else here? Possibly generate list of env var objects someohow?

WDYT?

[potiuk]

Also @dstandish ^^

[potiuk]

No response yet from Jed and Daniel, but I think you should create a feature request for that one @barrywhart.

[barrywhart]

I had a PR, but the review wasn't getting prompt, helpful feedback, so I closed it. We're using the proposed solution in a custom subclass instead. Thanks!