Repository: dask/distributed

Issue:
distributed.nanny.environ.MALLOC_TRIM_THRESHOLD_ is ineffective
Ubuntu 21.10 x86/64
distributed 2022.3.0

The MALLOC_TRIM_THRESHOLD_ env variable seems to be effective at making memory deallocation more reactive.
However, the config variable that sets it doesn't seem to do anything - which indicates that the variable is being set after the worker process is started, whereas it should be set before spawning it.

```python
import dask.array
import distributed

client = distributed.Client(n_workers=1, memory_limit="2 GiB")

N = 7_000
S = 160 * 1024

a = dask.array.random.random(N * S // 8, chunks=S // 8)
a = a.persist()
distributed.wait(a)
del a
```
Result:
Managed: 0
Unmanaged: 1.16 GiB

```python
import os
import dask.array
import dask.config
import distributed

os.environ["MALLOC_TRIM_THRESHOLD_"] = str(dask.config.get("distributed.nanny.environ.MALLOC_TRIM_THRESHOLD_"))
client = distributed.Client(n_workers=1, memory_limit="2 GiB")

N = 7_000
S = 160 * 1024

a = dask.array.random.random(N * S // 8, chunks=S // 8)
a = a.persist()
distributed.wait(a)
del a
```
Result:
Managed: 0
Unmanaged: 151 MiB

# Production Workaround
Set the env variable on the shell, before starting ``dask-worker``:

```bash
export MALLOC_TRIM_THRESHOLD_=65536
dask-worker <address>
```


Discussion:
I'm uncertain about how to solve this. The simple solution, to change ``os.environ`` in ``Nanny.__init__`` instead of passing them down to Worker, would also mean poisoning the whole process of the nanny. This is annoying for unit tests but I'm not sure if anybody cares in production?

The alternative is to have an intermediate process which sets the variables and then invokes python again, which however is very expensive.

This issue also impacts the other two variables set by the config:

      OMP_NUM_THREADS: 1
      MKL_NUM_THREADS: 1

AFAIK, if for any reason numpy is imported on the worker process before the config, these two variables will not be picked up.
> if for any reason numpy is imported on the worker process before the config

Quite possible: https://github.com/dask/distributed/issues/5729

> The simple solution, to change `os.environ` in `Nanny.__init__` instead of passing them down to Worker, would also mean poisoning the whole process of the nanny

To be fair, all the environment variables we're currently talking about setting (malloc_trim and num_threads) basically don't have an impact unless they're set before the interpreter starts. So for these specifically, setting them in the Nanny shouldn't actually change anything in practice. I still dislike the uncleanliness of setting them in the Nanny process, though.

> The alternative is to have an intermediate process which sets the variables and then invokes python again

Could also have a process-wide lock for `Nanny`, and set `os.environ` in the parent process while holding that lock, then reset it before releasing. You'd want to be clever and still allow subprocesses with the same `env` to be spawned in parallel, so `dask-worker --nworkers=100` is still performant. Still not perfectly clean, but maybe an acceptable tradeoff between cleanliness and performance.
> setting them in the Nanny shouldn't actually change anything in practice.

They'll be set before the worker process starts. the worker process is where it matters
I meant having them set in the Nanny process won't really affect things for the Nanny itself. Guido and I don't like the poor hygiene of leaving them set on the Nanny, but I'm just noting that that poor hygiene shouldn't affect anything on the Nanny in practice because the particular variables we're setting only have an effect at interpreter startup/NumPy import time.
I was worried about potential user-defined variables, not the three we set. But I'm leaning towards not over-engineering this just to cover purely hypothetical use cases.

Resolving patch:
diff --git a/distributed/nanny.py b/distributed/nanny.py
--- a/distributed/nanny.py
+++ b/distributed/nanny.py
@@ -87,6 +87,19 @@ class Nanny(ServerNode):
             2. Existing environment variables
             3. Dask configuration
 
+        Note
+        ----
+        Some environment variables, like ``OMP_NUM_THREADS``, must be set before
+        importing numpy to have effect. Others, like ``MALLOC_TRIM_THRESHOLD_`` (see
+        :ref:`memtrim`), must be set before starting the Linux process. So we need to
+        set them before spawning the subprocess, even if this means poisoning the
+        process running the Nanny.
+
+        For the same reason, be warned that changing
+        ``distributed.worker.multiprocessing-method`` from ``spawn`` to ``fork`` or
+        ``forkserver`` may inhibit some environment variables; if you do, you should
+        set the variables yourself in the shell before you start ``dask-worker``.
+
     See Also
     --------
     Worker
@@ -666,6 +679,10 @@ async def start(self) -> Status:
         self.stopped = asyncio.Event()
         self.status = Status.starting
 
+        # Must set env variables before spawning the subprocess.
+        # See note in Nanny docstring.
+        os.environ.update(self.env)
+
         try:
             await self.process.start()
         except OSError:
@@ -811,6 +828,9 @@ def _run(
         Worker,
     ):  # pragma: no cover
         try:
+            # Set the environment variables again. This is to avoid race conditions
+            # where different nannies in the same process set different environment
+            # variables.
             os.environ.update(env)
             dask.config.set(config)
 


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
@@ -8,7 +8,6 @@
 import random
 import sys
 from contextlib import suppress
-from unittest import mock
 
 import psutil
 import pytest
@@ -275,42 +274,50 @@ async def test_wait_for_scheduler():
     assert "restart" not in log.lower(), log
 
 
-@gen_cluster(nthreads=[], client=True)
+@gen_cluster(client=True, nthreads=[])
 async def test_environment_variable(c, s):
-    a = Nanny(s.address, memory_limit=0, env={"FOO": "123"})
-    b = Nanny(s.address, memory_limit=0, env={"FOO": "456"})
-    await asyncio.gather(a, b)
-    results = await c.run(lambda: os.environ["FOO"])
-    assert results == {a.worker_address: "123", b.worker_address: "456"}
-    await asyncio.gather(a.close(), b.close())
-
-
-@gen_cluster(nthreads=[], client=True)
-async def test_environment_variable_by_config(c, s, monkeypatch):
-
     with dask.config.set({"distributed.nanny.environ": "456"}):
         with pytest.raises(TypeError, match="configuration must be of type dict"):
-            Nanny(s.address, memory_limit=0)
-
-    with dask.config.set({"distributed.nanny.environ": {"FOO": "456"}}):
-
-        # precedence
-        # kwargs > env var > config
+            Nanny(s.address)
+
+    try:
+        # precedence: kwargs > env var > config
+        with dask.config.set({"distributed.nanny.environ": {"FOO": "123"}}):
+            os.environ["FOO"] = "456"
+            async with Nanny(s.address, env={"FOO": "789"}) as n:
+                results = await c.run(lambda: os.environ["FOO"])
+                assert results == {n.worker_address: "789"}
+
+            # Nanny sets the env variables in the Nanny's own process.
+            # Need to manually clean up os.environ afterwards.
+            os.environ["FOO"] = "456"
+            async with Nanny(s.address) as n:
+                results = await c.run(lambda: os.environ["FOO"])
+                assert results == {n.worker_address: "456"}
 
-        with mock.patch.dict(os.environ, {"FOO": "BAR"}, clear=True):
-            a = Nanny(s.address, memory_limit=0, env={"FOO": "123"})
-            x = Nanny(s.address, memory_limit=0)
+            del os.environ["FOO"]
+            async with Nanny(s.address) as n:
+                results = await c.run(lambda: os.environ["FOO"])
+                assert results == {n.worker_address: "123"}
+    finally:
+        del os.environ["FOO"]
 
-        b = Nanny(s.address, memory_limit=0)
 
-        await asyncio.gather(a, b, x)
+@gen_cluster(client=True, nthreads=[])
+async def test_environment_variable_race_condition(c, s):
+    """Nannies need to set the environment variables in their own process so that
+    special variables such as MALLOC_TRIM_THRESHOLD_ are picked up when the Worker's
+    interpreter is started. However, this creates a race condition with multiple Nannies
+    in the same process with different variables. Test that the environment variables
+    are set again inside their own process."""
+    nannies = [Nanny(s.address, env={"FOO": str(i)}) for i in range(8)]
+    try:
+        await asyncio.gather(*nannies)
         results = await c.run(lambda: os.environ["FOO"])
-        assert results == {
-            a.worker_address: "123",
-            b.worker_address: "456",
-            x.worker_address: "BAR",
-        }
-        await asyncio.gather(a.close(), b.close(), x.close())
+        assert sorted(results.values()) == [str(i) for i in range(8)]
+        await asyncio.gather(*(n.close() for n in nannies))
+    finally:
+        del os.environ["FOO"]
 
 
 @gen_cluster(
@@ -585,3 +592,53 @@ async def test_scheduler_crash_doesnt_restart(s, a):
     await a.finished()
     assert a.status == Status.closed
     assert a.process is None
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(not LINUX, reason="Requires GNU libc")
+@gen_cluster(
+    client=True,
+    Worker=Nanny,
+    nthreads=[("", 2)],
+    worker_kwargs={"memory_limit": "2GiB"},
+)
+async def test_malloc_trim_threshold(c, s, a):
+    """Test that the nanny sets the MALLOC_TRIM_THRESHOLD_ environment variable before
+    starting the worker process.
+
+    This test relies on these settings to work:
+
+        distributed.nanny.environ.MALLOC_TRIM_THRESHOLD_: 65536
+        distributed.worker.multiprocessing-method: spawn
+
+    We're deliberately not setting them explicitly in @gen_cluster above, as we want
+    this test to trip if somebody changes distributed.yaml.
+
+    Note
+    ----
+    This test may start failing in a future Python version if CPython switches to
+    using mimalloc by default. If it does, a thorough benchmarking exercise is needed.
+    """
+    da = pytest.importorskip("dask.array")
+
+    a = da.random.random(
+        2**29 // 8,  # 0.5 GiB,
+        chunks=160 * 2**10 // 8,  # 160 kiB
+    ).persist()
+    await wait(a)
+    # Wait for heartbeat
+    while s.memory.process < 2**29:
+        await asyncio.sleep(0.01)
+    del a
+
+    # This is the delicate bit, as it relies on
+    # 1. PyMem_Free() to be quick to invoke glibc free() when memory becomes available
+    # 2. glibc free() to be quick to invoke the kernel's sbrk() when the same happens
+    #
+    # At the moment of writing, the readings are:
+    # - 122 MiB after starting a new worker
+    # - 139 MiB after computing a trivial dask.array collection
+    # - 185 MiB at the end of this test, with MALLOC_TRIM_THRESHOLD_=65536
+    # - 698 MiB at the end of this test, without MALLOC_TRIM_THRESHOLD_
+    while s.memory.process > 250 * 2**20:
+        await asyncio.sleep(0.01)