-
-
Notifications
You must be signed in to change notification settings - Fork 719
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Refactor all event handlers #6410
Conversation
933127a
to
d16cf26
Compare
Unit Test Results 15 files + 15 15 suites +15 6h 35m 10s ⏱️ + 6h 35m 10s For more details on these failures and errors, see this check. Results for commit 42d399f. ± Comparison against base commit 7d280fd. ♻️ This comment has been updated with latest results. |
d16cf26
to
64d4b4a
Compare
@functools.singledispatchmethod | ||
def handle_event(self, ev: StateMachineEvent) -> RecsInstrs: | ||
raise TypeError(ev) # pragma: nocover | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Moved so that I didn't have to move the various refactored functions, thus minimising the diff.
This is temporary - this and all the registered methods will be moved to worker_state_machine.py.
rerun tests Note: rerunning gpuCI tests since those errors should be fixed by #6434 |
cd37b28
to
052828c
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM overall. All comments are nits and could be ignored, besides the list->set question and using descriptive function names.
ts = self.tasks[key] | ||
who_has[key] = {ws.address for ws in ts.who_has} | ||
|
||
who_has = {key: [ws.address for ws in self.tasks[key].who_has] for key in keys} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
who_has = {key: [ws.address for ws in self.tasks[key].who_has] for key in keys} | |
who_has = {key: {ws.address for ws in self.tasks[key].who_has} for key in keys} |
This was changed from a list to a set, was that intentional?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
from a set to a list. yes, as the set doesn't offer any useful feature here and a list is fractionally faster. This also makes it coherent with Scheduler.get_who_has.
"who_has": who_has, | ||
"stimulus_id": stimulus_id, | ||
}, | ||
) | ||
|
||
def request_remove_replicas(self, addr: str, keys: list, *, stimulus_id: str): | ||
def request_remove_replicas( | ||
self, addr: str, keys: list[str], *, stimulus_id: str |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
self, addr: str, keys: list[str], *, stimulus_id: str | |
self, addr: str, keys: Iterable[str], *, stimulus_id: str |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
No, because the variable is sent as-is through the msgpack RPC, so only a list, set, or tuple are accepted.
distributed/worker.py
Outdated
return {"nbytes": {k: sizeof(v) for k, v in data.items()}, "status": "OK"} | ||
|
||
@handle_event.register | ||
def _(self, ev: UpdateDataEvent) -> RecsInstrs: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
def _(self, ev: UpdateDataEvent) -> RecsInstrs: | |
def handle_update_data(self, ev: UpdateDataEvent) -> RecsInstrs: |
Could we please name these functions recognizably, even though they're singledispatch handlers? It makes it much easier to navigate the code. I usually fuzzy-find search for terms like update data
to find the handler.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done
if isinstance(self.run_spec, dict): | ||
self.run_spec = SerializedTask(**self.run_spec) | ||
elif not isinstance(self.run_spec, SerializedTask): | ||
self.run_spec = SerializedTask(task=self.run_spec) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
A bit nitpicky, feel free to ignore. But it might be nice to use init-only variables for function
, args
, kwargs
, and task
, make run_spec
a non-init field, and go back to this when constructing the message:
if isinstance(ts.run_spec, dict):
msg.update(ts.run_spec)
else:
msg["task"] = ts.run_spec
Just a little more explicitness and type checking around what has to be passed.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The reason why I packed it this way was to avoid replicating all the (fairly convoluted) annotations in the SerializedTask class. The SerializedTask class is messy to begin with - given more time, I'd redesign it thoroughly.
ev2 = ev.to_loggable(handled=11.22) | ||
assert ev2.handled == 11.22 | ||
assert ev2.run_spec == SerializedTask(task=None) | ||
assert ev.run_spec == SerializedTask(function=b"blob", args=b"blob") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
assert ev.run_spec == SerializedTask(function=b"blob", args=b"blob") | |
assert ev.run_spec == SerializedTask(function=b"blob", args=b"blob") | |
assert not ev.handled |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This isn't a feature we want to test, just an implementation detail. As a matter of fact, the base method StateMachineEvent.to_loggable overwrites the original object for the sake of performance.
) | ||
ev2 = ev.to_loggable(handled=11.22) | ||
assert ev2.handled == 11.22 | ||
assert ev2.data == {"x": None, "y": None} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
assert ev2.data == {"x": None, "y": None} | |
assert ev2.data == {"x": None, "y": None} | |
assert not ev.handled | |
assert ev.data == {"x": "foo", "y": "bar"} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
same
Part of #5736