Raise helpful error when using the wrong plugin base classes

The Client.register_plugin() method uses @functools.singledispatchmethod to dispatch to the correct plugin registration method. However singledispatchmethod doesn't like the dask.distributed import shim and can't determine that dask.distributed.diagnostics.plugins.WorkerPlugin is the same as distributed.diagnostics.plugins.WorkerPlugin.

This means that plugins based on classes like dask.distributed.diagnostics.plugins.WorkerPlugin raise the TypeError("Registering duck-typed plugins is not allowed. Please inherit from NannyPlugin, WorkerPlugin, or SchedulerPlugin to create a plugin.") exception.

I explored importing the base classes from the dask.distributed path and registering those for dispatch too, but it's not possible to import them from dask.distributed.diagnostics.plugins while distributed is being imported. And it's not possible to register them at runtime, only at import time.

Therefore the best option here appears to be to improve the error message. So this PR checks for plugins based on classes from dask.distributed.diagnostics.plugins and makes the error message more helpful.

In [1]: from dask.distributed.diagnostics.plugin import WorkerPlugin
   ...: import dask.distributed as dd
   ...: 
   ...: class PolitePlugin(WorkerPlugin):
   ...:     def setup(self, worker):
   ...:         print("hi", worker)
   ...:         self.worker = worker
   ...: 
   ...:     def teardown(self, worker):
   ...:         print("bye", worker)
   ...: 
   ...: cluster = dd.LocalCluster()
   ...: client = dd.Client(cluster)
   ...: client.register_plugin(PolitePlugin())
---------------------------------------------------------------------------
TypeError: Importing plugin base classes from `dask.distributed.diagnostics.plugin` is not supported. Please import directly from `distributed.diagnostics.plugin` instead.

Closes #8778

[ ] Tests added / passed
[ ] Passes pre-commit run --all-files

dask / distributed

Raise helpful error when using the wrong plugin base classes #8893

Unit Test Results