While registering a scheduler plugin get TypeError: PooledRPCCall.__getattr__.<locals>.send_recv_from_rpc() takes 0 positional arguments but 1 was given

[markcoletti]

Describe the issue:

When registering a Dask scheduler plugin via Scheduler.add_plugin, we get the following error:

TypeError: PooledRPCCall.__getattr__.<locals>.send_recv_from_rpc() takes 0 positional arguments but 1 was given

The offending code is in distributed's core.py:

     def __getattr__(self, key):
        async def send_recv_from_rpc(**kwargs): # THIS IS WHERE THE ERROR IS RAISED
            if self.serializers is not None and kwargs.get("serializers") is None:
                kwargs["serializers"] = self.serializers
            if self.deserializers is not None and kwargs.get("deserializers") is None:
                kwargs["deserializers"] = self.deserializers
            comm = await self.pool.connect(self.addr)
            prev_name, comm.name = comm.name, "ConnectionPool." + key
            try:
                return await send_recv(comm=comm, op=key, **kwargs)
            finally:
                self.pool.reuse(self.addr, comm)
                comm.name = prev_name

Minimal Complete Verifiable Example:

Taken from your own example of a scheduler plugin, which DOES NOT WORK NOW.

from distributed import Client, LocalCluster, SchedulerPlugin

class MySchedulerPlugin(SchedulerPlugin):
    def __init__(self):
        self.counter = 0

    def transition(self, key, start, finish, *args, **kwargs):
        if start == 'processing' and finish == 'memory':
            self.counter += 1

    def restart(self, scheduler):
        self.counter = 0

if __name__ == '__main__':
    with Client() as client:
        print(f'scheduler={client.scheduler}')
        my_scheduler_plugin = MySchedulerPlugin()
        client.scheduler.add_plugin(my_scheduler_plugin)
    print('Done')

This produces the following;

scheduler=<pooled rpc to 'tcp://127.0.0.1:57991'>
Traceback (most recent call last):
  File "/Users/may/Projects/scratch/dask_scheduler/scheduler.py", line 22, in <module>
    client.scheduler.add_plugin(my_scheduler_plugin)
TypeError: PooledRPCCall.__getattr__.<locals>.send_recv_from_rpc() takes 0 positional arguments but 1 was given

Anything else we need to know?:
N/A

Environment:

• Dask version:

dask                    2024.10.0
distributed             2024.10.0

And the same behavior observed for 2024.11.2

• Python version: 3.10.15 and 3.11.7
• Operating System:
  • Darwin mac135909 24.2.0 Darwin Kernel Version 24.2.0: Fri Dec 6 18:56:34 PST 2024; root:xnu-11215.61.5~2/RELEASE_ARM64_T6020 arm64
  • Linux login05 5.14.21-150500.55.49_13.0.57-cray_shasta_c #1 SMP Sun May 12 13:35:37 UTC 2024 (33add2b) x86_64 x86_64 x86_64 GNU/Linux
• Install method (conda, pip, source):
  • conda and pip

[hendrikmakait]

Hi, please use Client.register_plugin(...) instead of the undocumened Client.scheduler.add_plugin. FWIW, Client.scheduler is considered private as it's not documented, but we may want to mark it explicitly as private.

[hendrikmakait]

Would you be interested in submitting a PR that fixes the outdated example?