Repository: PrefectHQ/prefect

Issue:
Environment vars for custom worker and scheduler specs in Dask K8s are ignored
## Description

Environment variables that collide with certain prescribed Prefect variables in a custom worker or scheduler pod spec for Dask-K8s don't have an effect.

Consider a worker pod spec like the one provided in the library:
```
    ...
    env:
      ...
      - name: PREFECT__LOGGING__LEVEL
        value: "DEBUG"
      ...
```

If you create a custom worker spec like this, the log level will get overridden and you will not be able to get debug logs.

## Expected Behavior

If you set environment variables in a pod spec, it will not get overridden.

## Reproduction

Pass the file `prefect/src/prefect/environments/execution/dask/worker_pod.yaml` as the argument `worker_spec_file` to `DaskKubernetesEnvironment` and the eventual Prefect worker will not log at the debug level.

## Explanation

I think the reason is actually pretty simple and while I don't know this exactly, it seems right.  

There are four methods in `DaskKubernetesEnvironment` that populate yaml specs.  The first two are provided for reference and if someone wants to compare how those work to the custom spec methods.

+ `_populate_job_yaml`: Populate the provided `job.yaml`
+ `_populate_worker_pod_yaml`: Populate the provided `worker_pod.yaml`
+ `_populate_scheduler_spec_yaml`: Populate a custom scheduler spec file provided by the user.
+ `_populate_worker_spec_yaml`: Populate a custom worker spec file provided by the user.

If you provide custom files, these parts here for the respective files will make sure the environment variables in the specs have all the ones needed for Prefect, just in case you forgot to provide them.
https://github.com/PrefectHQ/prefect/blob/0.13.4/src/prefect/environments/execution/dask/k8s.py#L460-L505
https://github.com/PrefectHQ/prefect/blob/0.13.4/src/prefect/environments/execution/dask/k8s.py#L526-L569

However, the logic is that `env.extend(env_values)` will extend the list of environment variables, even if there are duplicates.  Sticking with the above example with trying to set `PREFECT__LOGGING__LEVEL` to `"DEBUG"`, when looking at the output of `kubectl describe` on a pod, we'll see the variable listed twice.  The first time is our attempted specification and the second time is the one added by the Prefect code.  K8s only seems to consider the latter value.

Furthermore, since the log level populated by Prefect comes from the config and `_populate_scheduler_spec_yaml` and `_populate_worker_spec_yaml` are called in `execute` and `run`, respectively, you can't set the config locally to affect populating the file.  They get populated at the moment they are needed.  This means there's a bit of a cascade effect: we had go all the way to the agent and change `PREFECT__LOGGING__LEVEL` in order to get the log level to propagate down to the worker.

## Proposed Solution

Adjust the logic so that all Prefect default values are at the beginning of the list (if this is indeed the logic K8s uses).  Otherwise, only add the Prefect values if they are not specified in the spec file.



Discussion:
@alexifm Good call out! Yeah those places in population functions where the env vars are extended we could do a simple check to only add them if they are not already provided.
As mentioned on Slack, I have the exact same issue!

dask-kubernetes cluster running on AWS EKS with the following minimal flow example:

<details>

```python
import os
from datetime import datetime
​
import prefect
from prefect import Flow, task
from prefect.environments import DaskKubernetesEnvironment
from prefect.environments.storage import Docker
​
​
@task
def debug_logging_task():
    logger = prefect.context.get("logger")
​
    logger.debug("a debug message")
    logger.info(f"PREFECT__LOGGING__LEVEL: {os.environ.get('PREFECT__LOGGING__LEVEL')}")
​
​
with Flow("minimal_flow") as flow:
​
    debug_logging_task()
​
​
flow.environment = DaskKubernetesEnvironment(
    min_workers=1, max_workers=10, labels=["k8s"]
)
​
flow.storage = Docker(
    python_dependencies=[
        # "numpy",
        # "pandas",
        # "snowflake-connector-python[pandas]==2.3.2",
        # "snowflake-sqlalchemy>=1.2.4",
        "tqdm",
    ],
    registry_url="495775544086.dkr.ecr.eu-central-1.amazonaws.com",
    image_name="minimal_flow",
    image_tag="beta_" + datetime.now().strftime("%Y%m%d_%H%M%S"),
    env_vars={"PREFECT__LOGGING__LEVEL": "DEBUG"},
)
flow.register(project_name="eks_test_01")
```

</details>

PS: We additionally added the environment variable to the whole prefect agent when spinning it up using pulumi, but the error remains.

![image](https://user-images.githubusercontent.com/60844432/95783667-a1136380-0cd2-11eb-9a49-7c55b7fac5f6.png)

Just read your thread on the slack.  Yea, you're coming across the same problem we faced.  The solution should be pretty straightforward.  Sorry haven't had the bandwidth to submit a PR on this.
It boils down to two locations (lines following 505 and 569) in `src\prefect\environments\execution\dask\k8s.py` with the following code, right?

```python
        # set environment variables
        env = yaml_obj["spec"]["template"]["spec"]["containers"][0].get("env")
        if not env:
            yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"] = []
            env = yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"]

        env.extend(env_values)

        # set image
        yaml_obj["spec"]["template"]["spec"]["containers"][0]["image"] = docker_name

        return yaml_obj
```

with `env_values` defined as:

```python
env_values = [
            ...,
            {
                "name": "PREFECT__LOGGING__LEVEL",
                "value": str(prefect.config.logging.level),
            },
            ...,
]
```
I think your approach that sets the environment variables on the Docker image might need some other fix.  The execution environment is setting the variables in the pod specification and has no awareness of what's set on the image.  For us, if we fix how the pod spec is being handled, we'll be fine.  But I think you'll still run into the problem.

ETA: You could change how you do things and set the variable on the execution side instead of the storage side.  It would perhaps help people if the storage raised a warning that setting one of those special prefect environment variables might get superseded.
@joshmeek, does it make sense to provide the environment variables as a dictionary and write a small `dict_to_yaml` function that transforms it into a yaml file with some check whether the environment variables are already defined?


> I think your approach that sets the environment variables on the Docker image might need some other fix. The execution environment is setting the variables in the pod specification and has no awareness of what's set on the image. For us, if we fix how the pod spec is being handled, we'll be fine. But I think you'll still run into the problem.

Oh, good point!

In my opinion, if it raises a warning that describes how to set it in the execution environment, that would be a sufficient solution.

I already added the `PREFECT__LOGGING__LEVEL` environment variable to the prefect agent k8s deployment (using pulumi for infrastructure as code), is that identical with the `execution environment`?

ETA: How do you `set the variable on the execution side`?
We're both interested in affecting the environment in which the workers are running.  One approach is with an execution specification, ie the yaml spec which is what our group did, the other is via Prefect's approach to storage, ie the Docker image as you did.  Right now, as far as I can tell the only way to get an environment variable like `PREFECT__LOGGING__LEVEL` set is to set it on the agent.  Setting it any other way doesn't appear to work.

> How do you set the variable on the execution side?

This is the default worker spec that we tried to override (it mostly works, just not the environment variables)
https://github.com/PrefectHQ/prefect/blob/0.13.4/src/prefect/environments/execution/dask/worker_pod.yaml

Does that clarify what I was saying?
Thanks for the clarifications, I think I understood.

Possibilities to set `PREFECT__LOGGING__LEVEL`:
1. currently unsuccessful, via `k8s prefect agent deploy` (infrastructure specification)
2. currently unsuccessful, via `flow.storage = Docker(..., env_vars={})` (storage specification)
3. currently unsuccessful, via `worker_pod.yaml` (execution specification)

Hope I got it right? :D

ETA: 
Normally, you would not want to do it on the infrastructure specification, because you want to change it without having to change infrastructure, right? However, I tried it out of desperation... x)


Resolving patch:
diff --git a/src/prefect/environments/execution/dask/k8s.py b/src/prefect/environments/execution/dask/k8s.py
--- a/src/prefect/environments/execution/dask/k8s.py
+++ b/src/prefect/environments/execution/dask/k8s.py
@@ -45,9 +45,12 @@ class DaskKubernetesEnvironment(Environment):
     - `PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS`
     - `PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS`
     - `PREFECT__ENGINE__EXECUTOR__DEFAULT_CLASS`
+    - `PREFECT__LOGGING__LEVEL`
     - `PREFECT__LOGGING__LOG_TO_CLOUD`
     - `PREFECT__LOGGING__EXTRA_LOGGERS`
 
+    Note: the logging attributes are only populated if they are not already provided.
+
     Args:
         - min_workers (int, optional): the minimum allowed number of Dask worker pods; defaults to 1
         - max_workers (int, optional): the maximum allowed number of Dask worker pods; defaults to 1
@@ -488,6 +491,10 @@ def _populate_scheduler_spec_yaml(self, yaml_obj: dict, docker_name: str) -> dic
                 "name": "PREFECT__ENGINE__EXECUTOR__DEFAULT_CLASS",
                 "value": "prefect.engine.executors.DaskExecutor",
             },
+        ]
+
+        # Logging env vars
+        log_vars = [
             {
                 "name": "PREFECT__LOGGING__LOG_TO_CLOUD",
                 "value": str(prefect.config.logging.log_to_cloud).lower(),
@@ -510,6 +517,11 @@ def _populate_scheduler_spec_yaml(self, yaml_obj: dict, docker_name: str) -> dic
 
         env.extend(env_values)
 
+        # Append logging env vars if not already present
+        for var in log_vars:
+            if not any(d.get("name") == var.get("name") for d in env):
+                env.append(var)
+
         # set image
         yaml_obj["spec"]["template"]["spec"]["containers"][0]["image"] = docker_name
 
@@ -552,6 +564,10 @@ def _populate_worker_spec_yaml(self, yaml_obj: dict) -> dict:
                 "name": "PREFECT__ENGINE__EXECUTOR__DEFAULT_CLASS",
                 "value": "prefect.engine.executors.DaskExecutor",
             },
+        ]
+
+        # Logging env vars
+        log_vars = [
             {
                 "name": "PREFECT__LOGGING__LOG_TO_CLOUD",
                 "value": str(prefect.config.logging.log_to_cloud).lower(),
@@ -574,6 +590,11 @@ def _populate_worker_spec_yaml(self, yaml_obj: dict) -> dict:
 
         env.extend(env_values)
 
+        # Append logging env vars if not already present
+        for var in log_vars:
+            if not any(d.get("name") == var.get("name") for d in env):
+                env.append(var)
+
         # set image
         yaml_obj["spec"]["containers"][0]["image"] = prefect.context.get(
             "image", "daskdev/dask:latest"


Test patch:
diff --git a/tests/environments/execution/test_dask_k8s_environment.py b/tests/environments/execution/test_dask_k8s_environment.py
--- a/tests/environments/execution/test_dask_k8s_environment.py
+++ b/tests/environments/execution/test_dask_k8s_environment.py
@@ -478,6 +478,87 @@ def test_populate_custom_scheduler_spec_yaml(log_flag):
     )
 
 
+@pytest.mark.parametrize("log_flag", [True, False])
+def test_populate_custom_yaml_specs_with_logging_vars(log_flag):
+    environment = DaskKubernetesEnvironment()
+
+    file_path = os.path.dirname(prefect.environments.execution.dask.k8s.__file__)
+
+    log_vars = [
+        {
+            "name": "PREFECT__LOGGING__LOG_TO_CLOUD",
+            "value": "YES",
+        },
+        {
+            "name": "PREFECT__LOGGING__LEVEL",
+            "value": "NO",
+        },
+        {
+            "name": "PREFECT__LOGGING__EXTRA_LOGGERS",
+            "value": "MAYBE",
+        },
+    ]
+
+    with open(path.join(file_path, "job.yaml")) as job_file:
+        job = yaml.safe_load(job_file)
+        job["spec"]["template"]["spec"]["containers"][0]["env"] = []
+        job["spec"]["template"]["spec"]["containers"][0]["env"].extend(log_vars)
+
+    with set_temporary_config(
+        {
+            "cloud.graphql": "gql_test",
+            "cloud.auth_token": "auth_test",
+            "logging.log_to_cloud": log_flag,
+            "logging.extra_loggers": ["test_logger"],
+        }
+    ):
+        with prefect.context(flow_run_id="id_test", namespace="namespace_test"):
+            yaml_obj = environment._populate_scheduler_spec_yaml(
+                yaml_obj=job, docker_name="test1/test2:test3"
+            )
+
+    assert yaml_obj["metadata"]["name"] == "prefect-dask-job-{}".format(
+        environment.identifier_label
+    )
+
+    env = yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"]
+
+    assert env[0]["value"] == "YES"
+    assert env[1]["value"] == "NO"
+    assert env[2]["value"] == "MAYBE"
+    assert len(env) == 12
+
+    # worker
+    with open(path.join(file_path, "worker_pod.yaml")) as pod_file:
+        pod = yaml.safe_load(pod_file)
+        pod["spec"]["containers"][0]["env"] = []
+        pod["spec"]["containers"][0]["env"].extend(log_vars)
+
+    with set_temporary_config(
+        {
+            "cloud.graphql": "gql_test",
+            "cloud.auth_token": "auth_test",
+            "logging.log_to_cloud": log_flag,
+            "logging.extra_loggers": ["test_logger"],
+        }
+    ):
+        with prefect.context(flow_run_id="id_test", image="my_image"):
+            yaml_obj = environment._populate_worker_spec_yaml(yaml_obj=pod)
+
+    assert (
+        yaml_obj["metadata"]["labels"]["prefect.io/identifier"]
+        == environment.identifier_label
+    )
+    assert yaml_obj["metadata"]["labels"]["prefect.io/flow_run_id"] == "id_test"
+
+    env = yaml_obj["spec"]["containers"][0]["env"]
+
+    assert env[0]["value"] == "YES"
+    assert env[1]["value"] == "NO"
+    assert env[2]["value"] == "MAYBE"
+    assert len(env) == 10
+
+
 def test_roundtrip_cloudpickle():
     with tempfile.TemporaryDirectory() as directory: