Can't register WorkerPlugin subclass with Client.register_plugin

[ivirshup]

Describe the issue:

Registering a WorkerPlugin subclass with Client.register_plugin raises an error. Client.register_worker_plugin works fine, but warns of deprecation.

My real use case is trying to run a "start collecting performance stats" and "stop collecting performance stats" command on each worker for a storage backend.

Minimal Complete Verifiable Example:

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                                 Traceback (most recent call last)
Cell In[1], line 14
     12 cluster = dd.LocalCluster()
     13 client = dd.Client(cluster)
---> 14 client.register_plugin(PolitePlugin())

File ~/miniforge3/envs/test-dask-worker-plugin/lib/python3.11/site-packages/distributed/client.py:4974, in Client.register_plugin(self, plugin, name, idempotent)
   4972     idempotent = getattr(plugin, "idempotent", False)
   4973 assert isinstance(idempotent, bool)
-> 4974 return self._register_plugin(plugin, name, idempotent)

File ~/miniforge3/envs/test-dask-worker-plugin/lib/python3.11/functools.py:946, in singledispatchmethod.__get__.<locals>._method(*args, **kwargs)
    944 def _method(*args, **kwargs):
    945     method = self.dispatcher.dispatch(args[0].__class__)
--> 946     return method.__get__(obj, cls)(*args, **kwargs)

File ~/miniforge3/envs/test-dask-worker-plugin/lib/python3.11/site-packages/distributed/client.py:4985, in Client._register_plugin(self, plugin, name, idempotent)
   4983 if isinstance(plugin, type):
   4984     raise TypeError("Please provide an instance of a plugin, not a type.")
-> 4985 raise TypeError(
   4986     "Registering duck-typed plugins is not allowed. Please inherit from "
   4987     "NannyPlugin, WorkerPlugin, or SchedulerPlugin to create a plugin."
   4988 )

TypeError: Registering duck-typed plugins is not allowed. Please inherit from NannyPlugin, WorkerPlugin, or SchedulerPlugin to create a plugin.

Anything else we need to know?:

This seems to work fine with register_worker_plugin

client.register_worker_plugin(PolitePlugin())

Environment:

• Dask version: 2024.7.0
• Python version: 3.11
• Operating System: ubuntu and macos
• Install method (conda, pip, source): pip

[jrbourbeau]

Thanks for the issue @ivirshup. Something strange is going on with the WorkerPlugin import. I'm able to reproduce with your minimal example (thanks for that), but things work as expected if I use

# Removed `dask.`
from distributed.diagnostics.plugin import WorkerPlugin

or

# Shorter version (also without `dask.` prefix)
from distributed import WorkerPlugin

instead