Repository: dask/distributed

Issue:
Nanny Plugins
We currently have Worker plugins and Scheduler plugins.  I think that there are a few situations where having a Nanny plugin would be useful.

### Motivation

Nannies can run code before the process starts up.  This makes them helpful in owning the environment of a dask worker.  In particular, the following become easy to manage.

1. Environment variables 
2. Software environment, including uploaded scripts, directories, or even conda environments

### Implementation

A lot of this is already done in the worker an can be easily copied over.  For example on the scheduler side we can copy-paste this and add the `nanny=True` flag to broadcast

https://github.com/dask/distributed/blob/117d27781e05149c0d745376641ed9ca33269245/distributed/scheduler.py#L6922-L6929

On the worker side if one searchs for `plugin` we find calls like `plugin_add` and `plugin_teardown` both as methods of the worker that the scheduler calls, and within initialization.  On the nanny we probably only care about when the scheduler calls us.  

We also care about possibly restarting the worker when we get a new plugin (this should be user-configurable).  Situations like setting environment variables and software environments will often reqiure a restart for currently running workers, although they won't require a restart if we're just creating this nanny in the first place.

Discussion:
mrocklin: I tried the UploadDirectory plugin in practice and it works great for existing workers, but for some reason new workers that are added after the plugin is registered aren't picking up the change.  

Patch:
diff --git a/distributed/__init__.py b/distributed/__init__.py
index 92a62a30..9f7a8d6f 100644
--- a/distributed/__init__.py
+++ b/distributed/__init__.py
@@ -21,7 +21,15 @@ from .client import (
 )
 from .core import Status, connect, rpc
 from .deploy import Adaptive, LocalCluster, SpecCluster, SSHCluster
-from .diagnostics.plugin import PipInstall, SchedulerPlugin, WorkerPlugin
+from .diagnostics.plugin import (
+    Environ,
+    NannyPlugin,
+    PipInstall,
+    SchedulerPlugin,
+    UploadDirectory,
+    UploadFile,
+    WorkerPlugin,
+)
 from .diagnostics.progressbar import progress
 from .event import Event
 from .lock import Lock
diff --git a/distributed/client.py b/distributed/client.py
index b93b67f3..75937435 100644
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -59,7 +59,12 @@ from .core import (
     connect,
     rpc,
 )
-from .diagnostics.plugin import UploadFile, WorkerPlugin, _get_worker_plugin_name
+from .diagnostics.plugin import (
+    NannyPlugin,
+    UploadFile,
+    WorkerPlugin,
+    _get_worker_plugin_name,
+)
 from .metrics import time
 from .objects import HasWhat, SchedulerInfo, WhoHas
 from .protocol import to_serialize
@@ -4089,10 +4094,13 @@ class Client:
         """
         return self.register_worker_plugin(_WorkerSetupPlugin(setup))
 
-    async def _register_worker_plugin(self, plugin=None, name=None):
-        responses = await self.scheduler.register_worker_plugin(
-            plugin=dumps(plugin, protocol=4), name=name
-        )
+    async def _register_worker_plugin(self, plugin=None, name=None, nanny=None):
+        if nanny or nanny is None and isinstance(plugin, NannyPlugin):
+            method = self.scheduler.register_nanny_plugin
+        else:
+            method = self.scheduler.register_worker_plugin
+
+        responses = await method(plugin=dumps(plugin, protocol=4), name=name)
         for response in responses.values():
             if response["status"] == "error":
                 exc = response["exception"]
@@ -4100,7 +4108,7 @@ class Client:
                 raise exc.with_traceback(tb)
         return responses
 
-    def register_worker_plugin(self, plugin=None, name=None, **kwargs):
+    def register_worker_plugin(self, plugin=None, name=None, nanny=None, **kwargs):
         """
         Registers a lifecycle worker plugin for all current and future workers.
 
@@ -4124,12 +4132,14 @@ class Client:
 
         Parameters
         ----------
-        plugin : WorkerPlugin
-            The plugin object to pass to the workers
+        plugin : WorkerPlugin or NannyPlugin
+            The plugin object to register.
         name : str, optional
             A name for the plugin.
             Registering a plugin with the same name will have no effect.
             If plugin has no name attribute a random name is used.
+        nanny : bool, optional
+            Whether to register the plugin with workers or nannies.
         **kwargs : optional
             If you pass a class as the plugin, instead of a class instance, then the
             class will be instantiated with any extra keyword arguments.
@@ -4174,10 +4184,15 @@ class Client:
 
         assert name
 
-        return self.sync(self._register_worker_plugin, plugin=plugin, name=name)
+        return self.sync(
+            self._register_worker_plugin, plugin=plugin, name=name, nanny=nanny
+        )
 
-    async def _unregister_worker_plugin(self, name):
-        responses = await self.scheduler.unregister_worker_plugin(name=name)
+    async def _unregister_worker_plugin(self, name, nanny=None):
+        if nanny:
+            responses = await self.scheduler.unregister_nanny_plugin(name=name)
+        else:
+            responses = await self.scheduler.unregister_worker_plugin(name=name)
 
         for response in responses.values():
             if response["status"] == "error":
@@ -4186,7 +4201,7 @@ class Client:
                 raise exc.with_traceback(tb)
         return responses
 
-    def unregister_worker_plugin(self, name):
+    def unregister_worker_plugin(self, name, nanny=None):
         """Unregisters a lifecycle worker plugin
 
         This unregisters an existing worker plugin. As part of the unregistration process
@@ -4220,7 +4235,7 @@ class Client:
         --------
         register_worker_plugin
         """
-        return self.sync(self._unregister_worker_plugin, name=name)
+        return self.sync(self._unregister_worker_plugin, name=name, nanny=nanny)
 
 
 class _WorkerSetupPlugin(WorkerPlugin):
diff --git a/distributed/comm/tcp.py b/distributed/comm/tcp.py
index de6cee01..0088ed65 100644
--- a/distributed/comm/tcp.py
+++ b/distributed/comm/tcp.py
@@ -1,3 +1,4 @@
+import ctypes
 import errno
 import functools
 import logging
@@ -14,6 +15,7 @@ try:
 except ImportError:
     ssl = None
 
+from tlz import sliding_window
 from tornado import netutil
 from tornado.iostream import StreamClosedError
 from tornado.tcpclient import TCPClient
@@ -34,6 +36,7 @@ from .utils import ensure_concrete_host, from_frames, get_tcp_server_address, to
 logger = logging.getLogger(__name__)
 
 
+C_INT_MAX = 256 ** ctypes.sizeof(ctypes.c_int) // 2 - 1
 MAX_BUFFER_SIZE = MEMORY_LIMIT / 2
 
 
@@ -193,9 +196,15 @@ class TCP(Comm):
             frames_nbytes = await stream.read_bytes(fmt_size)
             (frames_nbytes,) = struct.unpack(fmt, frames_nbytes)
 
-            frames = bytearray(frames_nbytes)
-            n = await stream.read_into(frames)
-            assert n == frames_nbytes, (n, frames_nbytes)
+            frames = memoryview(bytearray(frames_nbytes))
+            # Workaround for OpenSSL 1.0.2 (can drop with OpenSSL 1.1.1)
+            for i, j in sliding_window(
+                2, range(0, frames_nbytes + C_INT_MAX, C_INT_MAX)
+            ):
+                chunk = frames[i:j]
+                chunk_nbytes = len(chunk)
+                n = await stream.read_into(chunk)
+                assert n == chunk_nbytes, (n, chunk_nbytes)
         except StreamClosedError as e:
             self.stream = None
             self._closed = True
diff --git a/distributed/comm/ucx.py b/distributed/comm/ucx.py
index 457a5c96..0e048c9f 100644
--- a/distributed/comm/ucx.py
+++ b/distributed/comm/ucx.py
@@ -108,7 +108,7 @@ def init_once():
                     "In order to send/recv CUDA arrays, Numba or RMM is required"
                 )
 
-    pool_size_str = dask.config.get("rmm.pool-size")
+    pool_size_str = dask.config.get("distributed.rmm.pool-size")
     if pool_size_str is not None:
         pool_size = parse_bytes(pool_size_str)
         rmm.reinitialize(
@@ -121,7 +121,7 @@ def init_once():
         ucx_create_endpoint = ucp.create_endpoint
         ucx_create_listener = ucp.create_listener
     else:
-        reuse_endpoints = dask.config.get("ucx.reuse-endpoints")
+        reuse_endpoints = dask.config.get("distributed.comm.ucx.reuse-endpoints")
         if (
             reuse_endpoints is None and ucp.get_ucx_version() >= (1, 11, 0)
         ) or reuse_endpoints is False:
@@ -491,12 +491,12 @@ def _scrub_ucx_config():
     # leave UCX to its default configuration
     if any(
         [
-            dask.config.get("ucx.tcp"),
-            dask.config.get("ucx.nvlink"),
-            dask.config.get("ucx.infiniband"),
+            dask.config.get("distributed.comm.ucx.tcp"),
+            dask.config.get("distributed.comm.ucx.nvlink"),
+            dask.config.get("distributed.comm.ucx.infiniband"),
         ]
     ):
-        if dask.config.get("ucx.rdmacm"):
+        if dask.config.get("distributed.comm.ucx.rdmacm"):
             tls = "tcp" if ucx_110 else "tcp,rdmacm"
             tls_priority = "rdmacm"
         else:
@@ -506,17 +506,22 @@ def _scrub_ucx_config():
         # CUDA COPY can optionally be used with ucx -- we rely on the user
         # to define when messages will include CUDA objects.  Note:
         # defining only the Infiniband flag will not enable cuda_copy
-        if any([dask.config.get("ucx.nvlink"), dask.config.get("ucx.cuda_copy")]):
+        if any(
+            [
+                dask.config.get("distributed.comm.ucx.nvlink"),
+                dask.config.get("distributed.comm.ucx.cuda_copy"),
+            ]
+        ):
             tls = tls + ",cuda_copy"
 
-        if dask.config.get("ucx.infiniband"):
+        if dask.config.get("distributed.comm.ucx.infiniband"):
             tls = "rc," + tls
-        if dask.config.get("ucx.nvlink"):
+        if dask.config.get("distributed.comm.ucx.nvlink"):
             tls = tls + ",cuda_ipc"
 
         options = {"TLS": tls, "SOCKADDR_TLS_PRIORITY": tls_priority}
 
-        net_devices = dask.config.get("ucx.net-devices")
+        net_devices = dask.config.get("distributed.comm.ucx.net-devices")
         if net_devices is not None and net_devices != "":
             options["NET_DEVICES"] = net_devices
 
diff --git a/distributed/config.py b/distributed/config.py
index 4b1cfb35..3163c9e8 100644
--- a/distributed/config.py
+++ b/distributed/config.py
@@ -48,6 +48,8 @@ aliases = {
     "log-length": "distributed.admin.log-length",
     "log-format": "distributed.admin.log-format",
     "pdb-on-err": "distributed.admin.pdb-on-err",
+    "ucx": "distributed.comm.ucx",
+    "rmm": "distributed.rmm",
 }
 
 dask.config.rename(aliases)
diff --git a/distributed/dashboard/components/scheduler.py b/distributed/dashboard/components/scheduler.py
index 8cbf914e..83e54cf9 100644
--- a/distributed/dashboard/components/scheduler.py
+++ b/distributed/dashboard/components/scheduler.py
@@ -142,7 +142,7 @@ class Occupancy(DashboardComponent):
     @without_property_validation
     def update(self):
         with log_errors():
-            workers = list(self.scheduler.workers.values())
+            workers = self.scheduler.workers.values()
 
             y = list(range(len(workers)))
             occupancy = [ws.occupancy for ws in workers]
@@ -722,6 +722,98 @@ class BandwidthWorkers(DashboardComponent):
             update(self.source, result)
 
 
+class WorkerNetworkBandwidth(DashboardComponent):
+    """Worker network bandwidth chart
+
+    Plots horizontal bars with the read_bytes and write_bytes worker state
+    """
+
+    def __init__(self, scheduler, **kwargs):
+        with log_errors():
+            self.scheduler = scheduler
+            self.source = ColumnDataSource(
+                {
+                    "y_read": [],
+                    "y_write": [],
+                    "x_read": [],
+                    "x_write": [],
+                }
+            )
+            self.root = figure(
+                title="Worker Network Bandwidth",
+                tools="",
+                id="bk-worker-net-bandwidth",
+                name="worker_network_bandwidth",
+                **kwargs,
+            )
+
+            # read_bytes
+            self.root.hbar(
+                y="y_read",
+                right="x_read",
+                line_color=None,
+                left=0,
+                height=0.5,
+                fill_color="red",
+                legend_label="read",
+                source=self.source,
+            )
+
+            # write_bytes
+            self.root.hbar(
+                y="y_write",
+                right="x_write",
+                line_color=None,
+                left=0,
+                height=0.5,
+                fill_color="blue",
+                legend_label="write",
+                source=self.source,
+            )
+
+            self.root.axis[0].ticker = BasicTicker(**TICKS_1024)
+            self.root.xaxis[0].formatter = NumeralTickFormatter(format="0.0 b")
+            self.root.xaxis.major_label_orientation = XLABEL_ORIENTATION
+            self.root.xaxis.minor_tick_line_alpha = 0
+            self.root.x_range = Range1d(start=0)
+            self.root.yaxis.visible = False
+            self.root.ygrid.visible = False
+            self.root.toolbar_location = None
+            self.root.yaxis.visible = False
+
+    @without_property_validation
+    def update(self):
+        with log_errors():
+            workers = self.scheduler.workers.values()
+
+            h = 0.1
+            y_read = [i + 0.75 + i * h for i in range(len(workers))]
+            y_write = [i + 0.25 + i * h for i in range(len(workers))]
+
+            x_read = []
+            x_write = []
+
+            for ws in workers:
+                x_read.append(ws.metrics["read_bytes"])
+                x_write.append(ws.metrics["write_bytes"])
+
+            self.root.x_range.end = max(
+                max(x_read),
+                max(x_write),
+                100_000_000,
+                0.95 * self.root.x_range.end,
+            )
+
+            result = {
+                "y_read": y_read,
+                "y_write": y_write,
+                "x_read": x_read,
+                "x_write": x_write,
+            }
+
+            update(self.source, result)
+
+
 class ComputePerKey(DashboardComponent):
     """Bar chart showing time spend in action by key prefix"""
 
diff --git a/distributed/dashboard/components/shared.py b/distributed/dashboard/components/shared.py
index dbe6a8fa..ef857be6 100644
--- a/distributed/dashboard/components/shared.py
+++ b/distributed/dashboard/components/shared.py
@@ -482,8 +482,20 @@ class SystemMonitor(DashboardComponent):
             tools=tools,
             **kwargs,
         )
-        self.bandwidth.line(source=self.source, x="time", y="read_bytes", color="red")
-        self.bandwidth.line(source=self.source, x="time", y="write_bytes", color="blue")
+        self.bandwidth.line(
+            source=self.source,
+            x="time",
+            y="read_bytes",
+            color="red",
+            legend_label="read",
+        )
+        self.bandwidth.line(
+            source=self.source,
+            x="time",
+            y="write_bytes",
+            color="blue",
+            legend_label="write",
+        )
         self.bandwidth.yaxis.axis_label = "Bytes / second"
 
         # self.cpu.yaxis[0].formatter = NumeralTickFormatter(format='0%')
diff --git a/distributed/dashboard/scheduler.py b/distributed/dashboard/scheduler.py
index fb0efff1..cd321130 100644
--- a/distributed/dashboard/scheduler.py
+++ b/distributed/dashboard/scheduler.py
@@ -19,6 +19,7 @@ from .components.scheduler import (
     TaskGroupGraph,
     TaskProgress,
     TaskStream,
+    WorkerNetworkBandwidth,
     WorkersMemory,
     WorkerTable,
     events_doc,
@@ -67,6 +68,9 @@ applications = {
     "/individual-workers": individual_doc(WorkerTable, 500),
     "/individual-bandwidth-types": individual_doc(BandwidthTypes, 500),
     "/individual-bandwidth-workers": individual_doc(BandwidthWorkers, 500),
+    "/individual-workers-network-bandwidth": individual_doc(
+        WorkerNetworkBandwidth, 500
+    ),
     "/individual-memory-by-key": individual_doc(MemoryByKey, 500),
     "/individual-compute-time-per-key": individual_doc(ComputePerKey, 500),
     "/individual-aggregate-time-per-action": individual_doc(AggregateAction, 500),
diff --git a/distributed/diagnostics/plugin.py b/distributed/diagnostics/plugin.py
index 85be2642..f9077afd 100644
--- a/distributed/diagnostics/plugin.py
+++ b/distributed/diagnostics/plugin.py
@@ -4,8 +4,9 @@ import socket
 import subprocess
 import sys
 import uuid
+import zipfile
 
-from dask.utils import funcname
+from dask.utils import funcname, tmpfile
 
 logger = logging.getLogger(__name__)
 
@@ -175,6 +176,41 @@ class WorkerPlugin:
         """
 
 
+class NannyPlugin:
+    """Interface to extend the Nanny
+
+    A worker plugin enables custom code to run at different stages of the Workers'
+    lifecycle. A nanny plugin does the same thing, but benefits from being able
+    to run code before the worker is started, or to restart the worker if
+    necessary.
+
+    To implement a plugin implement some of the methods of this class and register
+    the plugin to your client in order to have it attached to every existing and
+    future nanny by passing ``nanny=True`` to
+    :meth:`Client.register_worker_plugin<distributed.Client.register_worker_plugin>`.
+
+    The ``restart`` attribute is used to control whether or not a running ``Worker``
+    needs to be restarted when registering the plugin.
+
+    See Also
+    --------
+    WorkerPlugin
+    SchedulerPlugin
+    """
+
+    restart = False
+
+    def setup(self, nanny):
+        """
+        Run when the plugin is attached to a nanny. This happens when the plugin is registered
+        and attached to existing nannies, or when a nanny is created after the plugin has been
+        registered.
+        """
+
+    def teardown(self, nanny):
+        """Run when the nanny to which the plugin is attached to is closed"""
+
+
 def _get_worker_plugin_name(plugin) -> str:
     """Returns the worker plugin name. If plugin has no name attribute
     a random name is used."""
@@ -289,3 +325,82 @@ class UploadFile(WorkerPlugin):
             comm=None, filename=self.filename, data=self.data, load=True
         )
         assert len(self.data) == response["nbytes"]
+
+
+class Environ(NannyPlugin):
+    restart = True
+
+    def __init__(self, environ={}):
+        self.environ = {k: str(v) for k, v in environ.items()}
+
+    async def setup(self, nanny):
+        nanny.env.update(self.environ)
+
+
+class UploadDirectory(NannyPlugin):
+    """A NannyPlugin to upload a local file to workers.
+
+    Parameters
+    ----------
+    path: str
+        A path to the directory to upload
+
+    Examples
+    --------
+    >>> from distributed.diagnostics.plugin import UploadDirectory
+    >>> client.register_worker_plugin(UploadDirectory("/path/to/directory"), nanny=True)  # doctest: +SKIP
+    """
+
+    def __init__(
+        self,
+        path,
+        restart=False,
+        update_path=False,
+        skip_words=(".git", ".github", ".pytest_cache", "tests", "docs"),
+        skip=(lambda fn: os.path.splitext(fn)[1] == ".pyc",),
+    ):
+        """
+        Initialize the plugin by reading in the data from the given file.
+        """
+        path = os.path.expanduser(path)
+        self.path = os.path.split(path)[-1]
+        self.restart = restart
+        self.update_path = update_path
+
+        self.name = "upload-directory-" + os.path.split(path)[-1]
+
+        with tmpfile(extension="zip") as fn:
+            with zipfile.ZipFile(fn, "w", zipfile.ZIP_DEFLATED) as z:
+                for root, dirs, files in os.walk(path):
+                    for file in files:
+                        filename = os.path.join(root, file)
+                        if any(predicate(filename) for predicate in skip):
+                            continue
+                        dirs = filename.split(os.sep)
+                        if any(word in dirs for word in skip_words):
+                            continue
+
+                        archive_name = os.path.relpath(
+                            os.path.join(root, file), os.path.join(path, "..")
+                        )
+                        z.write(filename, archive_name)
+
+            with open(fn, "rb") as f:
+                self.data = f.read()
+
+    async def setup(self, nanny):
+        fn = os.path.join(nanny.local_directory, f"tmp-{str(uuid.uuid4())}.zip")
+        with open(fn, "wb") as f:
+            f.write(self.data)
+
+        import zipfile
+
+        with zipfile.ZipFile(fn) as z:
+            z.extractall(path=nanny.local_directory)
+
+        if self.update_path:
+            path = os.path.join(nanny.local_directory, self.path)
+            if path not in sys.path:
+                sys.path.insert(0, path)
+
+        os.remove(fn)
diff --git a/distributed/distributed-schema.yaml b/distributed/distributed-schema.yaml
index a76e2f70..d10022f7 100644
--- a/distributed/distributed-schema.yaml
+++ b/distributed/distributed-schema.yaml
@@ -510,6 +510,11 @@ properties:
 
               See https://docs.dask.org/en/latest/setup/custom-startup.html for more information
 
+          environ:
+            type: object
+            description: |
+              Environment variables to set on all worker processes started by nannies
+
       client:
         type: object
         description: |
@@ -759,6 +764,53 @@ properties:
                       Alternatively, the key can be appended to the cert file
                       above, and this field left blank
 
+          ucx:
+            type: object
+            description: |
+              UCX provides access to other transport methods including NVLink and InfiniBand.
+            properties:
+              cuda_copy:
+                type: boolean
+                description: |
+                  Set environment variables to enable CUDA support over UCX. This may be used even if
+                  InfiniBand and NVLink are not supported or disabled, then transferring data over TCP.
+              tcp:
+                type: boolean
+                description: |
+                  Set environment variables to enable TCP over UCX, even if InfiniBand and NVLink
+                  are not supported or disabled.
+              nvlink:
+                type: boolean
+                description: |
+                  Set environment variables to enable UCX over NVLink, implies ``distributed.comm.ucx.tcp=True``.
+              infiniband:
+                type: boolean
+                description: |
+                  Set environment variables to enable UCX over InfiniBand, implies ``distributed.comm.ucx.tcp=True``.
+              rdmacm:
+                type: boolean
+                description: |
+                  Set environment variables to enable UCX RDMA connection manager support,
+                  requires ``distributed.comm.ucx.infiniband=True``.
+              net-devices:
+                type: [string, 'null']
+                description: |
+                  Interface(s) used by workers for UCX communication. Can be a string (like
+                  ``"eth0"`` for NVLink or ``"mlx5_0:1"``/``"ib0"`` for InfiniBand), ``"auto"``
+                  (requires ``distributed.comm.ucx.infiniband=True``) to pick the optimal interface per-worker based on
+                  the system's topology, or ``None`` to stay with the default value of ``"all"`` (use
+                  all available interfaces). Setting to ``"auto"`` requires UCX-Py to be installed
+                  and compiled with hwloc support. Unexpected errors can occur when using
+                  ``"auto"`` if any interfaces are disconnected or improperly configured.
+              reuse-endpoints:
+                type: [boolean, 'null']
+                description: |
+                  Enable UCX-Py reuse endpoints mechanism if ``True`` or if it's not specified and
+                  UCX < 1.11 is installed, otherwise disable reuse endpoints. This was primarily
+                  introduced to resolve an issue with CUDA IPC that has been fixed in UCX 1.10, but
+                  can cause establishing endpoints to be very slow, this is particularly noticeable in
+                  clusters of more than a few dozen workers.
+
           websockets:
             type: object
             properties:
@@ -895,58 +947,12 @@ properties:
                 type: string
                 description: Polling time to query cpu/memory statistics default 500ms
 
-  rmm:
-    type: object
-    description: |
-      Configuration options for the RAPIDS Memory Manager.
-    properties:
-      pool-size:
-        type: [integer, 'null']
-        description: |
-          The size of the memory pool in bytes.
-  ucx:
-    type: object
-    description: |
-      UCX provides access to other transport methods including NVLink and InfiniBand.
-    properties:
-      cuda_copy:
-        type: boolean
-        description: |
-          Set environment variables to enable CUDA support over UCX. This may be used even if
-          InfiniBand and NVLink are not supported or disabled, then transferring data over TCP.
-      tcp:
-        type: boolean
-        description: |
-          Set environment variables to enable TCP over UCX, even if InfiniBand and NVLink
-          are not supported or disabled.
-      nvlink:
-        type: boolean
-        description: |
-          Set environment variables to enable UCX over NVLink, implies ``ucx.tcp=True``.
-      infiniband:
-        type: boolean
-        description: |
-          Set environment variables to enable UCX over InfiniBand, implies ``ucx.tcp=True``.
-      rdmacm:
-        type: boolean
-        description: |
-          Set environment variables to enable UCX RDMA connection manager support,
-          requires ``ucx.infiniband=True``.
-      net-devices:
-        type: [string, 'null']
-        description: |
-          Interface(s) used by workers for UCX communication. Can be a string (like
-          ``"eth0"`` for NVLink or ``"mlx5_0:1"``/``"ib0"`` for InfiniBand), ``"auto"``
-          (requires ``ucx.infiniband=True``) to pick the optimal interface per-worker based on
-          the system's topology, or ``None`` to stay with the default value of ``"all"`` (use
-          all available interfaces). Setting to ``"auto"`` requires UCX-Py to be installed
-          and compiled with hwloc support. Unexpected errors can occur when using
-          ``"auto"`` if any interfaces are disconnected or improperly configured.
-      reuse-endpoints:
-        type: [boolean, 'null']
+      rmm:
+        type: object
         description: |
-          Enable UCX-Py reuse endpoints mechanism if ``True`` or if it's not specified and
-          UCX < 1.11 is installed, otherwise disable reuse endpoints. This was primarily
-          introduced to resolve an issue with CUDA IPC that has been fixed in UCX 1.10, but
-          can cause establishing endpoints to be very slow, this is particularly noticeable in
-          clusters of more than a few dozen workers.
+          Configuration options for the RAPIDS Memory Manager.
+        properties:
+          pool-size:
+            type: [integer, 'null']
+            description: |
+              The size of the memory pool in bytes.
diff --git a/distributed/distributed.yaml b/distributed/distributed.yaml
index bdee2b27..929b5867 100644
--- a/distributed/distributed.yaml
+++ b/distributed/distributed.yaml
@@ -142,6 +142,10 @@ distributed:
   nanny:
     preload: []             # Run custom modules with Nanny
     preload-argv: []        # See https://docs.dask.org/en/latest/setup/custom-startup.html
+    environ:
+      MALLOC_TRIM_THRESHOLD_: 65536
+      OMP_NUM_THREADS: 1
+      MKL_NUM_THREADS: 1
 
   client:
     heartbeat: 5s  # Interval between client heartbeats
@@ -170,6 +174,14 @@ distributed:
     default-scheme: tcp
     socket-backlog: 2048
     recent-messages-log-length: 0  # number of messages to keep for debugging
+    ucx:
+      cuda_copy: False  # enable cuda-copy
+      tcp: False  # enable tcp
+      nvlink: False  # enable cuda_ipc
+      infiniband: False # enable Infiniband
+      rdmacm: False # enable RDMACM
+      net-devices: null  # define what interface to use for UCX comm
+      reuse-endpoints: null  # enable endpoint reuse
 
     zstd:
       level: 3      # Compression level, between 1 and 22.
@@ -239,13 +251,5 @@ distributed:
     system-monitor:
       interval: 500ms
     event-loop: tornado
-rmm:
-  pool-size: null
-ucx:
-  cuda_copy: False  # enable cuda-copy
-  tcp: False  # enable tcp
-  nvlink: False  # enable cuda_ipc
-  infiniband: False # enable Infiniband
-  rdmacm: False # enable RDMACM
-  net-devices: null  # define what interface to use for UCX comm
-  reuse-endpoints: null  # enable endpoint reuse
+  rmm:
+    pool-size: null
diff --git a/distributed/nanny.py b/distributed/nanny.py
index a714489d..eea2f311 100644
--- a/distributed/nanny.py
+++ b/distributed/nanny.py
@@ -8,6 +8,7 @@ import uuid
 import warnings
 import weakref
 from contextlib import suppress
+from inspect import isawaitable
 from multiprocessing.queues import Empty
 from time import sleep as sync_sleep
 
@@ -22,15 +23,18 @@ from dask.utils import parse_timedelta
 from . import preloading
 from .comm import get_address_host, unparse_host_port
 from .comm.addressing import address_from_user_args
-from .core import CommClosedError, RPCClosed, Status, coerce_to_address
+from .core import CommClosedError, RPCClosed, Status, coerce_to_address, error_message
+from .diagnostics.plugin import _get_worker_plugin_name
 from .node import ServerNode
 from .process import AsyncProcess
 from .proctitle import enable_proctitle_on_children
+from .protocol import pickle
 from .security import Security
 from .utils import (
     TimeoutError,
     get_ip,
     json_load_robust,
+    log_errors,
     mp_context,
     parse_ports,
     silence_logging,
@@ -110,14 +114,14 @@ class Nanny(ServerNode):
 
         if local_directory is None:
             local_directory = dask.config.get("temporary-directory") or os.getcwd()
-            if not os.path.exists(local_directory):
-                os.makedirs(local_directory)
             self._original_local_dir = local_directory
             local_directory = os.path.join(local_directory, "dask-worker-space")
         else:
             self._original_local_dir = local_directory
 
         self.local_directory = local_directory
+        if not os.path.exists(self.local_directory):
+            os.makedirs(self.local_directory, exist_ok=True)
 
         self.preload = preload
         if self.preload is None:
@@ -162,7 +166,13 @@ class Nanny(ServerNode):
         self.death_timeout = parse_timedelta(death_timeout)
 
         self.Worker = Worker if worker_class is None else worker_class
-        self.env = env or {}
+        self.env = dask.config.get("distributed.nanny.environ")
+        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.config = config or dask.config.config
         worker_kwargs.update(
             {
@@ -199,8 +209,12 @@ class Nanny(ServerNode):
             "terminate": self.close,
             "close_gracefully": self.close_gracefully,
             "run": self.run,
+            "plugin_add": self.plugin_add,
+            "plugin_remove": self.plugin_remove,
         }
 
+        self.plugins = {}
+
         super().__init__(
             handlers=handlers, io_loop=self.loop, connection_args=self.connection_args
         )
@@ -294,6 +308,10 @@ class Nanny(ServerNode):
         for preload in self.preloads:
             await preload.start()
 
+        msg = await self.scheduler.register_nanny()
+        for name, plugin in msg["nanny-plugins"].items():
+            await self.plugin_add(plugin=plugin, name=name)
+
         logger.info("        Start Nanny at: %r", self.address)
         response = await self.instantiate()
         if response == Status.running:
@@ -384,6 +402,47 @@ class Nanny(ServerNode):
                 raise
         return result
 
+    async def plugin_add(self, comm=None, plugin=None, name=None):
+        with log_errors(pdb=False):
+            if isinstance(plugin, bytes):
+                plugin = pickle.loads(plugin)
+
+            if name is None:
+                name = _get_worker_plugin_name(plugin)
+
+            assert name
+
+            self.plugins[name] = plugin
+
+            logger.info("Starting Nanny plugin %s" % name)
+            if hasattr(plugin, "setup"):
+                try:
+                    result = plugin.setup(nanny=self)
+                    if isawaitable(result):
+                        result = await result
+                except Exception as e:
+                    msg = error_message(e)
+                    return msg
+            if getattr(plugin, "restart", False):
+                await self.restart()
+
+            return {"status": "OK"}
+
+    async def plugin_remove(self, comm=None, name=None):
+        with log_errors(pdb=False):
+            logger.info(f"Removing Nanny plugin {name}")
+            try:
+                plugin = self.plugins.pop(name)
+                if hasattr(plugin, "teardown"):
+                    result = plugin.teardown(nanny=self)
+                    if isawaitable(result):
+                        result = await result
+            except Exception as e:
+                msg = error_message(e)
+                return msg
+
+            return {"status": "OK"}
+
     async def restart(self, comm=None, timeout=30, executor_wait=True):
         async def _():
             if self.process is not None:
@@ -508,6 +567,14 @@ class Nanny(ServerNode):
         for preload in self.preloads:
             await preload.teardown()
 
+        teardowns = [
+            plugin.teardown(self)
+            for plugin in self.plugins.values()
+            if hasattr(plugin, "teardown")
+        ]
+
+        await asyncio.gather(*[td for td in teardowns if isawaitable(td)])
+
         self.stop()
         try:
             if self.process is not None:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index 2ba9f124..881ed7ab 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -3691,6 +3691,7 @@ class Scheduler(SchedulerState, ServerNode):
         )
         self.event_counts = defaultdict(int)
         self.worker_plugins = dict()
+        self.nanny_plugins = dict()
 
         worker_handlers = {
             "task-finished": self.handle_task_finished,
@@ -3721,6 +3722,7 @@ class Scheduler(SchedulerState, ServerNode):
             "register-client": self.add_client,
             "scatter": self.scatter,
             "register-worker": self.add_worker,
+            "register_nanny": self.add_nanny,
             "unregister": self.remove_worker,
             "gather": self.gather,
             "cancel": self.stimulus_cancel,
@@ -3760,6 +3762,8 @@ class Scheduler(SchedulerState, ServerNode):
             "register_scheduler_plugin": self.register_scheduler_plugin,
             "register_worker_plugin": self.register_worker_plugin,
             "unregister_worker_plugin": self.unregister_worker_plugin,
+            "register_nanny_plugin": self.register_nanny_plugin,
+            "unregister_nanny_plugin": self.unregister_nanny_plugin,
             "adaptive_target": self.adaptive_target,
             "workers_to_close": self.workers_to_close,
             "subscribe_worker_status": self.subscribe_worker_status,
@@ -4284,8 +4288,16 @@ class Scheduler(SchedulerState, ServerNode):
 
             if comm:
                 await comm.write(msg)
+
             await self.handle_worker(comm=comm, worker=address)
 
+    async def add_nanny(self, comm):
+        msg = {
+            "status": "OK",
+            "nanny-plugins": self.nanny_plugins,
+        }
+        return msg
+
     def update_graph_hlg(
         self,
         client=None,
@@ -6938,6 +6950,28 @@ class Scheduler(SchedulerState, ServerNode):
         responses = await self.broadcast(msg=dict(op="plugin-remove", name=name))
         return responses
 
+    async def register_nanny_plugin(self, comm, plugin, name=None):
+        """Registers a setup function, and call it on every worker"""
+        self.nanny_plugins[name] = plugin
+
+        responses = await self.broadcast(
+            msg=dict(op="plugin_add", plugin=plugin, name=name),
+            nanny=True,
+        )
+        return responses
+
+    async def unregister_nanny_plugin(self, comm, name):
+        """Unregisters a worker plugin"""
+        try:
+            self.nanny_plugins.pop(name)
+        except KeyError:
+            raise ValueError(f"The nanny plugin {name} does not exists")
+
+        responses = await self.broadcast(
+            msg=dict(op="plugin_remove", name=name), nanny=True
+        )
+        return responses
+
     def transition(self, key, finish: str, *args, **kwargs):
         """Transition a key from its current state to the finish state
 
diff --git a/docs/source/http_services.rst b/docs/source/http_services.rst
index 8e8ed1cd..75027a74 100644
--- a/docs/source/http_services.rst
+++ b/docs/source/http_services.rst
@@ -57,14 +57,15 @@ Individual bokeh plots
 - ``/individual-groups``
 - ``/individual-profile``
 - ``/individual-profile-server``
-- ``/individual-nbytes``
-- ``/individual-nbytes-cluster``
+- ``/individual-workers-memory``
+- ``/individual-cluster-memory``
 - ``/individual-cpu``
 - ``/individual-nprocessing``
 - ``/individual-occupancy``
 - ``/individual-workers``
 - ``/individual-bandwidth-types``
 - ``/individual-bandwidth-workers``
+- ``/individual-workers-network-bandwidth``
 - ``/individual-memory-by-key``
 - ``/individual-compute-time-per-key``
 - ``/individual-aggregate-time-per-action``
diff --git a/docs/source/plugins.rst b/docs/source/plugins.rst
index add010f8..bf91f6b5 100644
--- a/docs/source/plugins.rst
+++ b/docs/source/plugins.rst
@@ -75,7 +75,7 @@ the scheduler as so:
        scheduler.add_plugin(plugin)
 
 Worker Plugins
-=================
+==============
 
 :class:`distributed.diagnostics.plugin.WorkerPlugin` provides a base class
 for creating your own worker plugins. In addition, Dask provides some
@@ -91,3 +91,17 @@ Built-In Worker Plugins
 
 .. autoclass:: distributed.diagnostics.plugin.PipInstall
 .. autoclass:: distributed.diagnostics.plugin.UploadFile
+
+
+Nanny Plugins
+=============
+
+.. autoclass:: distributed.diagnostics.plugin.NannyPlugin
+   :members:
+
+
+Built-In Nanny Plugins
+----------------------
+
+.. autoclass:: distributed.diagnostics.plugin.Environ
+.. autoclass:: distributed.diagnostics.plugin.UploadDirectory
\ No newline at end of file


Test patch:
diff --git a/distributed/dashboard/tests/test_scheduler_bokeh.py b/distributed/dashboard/tests/test_scheduler_bokeh.py
index c271e5ac..f003c46d 100644
--- a/distributed/dashboard/tests/test_scheduler_bokeh.py
+++ b/distributed/dashboard/tests/test_scheduler_bokeh.py
@@ -35,6 +35,7 @@ from distributed.dashboard.components.scheduler import (
     TaskGroupGraph,
     TaskProgress,
     TaskStream,
+    WorkerNetworkBandwidth,
     WorkersMemory,
     WorkersMemoryHistogram,
     WorkerTable,
@@ -476,6 +477,37 @@ async def test_WorkerTable_with_memory_limit_as_0(c, s, a, b):
     assert wt.source.data["memory_percent"][0] == ""
 
 
+@gen_cluster(client=True)
+async def test_WorkerNetworkBandwidth(c, s, a, b):
+    nb = WorkerNetworkBandwidth(s)
+    nb.update()
+
+    assert all(len(v) == 2 for v in nb.source.data.values())
+
+    assert nb.source.data["y_read"] == [0.75, 1.85]
+    assert nb.source.data["y_write"] == [0.25, 1.35]
+
+
+@gen_cluster(client=True)
+async def test_WorkerNetworkBandwidth_metrics(c, s, a, b):
+    # Disable system monitor periodic callback to allow us to manually control
+    # when it is called below
+    a.periodic_callbacks["monitor"].stop()
+    b.periodic_callbacks["monitor"].stop()
+
+    # Update worker system monitors and send updated metrics to the scheduler
+    a.monitor.update()
+    b.monitor.update()
+    await asyncio.gather(a.heartbeat(), b.heartbeat())
+
+    nb = WorkerNetworkBandwidth(s)
+    nb.update()
+
+    for idx, ws in enumerate(s.workers.values()):
+        assert ws.metrics["read_bytes"] == nb.source.data["x_read"][idx]
+        assert ws.metrics["write_bytes"] == nb.source.data["x_write"][idx]
+
+
 @gen_cluster(client=True)
 async def test_TaskGraph(c, s, a, b):
     gp = TaskGraph(s)
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
index 82d0c3b8..7305ab1f 100644
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -6905,3 +6905,36 @@ async def test_computation_object_code_client_compute(c, s, a, b):
     assert len(comp.code) == 1
 
     assert comp.code[0] == test_function_code
+
+
+@gen_cluster(client=True, Worker=Nanny)
+async def test_upload_directory(c, s, a, b, tmp_path):
+    from dask.distributed import UploadDirectory
+
+    files = set(os.listdir())
+
+    with open(tmp_path / "foo.py", "w") as f:
+        f.write("x = 123")
+    with open(tmp_path / "bar.py", "w") as f:
+        f.write("from foo import x")
+
+    plugin = UploadDirectory(tmp_path, restart=True, update_path=True)
+    await c.register_worker_plugin(plugin)
+
+    [name] = a.plugins
+    assert os.path.split(tmp_path)[-1] in name
+
+    def f():
+        import bar
+
+        return bar.x
+
+    results = await c.run(f)
+    assert results[a.worker_address] == 123
+    assert results[b.worker_address] == 123
+
+    async with Nanny(s.address, local_directory=tmp_path / "foo", name="foo") as n:
+        results = await c.run(f)
+        assert results[n.worker_address] == 123
+
+    assert files == set(os.listdir())  # no change
diff --git a/distributed/tests/test_config.py b/distributed/tests/test_config.py
index 200493a8..cd2c5418 100644
--- a/distributed/tests/test_config.py
+++ b/distributed/tests/test_config.py
@@ -331,7 +331,7 @@ def test_schema_is_complete():
     with open(schema_fn) as f:
         schema = yaml.safe_load(f)
 
-    skip = {"default-task-durations", "bokeh-application"}
+    skip = {"default-task-durations", "bokeh-application", "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
index b23d6b1c..87ea7b9b 100644
--- a/distributed/tests/test_nanny.py
+++ b/distributed/tests/test_nanny.py
@@ -373,6 +373,21 @@ async def test_environment_variable(c, s):
     await asyncio.gather(a.close(), b.close())
 
 
+@gen_cluster(
+    nthreads=[],
+    client=True,
+    config={"distributed.nanny.environ": {"A": 1, "B": 2, "D": 4}},
+)
+async def test_environment_variable_config(c, s, monkeypatch):
+    monkeypatch.setenv("D", "123")
+    async with Nanny(s.address, env={"B": 3, "C": 4}) as n:
+        results = await c.run(lambda: os.environ)
+        assert results[n.worker_address]["A"] == "1"
+        assert results[n.worker_address]["B"] == "3"
+        assert results[n.worker_address]["C"] == "4"
+        assert results[n.worker_address]["D"] == "123"
+
+
 @gen_cluster(nthreads=[], client=True)
 async def test_data_types(c, s):
     w = await Nanny(s.address, data=dict)
@@ -548,12 +563,12 @@ class BrokenWorker(worker.Worker):
         raise StartException("broken")
 
 
-@pytest.mark.asyncio
-async def test_worker_start_exception(cleanup):
+@gen_cluster(nthreads=[])
+async def test_worker_start_exception(s):
     # make sure this raises the right Exception:
     with pytest.raises(StartException):
-        async with Nanny("tcp://localhost:1", worker_class=BrokenWorker) as n:
-            await n.start()
+        async with Nanny(s.address, worker_class=BrokenWorker) as n:
+            pass
 
 
 @pytest.mark.asyncio
@@ -564,3 +579,16 @@ async def test_failure_during_worker_initialization(cleanup):
                 async with Nanny(s.address, foo="bar") as n:
                     await n
         assert "Restarting worker" not in logs.getvalue()
+
+
+@gen_cluster(client=True, Worker=Nanny, timeout=10000000)
+async def test_environ_plugin(c, s, a, b):
+    from dask.distributed import Environ
+
+    await c.register_worker_plugin(Environ({"ABC": 123}))
+
+    async with Nanny(s.address, name="new") as n:
+        results = await c.run(os.getenv, "ABC")
+        assert results[a.worker_address] == "123"
+        assert results[b.worker_address] == "123"
+        assert results[n.worker_address] == "123"