Repository: dask/distributed

Issue:
Unsetting `MALLOC_TRIM_THRESHOLD_` is sensitive to import order
`import distributed` respects any previous calls to `dask.config.set` that *change a value*, but not those that *remove a key*:

```python
>>> import dask
>>> dask.config.set({"distributed.nanny.pre-spawn-environ": {'OMP_NUM_THREADS': 2}})
>>> dask.config.get("distributed.nanny.pre-spawn-environ")
{'OMP_NUM_THREADS': 2}
>>> import distributed
>>> dask.config.get("distributed.nanny.pre-spawn-environ")
{'OMP_NUM_THREADS': 2, 'MALLOC_TRIM_THRESHOLD_': 65536, 'MKL_NUM_THREADS': 1, 'OPENBLAS_NUM_THREADS': 1}
```

The problem here is that an unset `MALLOC_TRIM_THRESHOLD_` has a special meaning for the glibc so it can't just simply be overridden with a different explicit value.


Discussion:
For Coiled, it would be very helpful if the solution also worked for env vars and yaml, not just the in memory config dict.
> For Coiled, it would be very helpful if the solution also worked for env vars and yaml, not just the in memory config dict.

Yes, with the linked PR you can set:
```bash
$ export DASK_DISTRIBUTED__NANNY__PRE_SPAWN_ENVIRON__MALLOC_TRIM_THRESHOLD_=None
```
(case sensitive)

or you can use a yaml file:
```yaml
distributed:
  nanny:
    pre-spawn-environ:
      MALLOC_TRIM_THRESHOLD_: null
```

For clarity, the above should *not* be the default for coiled.

Resolving patch:
diff --git a/distributed/nanny.py b/distributed/nanny.py
--- a/distributed/nanny.py
+++ b/distributed/nanny.py
@@ -1008,5 +1008,6 @@ def _get_env_variables(config_key: str) -> dict[str, str]:
             f"{config_key} configuration must be of type dict. Instead got {type(cfg)}"
         )
     # Override dask config with explicitly defined env variables from the OS
-    cfg = {k: os.environ.get(k, str(v)) for k, v in cfg.items()}
+    # Allow unsetting a variable in a config override by setting its value to None.
+    cfg = {k: os.environ.get(k, str(v)) for k, v in cfg.items() if v is not None}
     return cfg


Test patch:
diff --git a/distributed/tests/test_nanny.py b/distributed/tests/test_nanny.py
--- a/distributed/tests/test_nanny.py
+++ b/distributed/tests/test_nanny.py
@@ -349,6 +349,40 @@ async def test_environment_variable_pre_post_spawn(c, s, n):
     assert "POST-SPAWN" not in os.environ
 
 
+@gen_cluster(
+    nthreads=[],
+    client=True,
+    config={
+        "distributed.nanny.pre-spawn-environ.PRE1": 1,
+        "distributed.nanny.pre-spawn-environ.PRE2": 2,
+        "distributed.nanny.pre-spawn-environ.PRE3": 3,
+        "distributed.nanny.environ.POST1": 4,
+        "distributed.nanny.environ.POST2": 5,
+        "distributed.nanny.environ.POST3": 6,
+    },
+)
+async def test_environment_variable_overlay(c, s):
+    """You can set a value to None to unset a variable in a config overlay"""
+    # Not the same as running Nanny(config=...), which would not work for pre-spawn
+    # variables
+    with dask.config.set(
+        {
+            "distributed.nanny.pre-spawn-environ.PRE2": 7,
+            "distributed.nanny.pre-spawn-environ.PRE3": None,
+            "distributed.nanny.environ.POST2": 8,
+            "distributed.nanny.environ.POST3": None,
+        },
+    ):
+        async with Nanny(s.address):
+            env = await c.submit(lambda: os.environ)
+            assert env["PRE1"] == "1"
+            assert env["PRE2"] == "7"
+            assert "PRE3" not in env
+            assert env["POST1"] == "4"
+            assert env["POST2"] == "8"
+            assert "POST3" not in env
+
+
 @gen_cluster(client=True, nthreads=[])
 async def test_config_param_overlays(c, s):
     with dask.config.set({"test123.foo": 1, "test123.bar": 2}):