Repository: dask/distributed

Issue:
Add `Client.register_scheduler_plugin`
Today scheduler plugins are added to a cluster by calling the `Scheduler.add_plugin` method. I often see plugins added at scheduler startup through a preload script.

However, it's also useful to be able to add a scheduler plugin to an existing, already running scheduler. In these cases one often doesn't have direct access to the `Scheduler` instance and ends up calling `client.run_on_scheduler` to add their plugin:

```python
def register_plugin(plugin):
    await plugin.start()
    dask_scheduler.add_plugin(plugin)

client.run_on_scheduler(register_plugin)
```

This works, but is a little awkward. Additionally, some users might not be aware that `client.run_on_scheduler` exists.

It would be nice if we had a `client.register_scheduler_plugin`, which was similar to our existing `client.register_worker_plugin`, to allow users to register their own scheduler plugins in a more straightforward way. 

cc @douglasdavis who may be interested in adding a `client.register_scheduler_plugin` method



Discussion:
Additionally, we're in the process of adding a `client.unregister_worker_plugin` method (xref https://github.com/dask/distributed/issues/4237). So we could also add a corresponding `client.unregister_scheduler_plugin` method as well. 

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -4007,6 +4007,34 @@ async def _get_task_stream(
         else:
             return msgs
 
+    async def _register_scheduler_plugin(self, plugin, **kwargs):
+        if isinstance(plugin, type):
+            plugin = plugin(**kwargs)
+
+        return await self.scheduler.register_scheduler_plugin(
+            plugin=dumps(plugin, protocol=4)
+        )
+
+    def register_scheduler_plugin(self, plugin, **kwargs):
+        """Register a scheduler plugin.
+
+        See https://distributed.readthedocs.io/en/latest/plugins.html#scheduler-plugins
+
+        Parameters
+        ----------
+        plugin : SchedulerPlugin
+            Plugin class or object to pass to the scheduler.
+        **kwargs : Any
+            Arguments passed to the Plugin class (if Plugin is an
+            instance kwargs are unused).
+
+        """
+        return self.sync(
+            self._register_scheduler_plugin,
+            plugin=plugin,
+            **kwargs,
+        )
+
     def register_worker_callbacks(self, setup=None):
         """
         Registers a setup callback function for all current and future workers.
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -58,6 +58,7 @@
 from .multi_lock import MultiLockExtension
 from .node import ServerNode
 from .proctitle import setproctitle
+from .protocol.pickle import loads
 from .publish import PublishExtension
 from .pubsub import PubSubSchedulerExtension
 from .queues import QueueExtension
@@ -3575,6 +3576,7 @@ def __init__(
             "heartbeat_worker": self.heartbeat_worker,
             "get_task_status": self.get_task_status,
             "get_task_stream": self.get_task_stream,
+            "register_scheduler_plugin": self.register_scheduler_plugin,
             "register_worker_plugin": self.register_worker_plugin,
             "unregister_worker_plugin": self.unregister_worker_plugin,
             "adaptive_target": self.adaptive_target,
@@ -5210,6 +5212,24 @@ def remove_plugin(self, plugin):
         """Remove external plugin from scheduler"""
         self.plugins.remove(plugin)
 
+    async def register_scheduler_plugin(self, comm=None, plugin=None):
+        """Register a plugin on the scheduler."""
+        if not dask.config.get("distributed.scheduler.pickle"):
+            raise ValueError(
+                "Cannot register a scheduler plugin as the scheduler "
+                "has been explicitly disallowed from deserializing "
+                "arbitrary bytestrings using pickle via the "
+                "'distributed.scheduler.pickle' configuration setting."
+            )
+        plugin = loads(plugin)
+
+        if hasattr(plugin, "start"):
+            result = plugin.start(self)
+            if inspect.isawaitable(result):
+                result = await result
+
+        self.add_plugin(plugin=plugin)
+
     def worker_send(self, worker, msg):
         """Send message to worker
 


Test patch:
diff --git a/distributed/diagnostics/tests/test_scheduler_plugin.py b/distributed/diagnostics/tests/test_scheduler_plugin.py
--- a/distributed/diagnostics/tests/test_scheduler_plugin.py
+++ b/distributed/diagnostics/tests/test_scheduler_plugin.py
@@ -125,3 +125,41 @@ async def close(self):
 
     assert plugin.history == ["started", "closed"]
     assert plugin.scheduler is s
+
+
+@gen_cluster(client=True)
+async def test_register_scheduler_plugin(c, s, a, b):
+    class Dummy1(SchedulerPlugin):
+        def start(self, scheduler):
+            scheduler.foo = "bar"
+
+    assert not hasattr(s, "foo")
+    await c.register_scheduler_plugin(Dummy1)
+    assert s.foo == "bar"
+
+    class Dummy2(SchedulerPlugin):
+        def start(self, scheduler):
+            raise RuntimeError("raising in start method")
+
+    n_plugins = len(s.plugins)
+    with pytest.raises(RuntimeError, match="raising in start method"):
+        await c.register_scheduler_plugin(Dummy2)
+    # total number of plugins should be unchanged
+    assert n_plugins == len(s.plugins)
+
+
+@gen_cluster(client=True, config={"distributed.scheduler.pickle": False})
+async def test_register_scheduler_plugin_pickle_disabled(c, s, a, b):
+    class Dummy1(SchedulerPlugin):
+        def start(self, scheduler):
+            scheduler.foo = "bar"
+
+    n_plugins = len(s.plugins)
+    with pytest.raises(ValueError) as excinfo:
+        await c.register_scheduler_plugin(Dummy1)
+
+    msg = str(excinfo.value)
+    assert "disallowed from deserializing" in msg
+    assert "distributed.scheduler.pickle" in msg
+
+    assert n_plugins == len(s.plugins)