Repository: dask/distributed

Issue:
Worker plugin registration with name is not idempotent after cluster scaling
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**What happened**:

Worker plugin registration with name is not idempotent after cluster scaling.

**What you expected to happen**:

Worker plugin registration with name is always idempotent.

**Minimal Complete Verifiable Example**:

```python
import time

from distributed import Client, LocalCluster


class FooWorkerPlugin:
    # name = "foo"  # Workaround.

    def setup(self, worker):
        if hasattr(worker, "foo"):
            raise RuntimeError(f"Worker {worker.address} already has foo!")

        worker.foo = True


cluster = LocalCluster(
    n_workers=1,
    threads_per_worker=1,
    processes=False)

client = Client(cluster)

for i in range(3):
    print("Workers:", list(cluster.scheduler.workers))

    responses = client.register_worker_plugin(
        name="foo",
        plugin=FooWorkerPlugin())

    print("Register foo plugin:", responses)

    if i == 1:
        cluster.scale(2)

    time.sleep(1)
    print()
```
```
Workers: ['inproc://192.168.10.20/34912/3']
Register foo plugin: {'inproc://192.168.10.20/34912/3': {'status': 'OK'}}

Workers: ['inproc://192.168.10.20/34912/3']
Register foo plugin: {'inproc://192.168.10.20/34912/3': {'status': 'repeat'}}

Workers: ['inproc://192.168.10.20/34912/10', 'inproc://192.168.10.20/34912/3']
Traceback (most recent call last):
  ...
RuntimeError: Worker inproc://192.168.10.20/34912/10 already has foo!
```
With the workaround:
```
Workers: ['inproc://192.168.10.20/32776/3']
Register foo plugin: {'inproc://192.168.10.20/32776/3': {'status': 'OK'}}

Workers: ['inproc://192.168.10.20/32776/3']
Register foo plugin: {'inproc://192.168.10.20/32776/3': {'status': 'repeat'}}

Workers: ['inproc://192.168.10.20/32776/10', 'inproc://192.168.10.20/32776/3']
Register foo plugin: {'inproc://192.168.10.20/32776/10': {'status': 'repeat'}, 'inproc://192.168.10.20/32776/3': {'status': 'repeat'}}
```

**Anything else we need to know?**:

I propose that `Scheduler.worker_plugins` should be a dict `{name: plugin}` instead of a list of plugins. This is a change to a public API.

**Environment**:

- Dask version: 2.25.0
- Python version: 3.8.5
- Operating System: Windows
- Install method (conda, pip, source): conda



Discussion:
Making the worker_plugins field a dict makes sense to me.  I'm not
concerned about changing this API.  Very few people modify worker_plugins
directly.

On Tue, Sep 8, 2020 at 7:59 AM Nick Evans <notifications@github.com> wrote:

> *What happened*:
>
> Worker plugin registration with name is not idempotent after cluster
> scaling.
>
> *What you expected to happen*:
>
> Worker plugin registration with name is always idempotent.
>
> *Minimal Complete Verifiable Example*:
>
> import time
> from distributed import Client, LocalCluster
>
> class FooWorkerPlugin:
>     # name = "foo"  # Workaround.
>
>     def setup(self, worker):
>         if hasattr(worker, "foo"):
>             raise RuntimeError(f"Worker {worker.address} already has foo!")
>
>         worker.foo = True
>
> cluster = LocalCluster(
>     n_workers=1,
>     threads_per_worker=1,
>     processes=False)
> client = Client(cluster)
> for i in range(3):
>     print("Workers:", list(cluster.scheduler.workers))
>
>     responses = client.register_worker_plugin(
>         name="foo",
>         plugin=FooWorkerPlugin())
>
>     print("Register foo plugin:", responses)
>
>     if i == 1:
>         cluster.scale(2)
>
>     time.sleep(1)
>     print()
>
> Workers: ['inproc://192.168.10.20/34912/3']
> Register foo plugin: {'inproc://192.168.10.20/34912/3': {'status': 'OK'}}
>
> Workers: ['inproc://192.168.10.20/34912/3']
> Register foo plugin: {'inproc://192.168.10.20/34912/3': {'status': 'repeat'}}
>
> Workers: ['inproc://192.168.10.20/34912/10', 'inproc://192.168.10.20/34912/3']
> Traceback (most recent call last):
>   ...
> RuntimeError: Worker inproc://192.168.10.20/34912/10 already has foo!
>
> With the workaround:
>
> Workers: ['inproc://192.168.10.20/32776/3']
> Register foo plugin: {'inproc://192.168.10.20/32776/3': {'status': 'OK'}}
>
> Workers: ['inproc://192.168.10.20/32776/3']
> Register foo plugin: {'inproc://192.168.10.20/32776/3': {'status': 'repeat'}}
>
> Workers: ['inproc://192.168.10.20/32776/10', 'inproc://192.168.10.20/32776/3']
> Register foo plugin: {'inproc://192.168.10.20/32776/10': {'status': 'repeat'}, 'inproc://192.168.10.20/32776/3': {'status': 'repeat'}}
>
> *Anything else we need to know?*:
>
> I propose that Scheduler.worker_plugins should be a dict {name: plugin}
> instead of a list of plugins. This is a change to a public API.
>
> *Environment*:
>
>    - Dask version: 2.25.0
>    - Python version: 3.8.5
>    - Operating System: Windows
>    - Install method (conda, pip, source): conda
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/distributed/issues/4100>, or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTCIUCFWW24TXV6JNELSEZBGHANCNFSM4RAC3WJQ>
> .
>


Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -3844,7 +3844,7 @@ def get_task_stream(self, comm=None, start=None, stop=None, count=None):
 
     async def register_worker_plugin(self, comm, plugin, name=None):
         """ Registers a setup function, and call it on every worker """
-        self.worker_plugins.append(plugin)
+        self.worker_plugins.append({"plugin": plugin, "name": name})
 
         responses = await self.broadcast(
             msg=dict(op="plugin-add", plugin=plugin, name=name)
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -845,8 +845,8 @@ async def _register_with_scheduler(self):
         else:
             await asyncio.gather(
                 *[
-                    self.plugin_add(plugin=plugin)
-                    for plugin in response["worker-plugins"]
+                    self.plugin_add(**plugin_kwargs)
+                    for plugin_kwargs in response["worker-plugins"]
                 ]
             )
 


Test patch:
diff --git a/distributed/diagnostics/tests/test_worker_plugin.py b/distributed/diagnostics/tests/test_worker_plugin.py
--- a/distributed/diagnostics/tests/test_worker_plugin.py
+++ b/distributed/diagnostics/tests/test_worker_plugin.py
@@ -136,6 +136,23 @@ async def test_release_dep_called(c, s, w):
     await async_wait_for(lambda: not (w.task_state or w.dep_state), timeout=10)
 
 
+@gen_cluster(nthreads=[("127.0.0.1", 1)], client=True)
+async def test_registering_with_name_arg(c, s, w):
+    class FooWorkerPlugin:
+        def setup(self, worker):
+            if hasattr(worker, "foo"):
+                raise RuntimeError(f"Worker {worker.address} already has foo!")
+
+            worker.foo = True
+
+    responses = await c.register_worker_plugin(FooWorkerPlugin(), name="foo")
+    assert list(responses.values()) == [{"status": "OK"}]
+
+    async with Worker(s.address, loop=s.loop):
+        responses = await c.register_worker_plugin(FooWorkerPlugin(), name="foo")
+        assert list(responses.values()) == [{"status": "repeat"}] * 2
+
+
 @gen_cluster(nthreads=[("127.0.0.1", 1)], client=True)
 async def test_empty_plugin(c, s, w):
     class EmptyPlugin: