Skip to content

Commit

Permalink
Allow deserialized plugins in register_scheduler_plugin (#6401)
Browse files Browse the repository at this point in the history
Previously we would always deserialize the input.
However in situations where we register the plugin directly on the
scheduler (such as within a preload) this isn't necessary.
  • Loading branch information
mrocklin authored May 20, 2022
1 parent 90efa71 commit 41a54ee
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 1 deletion.
11 changes: 11 additions & 0 deletions distributed/diagnostics/tests/test_scheduler_plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -198,3 +198,14 @@ def f():
await c.submit(f)

assert ("foo", 123) in s._recorded_events


@gen_cluster(client=True)
async def test_register_plugin_on_scheduler(c, s, a, b):
class MyPlugin(SchedulerPlugin):
async def start(self, scheduler: Scheduler) -> None:
scheduler._foo = "bar" # type: ignore

await s.register_scheduler_plugin(MyPlugin())

assert s._foo == "bar"
3 changes: 2 additions & 1 deletion distributed/scheduler.py
Original file line number Diff line number Diff line change
Expand Up @@ -4865,7 +4865,8 @@ async def register_scheduler_plugin(self, plugin, name=None, idempotent=None):
"arbitrary bytestrings using pickle via the "
"'distributed.scheduler.pickle' configuration setting."
)
plugin = loads(plugin)
if not isinstance(plugin, SchedulerPlugin):
plugin = loads(plugin)

if name is None:
name = _get_plugin_name(plugin)
Expand Down

0 comments on commit 41a54ee

Please sign in to comment.