Notify worker plugins when a task is released #3817

nre · 2020-05-20T13:53:41Z

nre · 2020-06-01T14:37:29Z

I have used asyncio.sleep to prevent the tests from being flaky because I see it being used in the other tests but please let me know if there is a better way.

martindurant

Basically, I think the change is nice, but I worry about the API breakage for the notification, and also had a couple of style thoughts.
@quasiben or @mrocklin should probably give a moment to consider whether the name change for the _notify_transition (which ought to be an internal method) could have ramifications.

martindurant · 2020-06-15T14:47:53Z

distributed/diagnostics/plugin.py

@@ -147,3 +148,34 @@ def transition(self, key, start, finish, **kwargs):
            Final state of the transition.
        kwargs: More options passed when transitioning
        """
+
+    def release_key(self, key, state, cause, reason, report):


Some of the variables are types in the docstring example, but not here. I know this is true of the existing methods, but seems a little odd. Also, cause does appear to have a type below.
reason should have a description.

martindurant · 2020-06-15T14:49:47Z

distributed/diagnostics/tests/test_worker_plugin.py

+        self.observed_notifications.append(ReleasedDep(dep, state))
+
+
+Transition = namedtuple("Transition", "key, start, finish")


Do you think these definitions might be useful outside of tests?

I don't think that we should generalize these further. In fact, I wouldn't mind them being scoped to just this one test if that's alright. I'd like to reduce the number of abstractions we create if possible.

Personally, I'd prefer just having a dictionary here rather than three new namedtuples, but that's purely subjective, and not something that should hold up this PR.

martindurant · 2020-06-15T14:53:03Z

distributed/diagnostics/tests/test_worker_plugin.py


    await c.register_worker_plugin(plugin)
    await c.submit(lambda x: x, 1, key="task")
+    await asyncio.sleep(DELAY_BEFORE_TEARDOWN)


Perhaps use distributed.utils_test.async_wait_for ?

+1 . In general we avoid writing tests that depend on sleeping a certain amount of time.

Also, if a global value like this is only being used once then I would greatly prefer to put the value directly in the test. Otherwise reviewers and future devs will have to dive through more code to see what is happening. See http://matthewrocklin.com/blog/work/2019/06/23/avoid-indirection for more thoughts.

I understand the desire to avoid flaky tests, however if the test depends on a sleep, then it may be sensitive to delays. For example, on small CI machines we frequently experience GC cleanups on the order of seconds. So if this test isn't robust to random multi-second-long pauses with no work then it is likely to fail eventually.

martindurant · 2020-06-15T14:55:15Z

distributed/diagnostics/tests/test_worker_plugin.py

            ):
-                assert expected == real
+                assert type(expected) is type(real) and expected == real


It seems to me a little weird to have the test as part of the teardown; could you not make an assert_ok method (or better name) to be called in the body of the test?

To be clear, this is the way that it was already done (this is my fault, not @nre's)

I agree that it's atypical, but it does isolate the error effectively. This is also one of the benefits of running everything in the same thread.

martindurant · 2020-06-15T14:58:23Z

distributed/worker.py

@@ -2825,15 +2829,16 @@ def get_call_stack(self, comm=None, keys=None):
        result = {k: profile.call_stack(frame) for k, frame in frames.items()}
        return result

-    def _notify_transition(self, key, start, finish, **kwargs):
+    def _notify_plugins(self, method_name, *args, **kwargs):


This is the potential breaking change, due to renaming the method. I'm not certain the rename is required, since the release modes could be seen as transitions too. I understand you also want to pass different arguments here. Should there be a deprecation cycle for the old name/signature.

I don't think that user-created WorkerPlugins were aware of the use of the _notify_transition method. This was a purely internal detail. The user API of having a def transition(...) method on the plugin itself remains the same. I think that that is what we care about for API stability.

mrocklin

Thank you for the PR @nre. I apologize for the delay in review.

I've left several comments here. If you have time to look through them and share your thoughts that would be welcome.

mrocklin · 2020-06-15T17:54:24Z

distributed/diagnostics/tests/test_worker_plugin.py

            ):
-                assert expected == real
+                assert type(expected) is type(real) and expected == real


To be clear, this is the way that it was already done (this is my fault, not @nre's)

I agree that it's atypical, but it does isolate the error effectively. This is also one of the benefits of running everything in the same thread.

mrocklin · 2020-06-15T17:55:51Z

distributed/diagnostics/tests/test_worker_plugin.py

+        self.observed_notifications.append(ReleasedDep(dep, state))
+
+
+Transition = namedtuple("Transition", "key, start, finish")


I don't think that we should generalize these further. In fact, I wouldn't mind them being scoped to just this one test if that's alright. I'd like to reduce the number of abstractions we create if possible.

Personally, I'd prefer just having a dictionary here rather than three new namedtuples, but that's purely subjective, and not something that should hold up this PR.

mrocklin · 2020-06-15T17:57:06Z

distributed/diagnostics/tests/test_worker_plugin.py


    await c.register_worker_plugin(plugin)
    await c.submit(lambda x: x, 1, key="task")
+    await asyncio.sleep(DELAY_BEFORE_TEARDOWN)


+1 . In general we avoid writing tests that depend on sleeping a certain amount of time.

Also, if a global value like this is only being used once then I would greatly prefer to put the value directly in the test. Otherwise reviewers and future devs will have to dive through more code to see what is happening. See http://matthewrocklin.com/blog/work/2019/06/23/avoid-indirection for more thoughts.

mrocklin · 2020-06-15T17:58:04Z

distributed/diagnostics/tests/test_worker_plugin.py

+    expected_notifications = [
+        Transition("task", "waiting", "ready"),
+        Transition("task", "ready", "executing"),
+        Transition("task", "executing", "error"),
    ]


I'd prefer these to not be Transition types if that's ok. I'd prefer us to stick to core data structures like dicts if possible. As a reviewer I don't know what this Transition type is or what I should expect of it, but I will immediately understand how to manage a dict.

mrocklin · 2020-06-15T17:59:28Z

distributed/worker.py

@@ -2825,15 +2829,16 @@ def get_call_stack(self, comm=None, keys=None):
        result = {k: profile.call_stack(frame) for k, frame in frames.items()}
        return result

-    def _notify_transition(self, key, start, finish, **kwargs):
+    def _notify_plugins(self, method_name, *args, **kwargs):


I don't think that user-created WorkerPlugins were aware of the use of the _notify_transition method. This was a purely internal detail. The user API of having a def transition(...) method on the plugin itself remains the same. I think that that is what we care about for API stability.

mrocklin · 2020-06-15T18:00:50Z

distributed/worker.py

+            try:
+                plugin_method = getattr(plugin, method_name)
+            except AttributeError:
+                continue


I have a slight preference for if hasattr(plugin, method_name) over a try-except. It seems more clear, and it's also marginally faster when we expect many misses (try-except is faster than if-else if the try block succeeds, but about 5x slower if the except path is taken).

nre · 2020-06-17T10:12:54Z

@mrocklin Thank you for the Bokeh fix.

Sorry for this force-push but the tests need the Bokeh fix.

The test failure is unrelated to this PR: test_tcp_comm_closed_implicit

Changes made:

Add missing type annotations and parameter description to docstrings.
Use async_wait_for instead of sleep.
Use dict instead of namedtuple.
Use hasattr instead of try.

Changes not made but I am happy to do if required:

Add type annotations to worker plugin methods.
Restore the _notify_transition method.

mrocklin · 2020-06-17T13:20:21Z

Sorry for this force-push but the tests need the Bokeh fix.

Not a problem. We squash-and-merge anyway.

mrocklin · 2020-06-18T13:31:41Z

This is in. Thanks @nre !

nre force-pushed the worker_plugin_release_key branch 3 times, most recently from e5493a3 to 65316c2 Compare June 1, 2020 12:57

nre marked this pull request as ready for review June 1, 2020 14:32

martindurant reviewed Jun 15, 2020

View reviewed changes

mrocklin reviewed Jun 15, 2020

View reviewed changes

nre added 5 commits June 17, 2020 10:05

Notify worker plugins when a task/dep is released

1fa8fed

Fix worker-plugin docs missing types

199dee5

Wait for worker to finish

c2d524a

Use dicts instead of named tuples

3199ef4

Use hasattr instead of try-except

61592de

nre force-pushed the worker_plugin_release_key branch from a45ef8b to 61592de Compare June 17, 2020 09:08

mrocklin merged commit 5172678 into dask:master Jun 18, 2020

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Notify worker plugins when a task is released #3817

Notify worker plugins when a task is released #3817

nre commented May 20, 2020

nre commented Jun 1, 2020

martindurant left a comment

martindurant Jun 15, 2020

martindurant Jun 15, 2020

mrocklin Jun 15, 2020

martindurant Jun 15, 2020

mrocklin Jun 15, 2020

mrocklin Jun 16, 2020

martindurant Jun 15, 2020

mrocklin Jun 15, 2020

martindurant Jun 15, 2020

mrocklin Jun 15, 2020

mrocklin left a comment

mrocklin Jun 15, 2020

mrocklin Jun 15, 2020

mrocklin Jun 15, 2020

mrocklin Jun 15, 2020

mrocklin Jun 15, 2020

mrocklin Jun 15, 2020

nre commented Jun 17, 2020

mrocklin commented Jun 17, 2020

mrocklin commented Jun 18, 2020

		self.observed_notifications.append(ReleasedDep(dep, state))


		Transition = namedtuple("Transition", "key, start, finish")

Notify worker plugins when a task is released #3817

Notify worker plugins when a task is released #3817

Conversation

nre commented May 20, 2020

nre commented Jun 1, 2020

martindurant left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mrocklin left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nre commented Jun 17, 2020

mrocklin commented Jun 17, 2020

mrocklin commented Jun 18, 2020