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.

[alexifm]

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.

[alexifm]

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.