Repository: dask/distributed

Issue:
AllProgress plugin not removed after progress_stream: `TypeError: remove_plugin() takes 0 positional arguments but 2 were given`
**What happened**:
```
distributed/diagnostics/tests/test_progress_stream.py::test_progress_quads_too_many PASSED
distributed.worker - WARNING - Compute Failed
Function:  div
args:      (1, 0)
kwargs:    {}
Exception: "ZeroDivisionError('division by zero')"
distributed.utils - ERROR - remove_plugin() takes 0 positional arguments but 2 were given
Traceback (most recent call last):
  File "/Users/runner/work/distributed/distributed/distributed/utils.py", line 680, in log_errors
    yield
  File "/Users/runner/work/distributed/distributed/distributed/scheduler.py", line 7057, in feed
    teardown(self, state)
TypeError: remove_plugin() takes 0 positional arguments but 2 were given
distributed.core - ERROR - Exception while handling op feed
Traceback (most recent call last):
  File "/Users/runner/work/distributed/distributed/distributed/core.py", line 520, in handle_comm
    result = await result
  File "/Users/runner/work/distributed/distributed/distributed/scheduler.py", line 7057, in feed
    teardown(self, state)
TypeError: remove_plugin() takes 0 positional arguments but 2 were given
tornado.application - ERROR - Exception in callback functools.partial(<function TCPServer._handle_connection.<locals>.<lambda> at 0x13f9a6f70>, <Task finished name='Task-49498' coro=<BaseTCPListener._handle_stream() done, defined at /Users/runner/work/distributed/distributed/distributed/comm/tcp.py:513> exception=TypeError('remove_plugin() takes 0 positional arguments but 2 were given')>)
Traceback (most recent call last):
  File "/Users/runner/miniconda3/envs/dask-distributed/lib/python3.9/site-packages/tornado/ioloop.py", line 741, in _run_callback
    ret = callback()
  File "/Users/runner/miniconda3/envs/dask-distributed/lib/python3.9/site-packages/tornado/tcpserver.py", line 331, in <lambda>
    gen.convert_yielded(future), lambda f: f.result()
  File "/Users/runner/work/distributed/distributed/distributed/comm/tcp.py", line 530, in _handle_stream
    await self.comm_handler(comm)
  File "/Users/runner/work/distributed/distributed/distributed/core.py", line 520, in handle_comm
    result = await result
  File "/Users/runner/work/distributed/distributed/distributed/scheduler.py", line 7057, in feed
    teardown(self, state)
TypeError: remove_plugin() takes 0 positional arguments but 2 were given
```
**What you expected to happen**:
The "all-progress" plugin should be removed when "progress_stream" is torn down.


Discussion:
`remove_plugin` is defined to only take kwargs:
https://github.com/dask/distributed/blob/39024291e429d983d7b73064c209701b68f41f71/distributed/diagnostics/progress_stream.py#L25-L29

but `teardown` is called with `self, state`: 
https://github.com/dask/distributed/blob/39024291e429d983d7b73064c209701b68f41f71/distributed/scheduler.py#L7056-L7057

Resolving patch:
diff --git a/distributed/diagnostics/progress_stream.py b/distributed/diagnostics/progress_stream.py
--- a/distributed/diagnostics/progress_stream.py
+++ b/distributed/diagnostics/progress_stream.py
@@ -1,10 +1,8 @@
 import logging
-from functools import partial
 
 from tlz import merge, valmap
 
 from ..core import coerce_to_address, connect
-from ..scheduler import Scheduler
 from ..utils import color_of, key_split
 from ..worker import dumps_function
 from .progress import AllProgress
@@ -22,10 +20,10 @@ def counts(scheduler, allprogress):
     )
 
 
-def remove_plugin(**kwargs):
+def _remove_all_progress_plugin(self, *args, **kwargs):
     # Wrapper function around `Scheduler.remove_plugin` to avoid raising a
     # `PicklingError` when using a cythonized scheduler
-    return Scheduler.remove_plugin(**kwargs)
+    self.remove_plugin(name=AllProgress.name)
 
 
 async def progress_stream(address, interval):
@@ -54,7 +52,7 @@ async def progress_stream(address, interval):
             "setup": dumps_function(AllProgress),
             "function": dumps_function(counts),
             "interval": interval,
-            "teardown": dumps_function(partial(remove_plugin, name=AllProgress.name)),
+            "teardown": dumps_function(_remove_all_progress_plugin),
         }
     )
     return comm


Test patch: