fix progress_stream teardown

dask · Feb 16, 2022 · 2ab5a76 · 2ab5a76
1 parent 8d0df89
commit 2ab5a76
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/distributed/diagnostics/progress_stream.py b/distributed/diagnostics/progress_stream.py
@@ -1,5 +1,4 @@
 import logging
-from functools import partial
 
 from tlz import merge, valmap
 
@@ -22,10 +21,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 +53,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