Repository: dask/distributed

Issue:
Performance regression in cluster memory usage
Recently, some of us have started tracking some performance metrics for Dask clusters under a variety of different usage patterns. The idea is to be able to identify performance regressions before they are released (especially ones at scale which might not show up in a unit test context).

An example of these metrics is at [this](https://coiled.github.io/coiled-runtime/coiled-upstream-py3.9.html) static site. It's only been collecting results for a few days, but already we seem to have come across a significant regression in cluster memory usage.  Here is a test which measures array rechunking:
```python
def test_rechunk_in_memory(small_client):
    x = da.random.random((50000, 50000))
    x.rechunk((50000, 20)).rechunk((20, 50000)).sum().compute()
```
and a screenshot of average cluster memory usage for that operation over the last week+ :
![image](https://user-images.githubusercontent.com/5728311/182952487-345018c3-c059-4031-880b-8078617a1a17.png)
(I encourage folks to click through, this same behavior appears on a lot of tests around July 26)

The above is based on a coiled cluster, but I've reproduced it using a `LocalCluster` with the following procedure:

1. Create a software environment with nightly dask versions from the `dask` conda channel:
    ```bash
    conda create -n memory-regression python=3.9 dask distributed numpy
    conda activate memory-regression
    # Install nightly from July 22nd
    conda install https://conda.anaconda.org/dask/label/dev/noarch/dask-2022.7.1a220722-py_ga55bfd36_21.tar.bz2 https://conda.anaconda.org/dask/label/dev/noarch/distributed-2022.7.1a220722-py_ga55bfd36_21.tar.bz2
   # Or install nightly from July 25th
   conda install https://conda.anaconda.org/dask/label/dev/noarch/dask-2022.7.2a220725-py_g55cc1a50_1.tar.bz2 https://conda.anaconda.org/dask/label/dev/noarch/distributed-2022.7.2a220725-py_g55cc1a50_1.tar.bz2
    ```
2. Run the following script
```python
import ctypes
import uuid

import dask.array as da
import distributed


sampler = distributed.diagnostics.MemorySampler()


def trim_memory() -> int:
    libc = ctypes.CDLL("libc.so.6")
    return libc.malloc_trim(0)


if __name__ == "__main__":
    print(distributed.__version__)
    client = distributed.Client()
    mems = []
    for i in range(20):
        label = str(uuid.uuid4())

        with sampler.sample(label=label, client=client, measure="process"):
            x = da.random.random((20000, 20000))
            x.rechunk((20000, 20)).rechunk((20, 20000)).sum().compute()

        df = sampler.to_pandas()
        mems.append(df[label].mean())
        client.run(trim_memory)
        client.restart()

    print(mems)
```

This produces results like the following:

![image](https://user-images.githubusercontent.com/5728311/182953643-010a5875-14bc-448c-9195-1e4c96121c53.png)

~~Timing-wise, this suggests to me that #6728 might have had some unintended side-effects in cluster memory usage, but I have not verified, nor do I know how it could be so drastic.~~
Edit, see below


Discussion:
Oh, right, it's clearly https://github.com/dask/distributed/pull/6777.

Nice to see a consistent story, I guess (I did a bisect to confirm)
Indeed, https://github.com/dask/distributed/pull/6777 is a pretty clear culprit!

Nice fire drill of the benchmarking. Clearly it works! Now we just need alerts.

Resolving patch:
diff --git a/distributed/nanny.py b/distributed/nanny.py
--- a/distributed/nanny.py
+++ b/distributed/nanny.py
@@ -17,6 +17,7 @@
 from time import sleep as sync_sleep
 from typing import TYPE_CHECKING, ClassVar
 
+from toolz import merge
 from tornado import gen
 from tornado.ioloop import IOLoop
 
@@ -82,12 +83,27 @@ class Nanny(ServerNode):
         ensured to be set in the Worker process as well. This argument allows to
         overwrite or otherwise set environment variables for the Worker. It is
         also possible to set environment variables using the option
-        `distributed.nanny.environ`. Precedence as follows
+        ``distributed.nanny.environ``. Precedence as follows
 
             1. Nanny arguments
             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. Such variables
+        would be ineffective if set here or in ``distributed.nanny.environ``; they
+        must be set in ``distributed.nanny.pre-spawn-environ`` so that they are set
+        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
@@ -97,6 +113,9 @@ class Nanny(ServerNode):
     process = None
     memory_manager: NannyMemoryManager
 
+    env: dict[str, str]
+    pre_spawn_env: dict[str, str]
+
     # Inputs to parse_ports()
     _given_worker_port: int | str | Collection[int] | None
     _start_port: int | str | Collection[int] | None
@@ -204,19 +223,14 @@ def __init__(  # type: ignore[no-untyped-def]
         self.death_timeout = parse_timedelta(death_timeout)
 
         self.Worker = Worker if worker_class is None else worker_class
-        config_environ = dask.config.get("distributed.nanny.environ", {})
-        if not isinstance(config_environ, dict):
-            raise TypeError(
-                "distributed.nanny.environ configuration must be of type dict. "
-                f"Instead got {type(config_environ)}"
-            )
-        self.env = config_environ.copy()
-        for k in self.env:
-            if k in os.environ:
-                self.env[k] = os.environ[k]
-        if env:
-            self.env.update(env)
-        self.env = {k: str(v) for k, v in self.env.items()}
+
+        self.pre_spawn_env = _get_env_variables("distributed.nanny.pre-spawn-environ")
+        self.env = merge(
+            self.pre_spawn_env,
+            _get_env_variables("distributed.nanny.environ"),
+            {k: str(v) for k, v in env.items()} if env else {},
+        )
+
         self.config = config or dask.config.config
         worker_kwargs.update(
             {
@@ -401,6 +415,7 @@ async def instantiate(self) -> Status:
                 on_exit=self._on_worker_exit_sync,
                 worker=self.Worker,
                 env=self.env,
+                pre_spawn_env=self.pre_spawn_env,
                 config=self.config,
             )
 
@@ -600,6 +615,9 @@ class WorkerProcess:
     running: asyncio.Event
     stopped: asyncio.Event
 
+    env: dict[str, str]
+    pre_spawn_env: dict[str, str]
+
     # The interval how often to check the msg queue for init
     _init_msg_interval = 0.05
 
@@ -610,6 +628,7 @@ def __init__(
         on_exit,
         worker,
         env,
+        pre_spawn_env,
         config,
     ):
         self.status = Status.init
@@ -619,6 +638,7 @@ def __init__(
         self.process = None
         self.Worker = worker
         self.env = env
+        self.pre_spawn_env = pre_spawn_env
         self.config = config
 
         # Initialized when worker is ready
@@ -660,6 +680,10 @@ async def start(self) -> Status:
         self.stopped = asyncio.Event()
         self.status = Status.starting
 
+        # Set selected environment variables before spawning the subprocess.
+        # See note in Nanny docstring.
+        os.environ.update(self.pre_spawn_env)
+
         try:
             await self.process.start()
         except OSError:
@@ -922,3 +946,14 @@ async def run():
                     child_stop_q.close()  # usually redundant
                 child_stop_q.join_thread()
                 thread.join(timeout=2)
+
+
+def _get_env_variables(config_key: str) -> dict[str, str]:
+    cfg = dask.config.get(config_key)
+    if not isinstance(cfg, dict):
+        raise TypeError(  # pragma: nocover
+            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()}
+    return cfg


Test patch:
diff --git a/distributed/tests/test_config.py b/distributed/tests/test_config.py
--- a/distributed/tests/test_config.py
+++ b/distributed/tests/test_config.py
@@ -337,7 +337,12 @@ def test_schema_is_complete():
     with open(schema_fn) as f:
         schema = yaml.safe_load(f)
 
-    skip = {"default-task-durations", "bokeh-application", "environ"}
+    skip = {
+        "default-task-durations",
+        "bokeh-application",
+        "environ",
+        "pre-spawn-environ",
+    }
 
     def test_matches(c, s):
         if set(c) != set(s["properties"]):
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
@@ -330,6 +330,25 @@ async def test_environment_variable_config(c, s, monkeypatch):
         assert results[n.worker_address]["D"] == "123"
 
 
+@gen_cluster(
+    nthreads=[("", 1)],
+    client=True,
+    Worker=Nanny,
+    config={
+        "distributed.nanny.pre-spawn-environ": {"PRE-SPAWN": 1},
+        "distributed.nanny.environ": {"POST-SPAWN": 2},
+    },
+)
+async def test_environment_variable_pre_post_spawn(c, s, n):
+    assert n.env == {"PRE-SPAWN": "1", "POST-SPAWN": "2"}
+    results = await c.run(lambda: os.environ)
+    assert results[n.worker_address]["PRE-SPAWN"] == "1"
+    assert results[n.worker_address]["POST-SPAWN"] == "2"
+
+    del os.environ["PRE-SPAWN"]
+    assert "POST-SPAWN" not in os.environ
+
+
 @gen_cluster(nthreads=[])
 async def test_local_directory(s):
     with tmpfile() as fn:
@@ -633,3 +652,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.pre-spawn-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)