Skip to content

Commit

Permalink
Add idempotent to register_scheduler_plugin client (#6328)
Browse files Browse the repository at this point in the history
Updates client-side interface of register_scheduler_plugin to pass through idempotent flag.

Extending #5545

Finalizing removal of kwargs support
deprecated in #5699
and removed in #6144
  • Loading branch information
asford authored May 13, 2022
1 parent b8b45c6 commit bab29c5
Showing 1 changed file with 6 additions and 5 deletions.
11 changes: 6 additions & 5 deletions distributed/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -4418,16 +4418,14 @@ async def _get_task_stream(
else:
return msgs

async def _register_scheduler_plugin(self, plugin, name, **kwargs):
if isinstance(plugin, type):
plugin = plugin(**kwargs)

async def _register_scheduler_plugin(self, plugin, name, idempotent=False):
return await self.scheduler.register_scheduler_plugin(
plugin=dumps(plugin, protocol=4),
name=name,
idempotent=idempotent,
)

def register_scheduler_plugin(self, plugin, name=None):
def register_scheduler_plugin(self, plugin, name=None, idempotent=False):
"""Register a scheduler plugin.
See https://distributed.readthedocs.io/en/latest/plugins.html#scheduler-plugins
Expand All @@ -4439,6 +4437,8 @@ def register_scheduler_plugin(self, plugin, name=None):
name : str
Name for the plugin; if None, a name is taken from the
plugin instance or automatically generated if not present.
idempotent : bool
Do not re-register if a plugin of the given name already exists.
"""
if name is None:
name = _get_plugin_name(plugin)
Expand All @@ -4447,6 +4447,7 @@ def register_scheduler_plugin(self, plugin, name=None):
self._register_scheduler_plugin,
plugin=plugin,
name=name,
idempotent=idempotent,
)

def register_worker_callbacks(self, setup=None):
Expand Down

0 comments on commit bab29c5

Please sign in to comment.